[BasicBlockSections] Split cold parts of custom-section functions.

[rlavaee]

This PR makes -basic-block-sections handle functions with custom non-dot-text sections correctly. Cold parts of such functions must be placed in the same section (not in .text.split) but with a unique id.

[llvmbot]

@llvm/pr-subscribers-backend-x86

Changes

This PR makes -basic-block-sections handle functions with custom non-dot-text sections correctly. Cold parts of such functions must be placed in the same section (not in .text.split) but with a unique id.

---

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

2 Files Affected:

• (modified) llvm/lib/CodeGen/TargetLoweringObjectFileImpl.cpp (+13-5)
• (modified) llvm/test/CodeGen/X86/basic-block-sections-pragma-sections.ll (+21-9)

diff --git a/llvm/lib/CodeGen/TargetLoweringObjectFileImpl.cpp b/llvm/lib/CodeGen/TargetLoweringObjectFileImpl.cpp
index 622c8addc548f4f..05d035dc697936b 100644
--- a/llvm/lib/CodeGen/TargetLoweringObjectFileImpl.cpp
+++ b/llvm/lib/CodeGen/TargetLoweringObjectFileImpl.cpp
@@ -1038,14 +1038,21 @@ MCSection *TargetLoweringObjectFileELF::getSectionForMachineBasicBlock(
   // under the .text.eh prefix. For regular sections, we either use a unique
   // name, or a unique ID for the section.
   SmallString<128> Name;
-  if (MBB.getSectionID() == MBBSectionID::ColdSectionID) {
-    Name += BBSectionsColdTextPrefix;
-    Name += MBB.getParent()->getName();
+  StringRef FunctionSectionName = MBB.getParent()->getSection()->getName();
+  if (!FunctionSectionName.equals(".text") && !FunctionSectionName.startswith(".text.")) {
+    // If the original function has a custom non-dot-text section, then split the cold part into that section too, but with a unique id.
+    Name = FunctionSectionName;
+    UniqueID = NextUniqueID++;
+  } else {
+  StringRef FunctionName = MBB.getParent()->getName();
+  if (MBB.getSectionID() == MBBSectionID::ColdSectionID){
+      Name += BBSectionsColdTextPrefix;
+      Name += FunctionName;
   } else if (MBB.getSectionID() == MBBSectionID::ExceptionSectionID) {
     Name += ".text.eh.";
-    Name += MBB.getParent()->getName();
+    Name += FunctionName;
   } else {
-    Name += MBB.getParent()->getSection()->getName();
+    Name += FunctionSectionName;
     if (TM.getUniqueBasicBlockSectionNames()) {
       if (!Name.endswith("."))
         Name += ".";
@@ -1054,6 +1061,7 @@ MCSection *TargetLoweringObjectFileELF::getSectionForMachineBasicBlock(
       UniqueID = NextUniqueID++;
     }
   }
+  }
 
   unsigned Flags = ELF::SHF_ALLOC | ELF::SHF_EXECINSTR;
   std::string GroupName;
diff --git a/llvm/test/CodeGen/X86/basic-block-sections-pragma-sections.ll b/llvm/test/CodeGen/X86/basic-block-sections-pragma-sections.ll
index d63fbdd7b362e07..13a4607100a057a 100644
--- a/llvm/test/CodeGen/X86/basic-block-sections-pragma-sections.ll
+++ b/llvm/test/CodeGen/X86/basic-block-sections-pragma-sections.ll
@@ -1,9 +1,15 @@
+;; Tests for basic block sections applied on a function in a custom section.
 ; RUN: llc < %s -mtriple=x86_64-pc-linux -basic-block-sections=all | FileCheck %s
 ; RUN: llc < %s -mtriple=x86_64-pc-linux -function-sections -basic-block-sections=all | FileCheck %s
-; RUN: echo "!_Z3fooi" > %t.list.txt
-; RUN: echo "!!2" >> %t.list.txt
-; RUN: llc < %s -mtriple=x86_64-pc-linux -basic-block-sections=%t.list.txt | FileCheck %s --check-prefix=LIST
-; RUN: llc < %s -mtriple=x86_64-pc-linux -function-sections -basic-block-sections=%t.list.txt | FileCheck %s --check-prefix=LIST
+; RUN: echo "!_Z3fooi" > %t1.list.txt
+; RUN: echo "!!2" >> %t1.list.txt
+; RUN: llc < %s -mtriple=x86_64-pc-linux -basic-block-sections=%t1.list.txt | FileCheck %s --check-prefix=LIST1
+; RUN: llc < %s -mtriple=x86_64-pc-linux -function-sections -basic-block-sections=%t1.list.txt | FileCheck %s --check-prefix=LIST1
+; RUN: echo "!_Z3fooi" > %t2.list.txt
+; RUN: echo "!!0" >> %t2.list.txt
+; RUN: llc < %s -mtriple=x86_64-pc-linux -basic-block-sections=%t2.list.txt | FileCheck %s --check-prefix=LIST2
+; RUN: llc < %s -mtriple=x86_64-pc-linux -function-sections -basic-block-sections=%t2.list.txt | FileCheck %s --check-prefix=LIST2
+
 
 ; CHECK: .section	foo_section,"ax",@progbits,unique,1
 ; CHECK-LABEL: _Z3fooi:
@@ -12,11 +18,17 @@
 ; CHECK: .section	foo_section,"ax",@progbits,unique,3
 ; CHECK-NEXT: _Z3fooi.__part.2:
 
-; LIST: .section	foo_section,"ax",@progbits,unique,1
-; LIST-LABEL: _Z3fooi:
-; LIST: .section	foo_section,"ax",@progbits,unique,2
-; LIST-NEXT: _Z3fooi.__part.0:
-; LIST-NOT: .section	foo_section,"ax",@progbits,unique,3
+; LIST1: .section	foo_section,"ax",@progbits,unique,1
+; LIST1-LABEL: _Z3fooi:
+; LIST1: .section	foo_section,"ax",@progbits,unique,2
+; LIST1-NEXT: _Z3fooi.__part.0:
+; LIST1-NOT: .section	foo_section,"ax",@progbits,unique,3
+
+; LIST2: .section	foo_section,"ax",@progbits,unique,1
+; LIST2-LABEL: _Z3fooi:
+; LIST2: .section	foo_section,"ax",@progbits,unique,2
+; LIST2-NEXT: _Z3fooi.cold:
+; LIST2-NOT: .section	foo_section,"ax",@progbits,unique,3
 
 ;; Source to generate the IR:
 ;; #pragma clang section text = "foo_section"

[snehasish]

It's a little easier to read if we drop the not on each, make this an OR clause and move the else part into this block.

Also it would be great to refactor this out into a separate function like std::string makeSectionName(MachineBasicBlock& MBB ...) then we could add some unit tests which check the returned string. What do you think?

[rlavaee]

Switched the conditional.

This piece of code also assigns UniqueID and modifies NextUniqueID based on naming. Not sure if it's worth creating that separate function.

[snehasish]

lgtm