[NFC][AsmPrinter] Refactor AsmPrinter and AArch64AsmPrinter to prepare for jump table partitions on aarch64 (llvm#125993)

With
llvm@3feb724,
AsmPrinter can place jump table entries into `.hot` or `.unlikely`
prefixed data sections. This change refactors AsmPrinter and
AArch64AsmPrinter to prepare for the aarch64 port.

* Before this patch, the AsmPrinter class exposes `emitJumpTableInfo` as
a virtual method, and AArch64AsmPrinter overrides `emitJumpTableInfo`
for jump table emission.
* After this patch, both AsmPrinter and AArch64AsmPrinter shares
`AsmPrinter::emitJumpTableInfo`, and class-specific code are moved
inside each class's `emitJumpTableImpl` respectively.

This is a follow-up of llvm#125987,
and llvm#126018 implements the port
for aarch64.

Author: mingmingl-llvm

llvm/include/llvm/CodeGen/AsmPrinter.h

@@ -908,9 +908,8 @@ class AsmPrinter : public MachineFunctionPass {
   // Internal Implementation Details
   //===------------------------------------------------------------------===//
 
-  void emitJumpTableImpl(const MachineJumpTableInfo &MJTI,
-                         ArrayRef<unsigned> JumpTableIndices,
-                         bool JTInDiffSection);
+  virtual void emitJumpTableImpl(const MachineJumpTableInfo &MJTI,
+                                 ArrayRef<unsigned> JumpTableIndices);
 
   void emitJumpTableSizesSection(const MachineJumpTableInfo &MJTI,
                                  const Function &F) const;

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

@@ -2860,22 +2860,12 @@ void AsmPrinter::emitConstantPool() {
 void AsmPrinter::emitJumpTableInfo() {
   const MachineJumpTableInfo *MJTI = MF->getJumpTableInfo();
   if (!MJTI) return;
-  if (MJTI->getEntryKind() == MachineJumpTableInfo::EK_Inline) return;
+
   const std::vector<MachineJumpTableEntry> &JT = MJTI->getJumpTables();
   if (JT.empty()) return;
 
-  // Pick the directive to use to print the jump table entries, and switch to
-  // the appropriate section.
-  const Function &F = MF->getFunction();
-  const TargetLoweringObjectFile &TLOF = getObjFileLowering();
-  bool JTInDiffSection = !TLOF.shouldPutJumpTableInFunctionSection(
-      MJTI->getEntryKind() == MachineJumpTableInfo::EK_LabelDifference32 ||
-          MJTI->getEntryKind() == MachineJumpTableInfo::EK_LabelDifference64,
-      F);
-
   if (!TM.Options.EnableStaticDataPartitioning) {
-    emitJumpTableImpl(*MJTI, llvm::to_vector(llvm::seq<unsigned>(JT.size())),
-                      JTInDiffSection);
+    emitJumpTableImpl(*MJTI, llvm::to_vector(llvm::seq<unsigned>(JT.size())));
     return;
   }
 
@@ -2891,33 +2881,40 @@ void AsmPrinter::emitJumpTableInfo() {
     }
   }
 
-  emitJumpTableImpl(*MJTI, HotJumpTableIndices, JTInDiffSection);
-  emitJumpTableImpl(*MJTI, ColdJumpTableIndices, JTInDiffSection);
+  emitJumpTableImpl(*MJTI, HotJumpTableIndices);
+  emitJumpTableImpl(*MJTI, ColdJumpTableIndices);
 }
 
 void AsmPrinter::emitJumpTableImpl(const MachineJumpTableInfo &MJTI,
-                                   ArrayRef<unsigned> JumpTableIndices,
-                                   bool JTInDiffSection) {
-  if (JumpTableIndices.empty())
+                                   ArrayRef<unsigned> JumpTableIndices) {
+  if (MJTI.getEntryKind() == MachineJumpTableInfo::EK_Inline ||
+      JumpTableIndices.empty())
     return;
 
   const TargetLoweringObjectFile &TLOF = getObjFileLowering();
   const Function &F = MF->getFunction();
   const std::vector<MachineJumpTableEntry> &JT = MJTI.getJumpTables();
   MCSection *JumpTableSection = nullptr;
-  if (TM.Options.EnableStaticDataPartitioning) {
-    JumpTableSection =
-        TLOF.getSectionForJumpTable(F, TM, &JT[JumpTableIndices.front()]);
-  } else {
-    JumpTableSection = TLOF.getSectionForJumpTable(F, TM);
-  }
 
-  const DataLayout &DL = MF->getDataLayout();
+  const bool UseLabelDifference =
+      MJTI.getEntryKind() == MachineJumpTableInfo::EK_LabelDifference32 ||
+      MJTI.getEntryKind() == MachineJumpTableInfo::EK_LabelDifference64;
+  // Pick the directive to use to print the jump table entries, and switch to
+  // the appropriate section.
+  const bool JTInDiffSection =
+      !TLOF.shouldPutJumpTableInFunctionSection(UseLabelDifference, F);
   if (JTInDiffSection) {
+    if (TM.Options.EnableStaticDataPartitioning) {
+      JumpTableSection =
+          TLOF.getSectionForJumpTable(F, TM, &JT[JumpTableIndices.front()]);
+    } else {
+      JumpTableSection = TLOF.getSectionForJumpTable(F, TM);
+    }
     OutStreamer->switchSection(JumpTableSection);
   }
 
-  emitAlignment(Align(MJTI.getEntryAlignment(MF->getDataLayout())));
+  const DataLayout &DL = MF->getDataLayout();
+  emitAlignment(Align(MJTI.getEntryAlignment(DL)));
 
   // Jump tables in code sections are marked with a data_region directive
   // where that's supported.

llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp

@@ -113,7 +113,8 @@ class AArch64AsmPrinter : public AsmPrinter {
   const MCExpr *lowerBlockAddressConstant(const BlockAddress &BA) override;
 
   void emitStartOfAsmFile(Module &M) override;
-  void emitJumpTableInfo() override;
+  void emitJumpTableImpl(const MachineJumpTableInfo &MJTI,
+                         ArrayRef<unsigned> JumpTableIndices) override;
   std::tuple<const MCSymbol *, uint64_t, const MCSymbol *,
              codeview::JumpTableEntrySize>
   getCodeViewJumpTableInfo(int JTI, const MachineInstr *BranchInstr,
@@ -1290,19 +1291,18 @@ void AArch64AsmPrinter::PrintDebugValueComment(const MachineInstr *MI,
   printOperand(MI, NOps - 2, OS);
 }
 
-void AArch64AsmPrinter::emitJumpTableInfo() {
-  const MachineJumpTableInfo *MJTI = MF->getJumpTableInfo();
-  if (!MJTI) return;
-
-  const std::vector<MachineJumpTableEntry> &JT = MJTI->getJumpTables();
-  if (JT.empty()) return;
-
+void AArch64AsmPrinter::emitJumpTableImpl(const MachineJumpTableInfo &MJTI,
+                                          ArrayRef<unsigned> JumpTableIndices) {
+  // Fast return if there is nothing to emit to avoid creating empty sections.
+  if (JumpTableIndices.empty())
+    return;
   const TargetLoweringObjectFile &TLOF = getObjFileLowering();
   MCSection *ReadOnlySec = TLOF.getSectionForJumpTable(MF->getFunction(), TM);
   OutStreamer->switchSection(ReadOnlySec);
 
+  const std::vector<MachineJumpTableEntry> &JT = MJTI.getJumpTables();
   auto AFI = MF->getInfo<AArch64FunctionInfo>();
-  for (unsigned JTI = 0, e = JT.size(); JTI != e; ++JTI) {
+  for (unsigned JTI : JumpTableIndices) {
     const std::vector<MachineBasicBlock*> &JTBBs = JT[JTI].MBBs;
 
     // If this jump table was deleted, ignore it.

llvm/test/CodeGen/AArch64/jump-table-partition.ll

@@ -17,48 +17,29 @@
 ; DEFAULT: .section .rodata.func_without_profile,"a",@progbits
 ; DEFAULT:   .LJTI1_0:
 ; DEFAULT: .section .rodata.bar_prefix.bar,"a",@progbits
-; DEFAULT: .LJTI2_0
-
-; RUN: llc -mtriple=aarch64-unknown-linux-gnu -enable-split-machine-functions \
-; RUN:     -partition-static-data-sections=true -function-sections=true \
+; DEFAULT:   .LJTI2_0
+
+; Test that section names are uniqufied by numbers but not function names with
+; {-function-sections, -unique-section-names=false}. Specifically, @foo jump
+; tables are emitted in two sections, one with unique ID 2 and the other with
+; unique ID 3.
+; RUN: llc -mtriple=aarch64-unknown-linux-gnu -partition-static-data-sections \
+; RUN:     -function-sections -unique-section-names=false \
 ; RUN:     -aarch64-enable-atomic-cfg-tidy=false -aarch64-min-jump-table-entries=2 \
-; RUN:     -unique-section-names=false %s -o - 2>&1 | FileCheck %s --check-prefixes=NUM,JT
+; RUN:     %s -o - 2>&1 | FileCheck %s --check-prefixes=NUM,JT
 
-; Section names will optionally have `.<func>` if -function-sections is enabled.
-; RUN: llc -mtriple=aarch64-unknown-linux-gnu -enable-split-machine-functions \
-; RUN:     -partition-static-data-sections=true -function-sections=true \
+; Section names will optionally have `.<func>` with {-function-sections, -unique-section-names}.
+; RUN: llc -mtriple=aarch64-unknown-linux-gnu -partition-static-data-sections \
+; RUN:     -function-sections -unique-section-names \
 ; RUN:     -aarch64-enable-atomic-cfg-tidy=false -aarch64-min-jump-table-entries=2  \
 ; RUN:     %s -o - 2>&1 | FileCheck %s --check-prefixes=FUNC,JT
 
-; RUN: llc -mtriple=aarch64-unknown-linux-gnu -enable-split-machine-functions \
-; RUN:     -partition-static-data-sections=true -function-sections=false \
+; Test that section names won't have `.<func>` with -function-sections=false.
+; RUN: llc -mtriple=aarch64-unknown-linux-gnu -partition-static-data-sections \
+; RUN:     -function-sections=false \
 ; RUN:     -aarch64-enable-atomic-cfg-tidy=false -aarch64-min-jump-table-entries=2 \
 ; RUN:     %s -o - 2>&1 | FileCheck %s --check-prefixes=FUNCLESS,JT
 
-; A function's section prefix is used for all jump tables of this function.
-; @foo is hot so its jump table data section has a hot prefix.
-; NUM:    .section .rodata.hot.,"a",@progbits,unique,2
-; FUNC:     .section .rodata.hot.foo,"a",@progbits
-; FUNCLESS: .section .rodata.hot.,"a",@progbits
-; JT: .LJTI0_0:
-; JT: .LJTI0_1:
-; JT: .LJTI0_2:
-; JT: .LJTI0_3:
-
-; func_without_profile doesn't have profiles, so its jumptable doesn't have
-; hotness-based prefix.
-; NUM: .section .rodata,"a",@progbits,unique,4
-; FUNC: .section .rodata.func_without_profile,"a",@progbits
-; FUNCLESS: .section .rodata,"a",@progbits
-; JT: .LJTI1_0:
-
-; @bar doesn't have profile information and it has a section prefix.
-; Tests that its jump tables are placed in sections with function prefixes.
-; NUM: .section .rodata.bar_prefix.,"a",@progbits,unique,
-; FUNC: .section .rodata.bar_prefix.bar
-; FUNCLESS: .section .rodata.bar_prefix.,"a"
-; JT: .LJTI2_0
-
 target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128"
 target triple = "aarch64-unknown-linux-gnu"
 
@@ -71,6 +52,19 @@ target triple = "aarch64-unknown-linux-gnu"
 @default = private constant [8 x i8] c"default\00"
 @jt3 = private constant [4 x i8] c"jt3\00"
 
+; A function's section prefix is used for all jump tables of this function.
+; @foo is hot so its jump table data section has a hot prefix.
+; NUM:          .section .rodata.hot.,"a",@progbits,unique,2
+; FUNC:         .section .rodata.hot.foo,"a",@progbits
+; FUNCLESS:     .section .rodata.hot.,"a",@progbits
+; JT:           .LJTI0_0:
+; JT:           .LJTI0_2:
+; NUM:          .section	.rodata.hot.,"a",@progbits,unique,3
+; FUNC-NOT:     .section .rodata.hot.foo
+; FUNCLESS-NOT: .section .rodata.hot.,"a",@progbits
+; JT:           .LJTI0_1:
+; JT:           .LJTI0_3:
+
 ; jt0 and jt2 are hot. jt1 and jt3 are cold.
 define i32 @foo(i32 %num) !prof !13 {
 entry:
@@ -168,6 +162,12 @@ return:
   ret i32 %mod3
 }
 
+; @func_without_profile doesn't have profiles, so its jumptable doesn't have
+; hotness-based prefix.
+; NUM:        .section .rodata,"a",@progbits,unique,5
+; FUNC:       .section .rodata.func_without_profile,"a",@progbits
+; FUNCLESS:   .section .rodata,"a",@progbits
+; JT:         .LJTI1_0:
 define void @func_without_profile(i32 %num) {
 entry:
   switch i32 %num, label %sw.default [
@@ -191,6 +191,12 @@ sw.epilog:
   ret void
 }
 
+; @bar doesn't have profile information and it has a section prefix.
+; Tests that its jump tables are placed in sections with function prefixes.
+; NUM:        .section .rodata.bar_prefix.,"a",@progbits,unique,7
+; FUNC:       .section .rodata.bar_prefix.bar
+; FUNCLESS:   .section .rodata.bar_prefix.,"a"
+; JT:         .LJTI2_0
 define void @bar(i32 %num) !section_prefix !20  {
 entry:
   switch i32 %num, label %sw.default [