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

mingmingl-llvm · 2025-02-06T03:30:26Z

With 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 #125987, and #126018 implements the port for aarch64.

mingmingl-llvm · 2025-02-06T06:36:08Z

cc @Colibrow fyi

llvmbot · 2025-02-06T06:37:32Z

@llvm/pr-subscribers-backend-aarch64

Author: Mingming Liu (mingmingl-llvm)

Changes

With 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 emitJumpTableImpl respectively.

This is a follow-up of #125987

Full diff: https://github.com/llvm/llvm-project/pull/125993.diff

4 Files Affected:

(modified) llvm/include/llvm/CodeGen/AsmPrinter.h (+2-3)
(modified) llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp (+21-25)
(modified) llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp (+9-9)
(modified) llvm/test/CodeGen/AArch64/jump-table-partition.ll (+11-8)

diff --git a/llvm/include/llvm/CodeGen/AsmPrinter.h b/llvm/include/llvm/CodeGen/AsmPrinter.h
index 3da63af5ba5716c..9ef9888af990c63 100644
--- a/llvm/include/llvm/CodeGen/AsmPrinter.h
+++ b/llvm/include/llvm/CodeGen/AsmPrinter.h
@@ -893,9 +893,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 emitJumpTableEntry(const MachineJumpTableInfo &MJTI,
                           const MachineBasicBlock *MBB, unsigned uid) const;
 
diff --git a/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp b/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
index 44b10c3ef997267..07fe589fcbbd065 100644
--- a/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
+++ b/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
@@ -2855,22 +2855,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;
   }
 
@@ -2886,33 +2876,39 @@ 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();
+  // Pick the directive to use to print the jump table entries, and switch to
+  // the appropriate section.
+  const bool JTInDiffSection = !TLOF.shouldPutJumpTableInFunctionSection(
+      MJTI.getEntryKind() == MachineJumpTableInfo::EK_LabelDifference32 ||
+          MJTI.getEntryKind() == MachineJumpTableInfo::EK_LabelDifference64,
+      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.
diff --git a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
index f1f25b65fc53fac..c92c203f247954b 100644
--- a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
+++ b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
@@ -112,7 +112,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,
@@ -1273,19 +1274,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.
diff --git a/llvm/test/CodeGen/AArch64/jump-table-partition.ll b/llvm/test/CodeGen/AArch64/jump-table-partition.ll
index 122bbaef09185a2..e0525d0384a978e 100644
--- a/llvm/test/CodeGen/AArch64/jump-table-partition.ll
+++ b/llvm/test/CodeGen/AArch64/jump-table-partition.ll
@@ -37,26 +37,29 @@
 
 ; 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
+; 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:
+; 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:
 
 ; 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
+; NUM:        .section .rodata,"a",@progbits,unique,5
+; 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"
+; NUM:        .section .rodata.bar_prefix.,"a",@progbits,unique,7
+; 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"

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

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

snehasish · 2025-04-08T00:10:34Z

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

+; 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.


Can you move these CHECKS (L38-L62) close to their appropriate functions?

done, and also did this for the subsequent groups of CHECK lines.

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

snehasish · 2025-04-08T00:13:59Z

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

+; RUN: llc -mtriple=aarch64-unknown-linux-gnu -enable-split-machine-functions \
+; RUN:     -partition-static-data-sections=true -function-sections=true \
+; 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


Why is unique-section-names false here?

With -function-sections=true and -unique-section-names=false, this line specifically tests that @foo jump tables are emitted in two sections (because of -partition-static-data-sections), one with unique ID 2 and the other with unique ID 3.

The two sections have identical name (.rodata.hot). The .hot substring comes from function prefix, like how jump table section names get .hot substring in L8-L10 (which doesn't enable -partition-static-data-sections).

auto& DL = MF->getDataLayout()`.

snehasish

lgtm

mingmingl-llvm · 2025-04-14T04:45:05Z

@davemgreen a friendly ping for review when you get a chance. Let me know if there are any questions!

davemgreen

I think this looks OK. LGTM

davemgreen · 2025-04-14T07:26:33Z

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp


-  const DataLayout &DL = MF->getDataLayout();
+  const bool UseLabelDifference =
+      (MJTI.getEntryKind() == MachineJumpTableInfo::EK_LabelDifference32 ||


Can drop () brackets.

mingmingl-llvm

thanks for the reviews!

mingmingl-llvm · 2025-04-14T16:49:47Z

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp


-  const DataLayout &DL = MF->getDataLayout();
+  const bool UseLabelDifference =
+      (MJTI.getEntryKind() == MachineJumpTableInfo::EK_LabelDifference32 ||


mingmingl-llvm · 2025-04-14T18:08:15Z

I'll merge as all presubmit tests pass except Windows Premerge Checks (Test Only - Please Ignore Results), which passed its Build and Test step (which I interpret as test passes) and failed at Upload Artifacts.

* origin/main: (199 commits) [NFC][AsmPrinter] Refactor AsmPrinter and AArch64AsmPrinter to prepare for jump table partitions on aarch64 (llvm#125993) [HEXAGON] Fix corner cases for hwloops pass (llvm#135439) [flang] Handle volatility in lowering and codegen (llvm#135311) [MLIR][Shape] Support >2 args in `shape.broadcast` folder (llvm#126808) [DirectX] Use scalar arguments for @llvm.dx.dot intrinsics (llvm#134570) Remove the redundant check for "WeakPtr" in isSmartPtrClass to fix the issue 135612. (llvm#135629) [BOLT] Support relative vtable (llvm#135449) [flang] Fix linking to libMLIR (llvm#135483) [AsmPrinter] Link .section_sizes to the correct section (llvm#135583) [ctxprof] Handle instrumenting functions with `musttail` calls (llvm#135121) [SystemZ] Consider VST/VL as SimpleBDXStore/Load (llvm#135623) [libc++][CI] Pin the XCode version. (llvm#135412) [lldb-dap] Fix win32 build. (llvm#135638) [Interp] Mark inline-virtual.cpp as unsupported with ASan (llvm#135402) [libc++] Removes the _LIBCPP_VERBOSE_ABORT_NOT_NOEXCEPT macro. (llvm#135494) [CVP] Add tests for ucmp/scmp with switch (NFC) [mlir][tosa] Align AbsOp example variable names (llvm#135268) [mlir][tosa] Align AddOp examples to spec (llvm#135266) [mlir][tosa] Align RFFT2d and FFT2d operator examples (llvm#135261) [flang][OpenMP][HLFIR] Support vector subscripted array sections for DEPEND (llvm#133892) ...

…sections for aarch64 (#126018) This is a follow-up patch of #125993 to port jump table partitions for aarch64. --------- Co-authored-by: Ellis Hoag <[email protected]>

…fixed data sections for aarch64 (#126018) This is a follow-up patch of llvm/llvm-project#125993 to port jump table partitions for aarch64. --------- Co-authored-by: Ellis Hoag <[email protected]>

…e 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.

…sections for aarch64 (llvm#126018) This is a follow-up patch of llvm#125993 to port jump table partitions for aarch64. --------- Co-authored-by: Ellis Hoag <[email protected]>

mingmingl-llvm added 4 commits February 5, 2025 19:08

[NFC][Test]Add a test case for aarch64 jump table sections

fb2e048

remove requires and stat

a5e6552

[nfc]asm printer jump table

169486b

rebase on top of precommit test

acc7e25

mingmingl-llvm force-pushed the users/mingmingl-llvm/spr/nfc branch from 411df61 to acc7e25 Compare February 6, 2025 06:01

mingmingl-llvm changed the base branch from main to users/mingmingl-llvm/spr/precommit February 6, 2025 06:02

mingmingl-llvm added 2 commits February 5, 2025 22:32

Fast return in aarch64 impl

bbec612

code style

eaa5fd5

mingmingl-llvm requested review from snehasish, ellishg, williamweixiao, HaohaiWen and david-xl February 6, 2025 06:35

mingmingl-llvm marked this pull request as ready for review February 6, 2025 06:36

llvmbot added the backend:AArch64 label Feb 6, 2025

mingmingl-llvm mentioned this pull request Feb 6, 2025

[AArch64AsmPrinter]Place jump tables into hot/unlikely-prefixed data sections for aarch64 #126018

Merged

Base automatically changed from users/mingmingl-llvm/spr/precommit to main April 7, 2025 23:21

snehasish reviewed Apr 8, 2025

View reviewed changes

mingmingl-llvm added 3 commits April 8, 2025 06:05

catch up with main branch and resolve feedback

a2e3812

Merge branch 'main' into users/mingmingl-llvm/spr/nfc

8a09ec1

drop -enable-split-machine-functions in test case, and add back `const

96bbb72

auto& DL = MF->getDataLayout()`.

mingmingl-llvm requested a review from snehasish April 8, 2025 06:40

snehasish approved these changes Apr 8, 2025

View reviewed changes

MaskRay requested a review from davemgreen April 8, 2025 16:46

davemgreen approved these changes Apr 14, 2025

View reviewed changes

drop brackets

bffe64d

mingmingl-llvm commented Apr 14, 2025

View reviewed changes

mingmingl-llvm merged commit c766d0c into main Apr 14, 2025
10 of 11 checks passed

mingmingl-llvm deleted the users/mingmingl-llvm/spr/nfc branch April 14, 2025 18:08

mingmingl-llvm restored the users/mingmingl-llvm/spr/nfc branch April 14, 2025 18:33

mingmingl-llvm deleted the users/mingmingl-llvm/spr/nfc branch April 14, 2025 20:29

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

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

Uh oh!

Conversation

mingmingl-llvm commented Feb 6, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

mingmingl-llvm commented Feb 6, 2025

Uh oh!

llvmbot commented Feb 6, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

snehasish Apr 8, 2025

Choose a reason for hiding this comment

Uh oh!

mingmingl-llvm Apr 8, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

snehasish Apr 8, 2025

Choose a reason for hiding this comment

Uh oh!

mingmingl-llvm Apr 8, 2025

Choose a reason for hiding this comment

Uh oh!

snehasish left a comment

Choose a reason for hiding this comment

Uh oh!

mingmingl-llvm commented Apr 14, 2025

Uh oh!

davemgreen left a comment

Choose a reason for hiding this comment

Uh oh!

davemgreen Apr 14, 2025

Choose a reason for hiding this comment

Uh oh!

mingmingl-llvm Apr 14, 2025

Choose a reason for hiding this comment

Uh oh!

mingmingl-llvm left a comment

Choose a reason for hiding this comment

Uh oh!

mingmingl-llvm Apr 14, 2025

Choose a reason for hiding this comment

Uh oh!

mingmingl-llvm commented Apr 14, 2025

Uh oh!

Uh oh!

Uh oh!

mingmingl-llvm commented Feb 6, 2025 •

edited

Loading