[llvm] [BasicBlockSections] Split cold parts of custom-section functions. (PR #66731)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-x86

<details>
<summary>Changes</summary>

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
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", at progbits,unique,1
 ; CHECK-LABEL: _Z3fooi:
@@ -12,11 +18,17 @@
 ; CHECK: .section	foo_section,"ax", at progbits,unique,3
 ; CHECK-NEXT: _Z3fooi.__part.2:
 
-; LIST: .section	foo_section,"ax", at progbits,unique,1
-; LIST-LABEL: _Z3fooi:
-; LIST: .section	foo_section,"ax", at progbits,unique,2
-; LIST-NEXT: _Z3fooi.__part.0:
-; LIST-NOT: .section	foo_section,"ax", at progbits,unique,3
+; LIST1: .section	foo_section,"ax", at progbits,unique,1
+; LIST1-LABEL: _Z3fooi:
+; LIST1: .section	foo_section,"ax", at progbits,unique,2
+; LIST1-NEXT: _Z3fooi.__part.0:
+; LIST1-NOT: .section	foo_section,"ax", at progbits,unique,3
+
+; LIST2: .section	foo_section,"ax", at progbits,unique,1
+; LIST2-LABEL: _Z3fooi:
+; LIST2: .section	foo_section,"ax", at progbits,unique,2
+; LIST2-NEXT: _Z3fooi.cold:
+; LIST2-NOT: .section	foo_section,"ax", at progbits,unique,3
 
 ;; Source to generate the IR:
 ;; #pragma clang section text = "foo_section"

``````````

</details>


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