Basic block sections for functions with implicit-section-name attribute

Functions can have section names set via #pragma or section attributes,
basic block sections should be correctly named for such functions.

With #pragma, the expectation is that all functions in that file are placed
in the same section in the final binary. Basic block sections should be
correctly named with the unique flag set so that the final binary has all the
basic blocks of the function in that named section. This patch fixes the bug
by calling getExplictSectionGlobal when implicit-section-name attribute is set
to make sure the function's basic blocks get the correct section name.

Differential Revision: https://reviews.llvm.org/D101311

diff --git a/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp b/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
index 0c5126d..581b877 100644 (file)
--- a/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
+++ b/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
@@ -710,9 +710,9 @@ void AsmPrinter::emitFunctionHeader() {
   emitConstantPool();
 
   // Print the 'header' of function.
-  // If basic block sections is desired and function sections is off,
-  // explicitly request a unique section for this function.
-  if (MF->front().isBeginSection() && !TM.getFunctionSections())
+  // If basic block sections are desired, explicitly request a unique section
+  // for this function's entry block.
+  if (MF->front().isBeginSection())
     MF->setSection(getObjFileLowering().getUniqueSectionForFunction(F, TM));
   else
     MF->setSection(getObjFileLowering().SectionForGlobal(&F, TM));

diff --git a/llvm/lib/CodeGen/TargetLoweringObjectFileImpl.cpp b/llvm/lib/CodeGen/TargetLoweringObjectFileImpl.cpp
index 6fed2ee..737a997 100644 (file)
--- a/llvm/lib/CodeGen/TargetLoweringObjectFileImpl.cpp
+++ b/llvm/lib/CodeGen/TargetLoweringObjectFileImpl.cpp
@@ -652,8 +652,10 @@ public:
 };
 }
 
-MCSection *TargetLoweringObjectFileELF::getExplicitSectionGlobal(
-    const GlobalObject *GO, SectionKind Kind, const TargetMachine &TM) const {
+static MCSection *selectExplicitSectionGlobal(
+    const GlobalObject *GO, SectionKind Kind, const TargetMachine &TM,
+    MCContext &Ctx, Mangler &Mang, unsigned &NextUniqueID,
+    bool Retain, bool ForceUnique) {
   StringRef SectionName = GO->getSection();
 
   // Check if '#pragma clang section' name is applicable.
@@ -696,23 +698,22 @@ MCSection *TargetLoweringObjectFileELF::getExplicitSectionGlobal(
   unsigned UniqueID = MCContext::GenericSectionID;
   const MCSymbolELF *LinkedToSym = getLinkedToSymbol(GO, TM);
   const bool Associated = GO->getMetadata(LLVMContext::MD_associated);
-  const bool Retain = Used.count(GO);
   if (Associated || Retain) {
     UniqueID = NextUniqueID++;
     if (Associated)
       Flags |= ELF::SHF_LINK_ORDER;
-    if (Retain && (getContext().getAsmInfo()->useIntegratedAssembler() ||
-                   getContext().getAsmInfo()->binutilsIsAtLeast(2, 36)))
+    if (Retain && (Ctx.getAsmInfo()->useIntegratedAssembler() ||
+                   Ctx.getAsmInfo()->binutilsIsAtLeast(2, 36)))
       Flags |= ELF::SHF_GNU_RETAIN;
   } else {
-    if (getContext().getAsmInfo()->useIntegratedAssembler() ||
-        getContext().getAsmInfo()->binutilsIsAtLeast(2, 35)) {
+    if (Ctx.getAsmInfo()->useIntegratedAssembler() ||
+        Ctx.getAsmInfo()->binutilsIsAtLeast(2, 35)) {
       // Symbols must be placed into sections with compatible entry
       // sizes. Generate unique sections for symbols that have not
       // been assigned to compatible sections.
       if (Flags & ELF::SHF_MERGE) {
-        auto maybeID = getContext().getELFUniqueIDForEntsize(SectionName, Flags,
-                                                             EntrySize);
+        auto maybeID = Ctx.getELFUniqueIDForEntsize(SectionName, Flags,
+                                                    EntrySize);
         if (maybeID)
           UniqueID = *maybeID;
         else {
@@ -721,9 +722,8 @@ MCSection *TargetLoweringObjectFileELF::getExplicitSectionGlobal(
           // to unique the section as the entry size for this symbol will be
           // compatible with implicitly created sections.
           SmallString<128> ImplicitSectionNameStem = getELFSectionNameForGlobal(
-              GO, Kind, getMangler(), TM, EntrySize, false);
-          if (!(getContext().isELFImplicitMergeableSectionNamePrefix(
-                    SectionName) &&
+              GO, Kind, Mang, TM, EntrySize, false);
+          if (!(Ctx.isELFImplicitMergeableSectionNamePrefix(SectionName) &&
                 SectionName.startswith(ImplicitSectionNameStem)))
             UniqueID = NextUniqueID++;
         }
@@ -731,8 +731,8 @@ MCSection *TargetLoweringObjectFileELF::getExplicitSectionGlobal(
         // We need to unique the section if the user has explicity
         // assigned a non-mergeable symbol to a section name for
         // a generic mergeable section.
-        if (getContext().isELFGenericMergeableSection(SectionName)) {
-          auto maybeID = getContext().getELFUniqueIDForEntsize(
+        if (Ctx.isELFGenericMergeableSection(SectionName)) {
+          auto maybeID = Ctx.getELFUniqueIDForEntsize(
               SectionName, Flags, EntrySize);
           UniqueID = maybeID ? *maybeID : NextUniqueID++;
         }
@@ -749,7 +749,13 @@ MCSection *TargetLoweringObjectFileELF::getExplicitSectionGlobal(
     }
   }
 
-  MCSectionELF *Section = getContext().getELFSection(
+  // Increment uniqueID if we are forced to emit a unique section.
+  // This works perfectly fine with section attribute or pragma section as the
+  // sections with the same name are grouped together by the assembler.
+  if (ForceUnique && UniqueID == MCContext::GenericSectionID)
+    UniqueID = NextUniqueID++;
+
+  MCSectionELF *Section = Ctx.getELFSection(
       SectionName, getELFSectionType(SectionName, Kind), Flags, EntrySize,
       Group, IsComdat, UniqueID, LinkedToSym);
   // Make sure that we did not get some other section with incompatible sh_link.
@@ -757,8 +763,8 @@ MCSection *TargetLoweringObjectFileELF::getExplicitSectionGlobal(
   assert(Section->getLinkedToSymbol() == LinkedToSym &&
          "Associated symbol mismatch between sections");
 
-  if (!(getContext().getAsmInfo()->useIntegratedAssembler() ||
-        getContext().getAsmInfo()->binutilsIsAtLeast(2, 35))) {
+  if (!(Ctx.getAsmInfo()->useIntegratedAssembler() ||
+        Ctx.getAsmInfo()->binutilsIsAtLeast(2, 35))) {
     // If we are using GNU as before 2.35, then this symbol might have
     // been placed in an incompatible mergeable section. Emit an error if this
     // is the case to avoid creating broken output.
@@ -777,6 +783,13 @@ MCSection *TargetLoweringObjectFileELF::getExplicitSectionGlobal(
   return Section;
 }
 
+MCSection *TargetLoweringObjectFileELF::getExplicitSectionGlobal(
+    const GlobalObject *GO, SectionKind Kind, const TargetMachine &TM) const {
+  return selectExplicitSectionGlobal(GO, Kind, TM, getContext(), getMangler(),
+                                     NextUniqueID, Used.count(GO),
+                                     /* ForceUnique = */false);
+}
+
 static MCSectionELF *selectELFSectionForGlobal(
     MCContext &Ctx, const GlobalObject *GO, SectionKind Kind, Mangler &Mang,
     const TargetMachine &TM, bool EmitUniqueSection, unsigned Flags,
@@ -859,9 +872,16 @@ MCSection *TargetLoweringObjectFileELF::getUniqueSectionForFunction(
     const Function &F, const TargetMachine &TM) const {
   SectionKind Kind = SectionKind::getText();
   unsigned Flags = getELFSectionFlags(Kind);
-  return selectELFSectionForGlobal(
-      getContext(), &F, Kind, getMangler(), TM, Used.count(&F),
-      /*EmitUniqueSection=*/true, Flags, &NextUniqueID);
+  // If the function's section names is pre-determined via pragma or a
+  // section attribute, call selectExplicitSectionGlobal.
+  if (F.hasSection() || F.hasFnAttribute("implicit-section-name"))
+    return selectExplicitSectionGlobal(
+        &F, Kind, TM, getContext(), getMangler(), NextUniqueID,
+        Used.count(&F), /* ForceUnique = */true);
+  else
+    return selectELFSectionForGlobal(
+        getContext(), &F, Kind, getMangler(), TM, Used.count(&F),
+        /*EmitUniqueSection=*/true, Flags, &NextUniqueID);
 }
 
 MCSection *TargetLoweringObjectFileELF::getSectionForJumpTable(

diff --git a/llvm/test/CodeGen/X86/basic-block-sections-named-section.ll b/llvm/test/CodeGen/X86/basic-block-sections-named-section.ll
new file mode 100644 (file)
index 0000000..213b1a0
--- /dev/null
+++ b/llvm/test/CodeGen/X86/basic-block-sections-named-section.ll
@@ -0,0 +1,42 @@
+; 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.order.txt
+; RUN: echo "!!2" >> %t.order.txt
+; RUN: llc < %s -mtriple=x86_64-pc-linux -basic-block-sections=%t.order.txt | FileCheck %s --check-prefix=LIST
+; RUN: llc < %s -mtriple=x86_64-pc-linux -function-sections -basic-block-sections=%t.order.txt | FileCheck %s --check-prefix=LIST
+
+; CHECK: .section      foo_section,"ax",@progbits,unique,1
+; CHECK-LABEL: _Z3fooi:
+; CHECK: .section      foo_section,"ax",@progbits,unique,2
+; CHECK-NEXT: _Z3fooi.__part.1:
+; 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
+
+;; Source to generate the IR:
+;; __attribute__((section("foo_section")))
+;; int foo(int n) {
+;;   if (n < 0)
+;;     exit(-1);
+;;   return 0;
+;; }
+
+define dso_local i32 @_Z3fooi(i32 %n) local_unnamed_addr section "foo_section" {
+entry:
+  %cmp = icmp slt i32 %n, 0
+  br i1 %cmp, label %if.then, label %if.end
+
+if.then:                                          ; preds = %entry
+  tail call void @exit(i32 -1) #2
+  unreachable
+
+if.end:                                           ; preds = %entry
+  ret i32 0
+}
+
+declare dso_local void @exit(i32) local_unnamed_addr

diff --git a/llvm/test/CodeGen/X86/basic-block-sections-pragma-sections.ll b/llvm/test/CodeGen/X86/basic-block-sections-pragma-sections.ll
new file mode 100644 (file)
index 0000000..d63fbdd
--- /dev/null
+++ b/llvm/test/CodeGen/X86/basic-block-sections-pragma-sections.ll
@@ -0,0 +1,44 @@
+; 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
+
+; CHECK: .section      foo_section,"ax",@progbits,unique,1
+; CHECK-LABEL: _Z3fooi:
+; CHECK: .section      foo_section,"ax",@progbits,unique,2
+; CHECK-NEXT: _Z3fooi.__part.1:
+; 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
+
+;; Source to generate the IR:
+;; #pragma clang section text = "foo_section"
+;; int foo(int n) {
+;;   if (n < 0)
+;;     exit(-1);
+;;   return 0;
+;; }
+
+define dso_local i32 @_Z3fooi(i32 %n) local_unnamed_addr #0 {
+entry:
+  %cmp = icmp slt i32 %n, 0
+  br i1 %cmp, label %if.then, label %if.end
+
+if.then:                                          ; preds = %entry
+  tail call void @exit(i32 -1) #2
+  unreachable
+
+if.end:                                           ; preds = %entry
+  ret i32 0
+}
+
+declare dso_local void @exit(i32) local_unnamed_addr
+
+attributes #0 = {"implicit-section-name"="foo_section" }