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

[via llvm-commits]

Author: Rahman Lavaee
Date: 2023-09-22T13:49:12-07:00
New Revision: 897a0b01d6a4f063b4fa2561a8aeaff14f411053

URL: https://github.com/llvm/llvm-project/commit/897a0b01d6a4f063b4fa2561a8aeaff14f411053
DIFF: https://github.com/llvm/llvm-project/commit/897a0b01d6a4f063b4fa2561a8aeaff14f411053.diff

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

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.

Added: 
    

Modified: 
    llvm/lib/CodeGen/TargetLoweringObjectFileImpl.cpp
    llvm/test/CodeGen/X86/basic-block-sections-pragma-sections.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/CodeGen/TargetLoweringObjectFileImpl.cpp b/llvm/lib/CodeGen/TargetLoweringObjectFileImpl.cpp
index 622c8addc548f4f..6210e7fc128a30c 100644
--- a/llvm/lib/CodeGen/TargetLoweringObjectFileImpl.cpp
+++ b/llvm/lib/CodeGen/TargetLoweringObjectFileImpl.cpp
@@ -1038,21 +1038,32 @@ 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();
-  } else if (MBB.getSectionID() == MBBSectionID::ExceptionSectionID) {
-    Name += ".text.eh.";
-    Name += MBB.getParent()->getName();
-  } else {
-    Name += MBB.getParent()->getSection()->getName();
-    if (TM.getUniqueBasicBlockSectionNames()) {
-      if (!Name.endswith("."))
-        Name += ".";
-      Name += MBB.getSymbol()->getName();
+  StringRef FunctionSectionName = MBB.getParent()->getSection()->getName();
+  if (FunctionSectionName.equals(".text") ||
+      FunctionSectionName.startswith(".text.")) {
+    // Function is in a regular .text section.
+    StringRef FunctionName = MBB.getParent()->getName();
+    if (MBB.getSectionID() == MBBSectionID::ColdSectionID) {
+      Name += BBSectionsColdTextPrefix;
+      Name += FunctionName;
+    } else if (MBB.getSectionID() == MBBSectionID::ExceptionSectionID) {
+      Name += ".text.eh.";
+      Name += FunctionName;
     } else {
-      UniqueID = NextUniqueID++;
+      Name += FunctionSectionName;
+      if (TM.getUniqueBasicBlockSectionNames()) {
+        if (!Name.endswith("."))
+          Name += ".";
+        Name += MBB.getSymbol()->getName();
+      } else {
+        UniqueID = NextUniqueID++;
+      }
     }
+  } else {
+    // If the original function has a custom non-dot-text section, then emit
+    // all basic block sections into that section too, each with a unique id.
+    Name = FunctionSectionName;
+    UniqueID = NextUniqueID++;
   }
 
   unsigned Flags = ELF::SHF_ALLOC | ELF::SHF_EXECINSTR;

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"