[llvm-branch-commits] [llvm] [AArch64AsmPrinter]Place jump tables into hot/unlikely-prefixed data sections for aarch64 (PR #126018)

[Mingming Liu via llvm-branch-commits]

https://github.com/mingmingl-llvm created https://github.com/llvm/llvm-project/pull/126018

This is a follow-up patch of https://github.com/llvm/llvm-project/pull/125993 to port jump table partitions for aarch64.

>From bcd1e154c97f57a6a7c00c5f137c301805eea0c2 Mon Sep 17 00:00:00 2001
From: mingmingl <mingmingl at google.com>
Date: Wed, 5 Feb 2025 23:08:11 -0800
Subject: [PATCH] [AArch64AsmPrinter]Place jump tables into
 hot/unlikely-prefixed sections for aarch64

---
 llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp     | 12 ++++++++++--
 llvm/test/CodeGen/AArch64/jump-table-partition.ll | 13 +++++++------
 2 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
index c92c203f247954..0593c1cd2b14d2 100644
--- a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
+++ b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp
@@ -1280,10 +1280,18 @@ void AArch64AsmPrinter::emitJumpTableImpl(const MachineJumpTableInfo &MJTI,
   if (JumpTableIndices.empty())
     return;
   const TargetLoweringObjectFile &TLOF = getObjFileLowering();
-  MCSection *ReadOnlySec = TLOF.getSectionForJumpTable(MF->getFunction(), TM);
+  const auto &F = MF->getFunction();
+  const std::vector<MachineJumpTableEntry> &JT = MJTI.getJumpTables();
+
+  MCSection *ReadOnlySec = nullptr;
+  if (TM.Options.EnableStaticDataPartitioning) {
+    ReadOnlySec =
+        TLOF.getSectionForJumpTable(F, TM, &JT[JumpTableIndices.front()]);
+  } else {
+    ReadOnlySec = TLOF.getSectionForJumpTable(F, TM);
+  }
   OutStreamer->switchSection(ReadOnlySec);
 
-  const std::vector<MachineJumpTableEntry> &JT = MJTI.getJumpTables();
   auto AFI = MF->getInfo<AArch64FunctionInfo>();
   for (unsigned JTI : JumpTableIndices) {
     const std::vector<MachineBasicBlock*> &JTBBs = JT[JTI].MBBs;
diff --git a/llvm/test/CodeGen/AArch64/jump-table-partition.ll b/llvm/test/CodeGen/AArch64/jump-table-partition.ll
index e0525d0384a978..d282c366578d70 100644
--- a/llvm/test/CodeGen/AArch64/jump-table-partition.ll
+++ b/llvm/test/CodeGen/AArch64/jump-table-partition.ll
@@ -35,16 +35,17 @@
 ; 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", at progbits,unique,2
+; In function @foo, the 2 switch instructions to jt0.* and jt1.* are placed in
+; hot-prefixed sections, and the 2 switch instructions to jt2.* and jt3.* are
+; placed in cold-prefixed sections.
+; NUM:    .section .rodata.hot.,"a", at progbits,unique,2
 ; FUNC:     .section .rodata.hot.foo,"a", at progbits
 ; FUNCLESS: .section .rodata.hot.,"a", at progbits
 ; JT: .LJTI0_0:
 ; JT: .LJTI0_2:
-; NUM:            .section	.rodata.hot.,"a", at progbits,unique,3
-; FUNC-NOT:       .section .rodata.hot.foo
-; FUNCLESS-NOT:   .section .rodata.hot.,"a", at progbits
+; NUM:        .section .rodata.unlikely.,"a", at progbits,unique,3
+; FUNC:       .section .rodata.unlikely.foo,
+; FUNCLESS:   .section .rodata.unlikely.,"a", at progbits
 ; JT: .LJTI0_1:
 ; JT: .LJTI0_3: