[BasicBlockSections] Make sure that the labels for address-taken blocks are emitted after switching the seciton.

[rahmanl]

Summary

Currently, AsmPrinter code is organized in a way in which the labels of address-taken blocks are emitted in the previous section, which makes the relocation incorrect.
This patch reorganizes the code to switch to the basic block section before handling address-taken blocks.

[snehasish]

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

3089

I would like a comment which explains the condition but this comment isn't very helpful. Consider rewording or removing.

3107

Wasn't there some concern about empty blocks being present which might confound this check?

llvm/test/CodeGen/X86/machine-function-splitter-blockaddress.ll

3

I think this should be part of basic block sections. MFS is using basic block sections and this addresses the underlying issue.

[tmsriram]

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

3099

Could you elaborate on what you mean by label is emitted prior to switching the section? Here it shows as label is emitted only after switching the section. Maybe a simple example would help.

[rahmanl]

• Change the test to -basic-block-sections.

[rahmanl]

Thanks for the comments @tmsriram and @snehasish.

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

3089

Thanks for the comment. If you look at D73674, this exact comment was initially there but we removed it because we wanted the same if-else block to serve other purposes as well. This patch kinds of reverts that change.

3099

It's actually about the ".Ltmp" label emitted in the "MBB.hasAddressTaken()" if block. This patch moves the section switching code before that block. One thing I can suggest is to run the included test on trunk and see the failure.

3107

It is OK for empty basic blocks to begin a section.

llvm/test/CodeGen/X86/machine-function-splitter-blockaddress.ll

3

Yes. Changed to a basic-block-sections test.

[rahmanl]

ping.

[snehasish]

lgtm with some minor comments.

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

3051

Can you extend the comment here to state that this needs to happen before the if (MBB.hasAddressTaken()) part below?

llvm/test/CodeGen/X86/basic-block-sections-blockaddress.ll

1 ↗

(On Diff #295180)

nit: perhaps a better name for this file is "basic-block-sections-address-taken.ll"

28 ↗

(On Diff #295180)

I think you meant to only indent the instructions, the labels should be indented less here and above?

[MaskRay]

llvm/test/CodeGen/X86/basic-block-sections-blockaddress.ll

28 ↗

(On Diff #295180)

The excess indentation is here to satisfy the longest directive CHECK-LABEL.

[snehasish]

llvm/test/CodeGen/X86/basic-block-sections-blockaddress.ll

28 ↗

(On Diff #295180)

Interesting, on my end I see that the indentation lines it up with the instructions below. Is this a phabricator quirk?
https://paste.pics/ABM6T

[rahmanl]

• Rename test, and fix indentation.

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

3051

Thanks for the comment, but I am not sure if it's necessary. (Likewise the alignment code above must be before switching the section, but we don't explicitly mention that).

llvm/test/CodeGen/X86/basic-block-sections-blockaddress.ll

1 ↗

(On Diff #295180)

Changed the test name to basic-block-sections-blockaddress-taken.ll.

28 ↗

(On Diff #295180)

The indentation for .Ltmp1: was incorrect. I fixed it.

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

if (MBB.isEHFuncletEntry()) {
     }
   }
 
   // Emit an alignment directive for this block, if needed.
   const Align Alignment = MBB.getAlignment();
   if (Alignment != Align(1))
     emitAlignment(Alignment);
 
+  // Switch to a new section if this basic block must begin a section. The
+  // entry block is always placed in the function section and is not handled
+  // here.
+  if (MBB.isBeginSection() && !MBB.pred_empty()) {
+    OutStreamer->SwitchSection(
+        getObjFileLowering().getSectionForMachineBasicBlock(MF->getFunction(),
+                                                            MBB, TM));
+    CurrentSectionBeginSym = MBB.getSymbol();
+  }
+
   // If the block has its address taken, emit any labels that were used to
   // reference the block.  It is possible that there is more than one label
   // here, because multiple LLVM BB's may have been RAUW'd to this block after
   // the references were generated.
   if (MBB.hasAddressTaken()) {
     const BasicBlock *BB = MBB.getBasicBlock();
     if (isVerbose())
       OutStreamer->AddComment("Block address taken");
if (const BasicBlock *BB = MBB.getBasicBlock()) {
         OutStreamer->GetCommentOS() << '\n';
       }
     }
 
     assert(MLI != nullptr && "MachineLoopInfo should has been computed");
     emitBasicBlockLoopComments(MBB, MLI, *this);
   }
 
+  // Emit label of the basic block if needed.
   if (MBB.pred_empty() ||
       (!MF->hasBBLabels() && isBlockOnlyReachableByFallthrough(&MBB) &&
        !MBB.isEHFuncletEntry() && !MBB.hasLabelMustBeEmitted())) {
     if (isVerbose()) {
       // NOTE: Want this comment at start of line, don't emit with AddComment.
       OutStreamer->emitRawComment(" %bb." + Twine(MBB.getNumber()) + ":",
                                   false);
     }
   } else {
-    if (isVerbose() && MBB.hasLabelMustBeEmitted()) {
+    if (isVerbose() && MBB.hasLabelMustBeEmitted())
       OutStreamer->AddComment("Label of block must be emitted");
-    }
-    auto *BBSymbol = MBB.getSymbol();
-    // Switch to a new section if this basic block must begin a section.
-    if (MBB.isBeginSection()) {
-      OutStreamer->SwitchSection(
-          getObjFileLowering().getSectionForMachineBasicBlock(MF->getFunction(),
-                                                              MBB, TM));
-      CurrentSectionBeginSym = BBSymbol;
-    }
-    OutStreamer->emitLabel(BBSymbol);
-    // With BB sections, each basic block must handle CFI information on its own
-    // if it begins a section.
-    if (MBB.isBeginSection())
-      for (const HandlerInfo &HI : Handlers)
-        HI.Handler->beginBasicBlock(MBB);
-  }
-}
+    OutStreamer->emitLabel(MBB.getSymbol());
+
+  // if it begins a section (The entry block is handled separately by
+  // AsmPrinterHandler::beginFunction).
+  if (MBB.isBeginSection() && !MBB.pred_empty())
 
 void AsmPrinter::emitBasicBlockEnd(const MachineBasicBlock &MBB) {
   // Check if CFI information needs to be updated for this MBB with basic block
   // sections.
   if (MBB.isEndSection())
     for (const HandlerInfo &HI : Handlers)
       HI.Handler->endBasicBlock(MBB);
 }

llvm/test/CodeGen/X86/machine-function-splitter-blockaddress.ll

+; RUN: llc < %s -mtriple=x86_64-unknown-linux-gnu -split-machine-functions | FileCheck %s
+
+define void @foo1(i1 zeroext %0) nounwind !prof !0 !section_prefix !1 {
+; CHECK:         .text
+; CHECK-LABEL: foo1:
+; CHECK:         movl    $.Ltmp0, %eax
+; CHECK-NEXT:    movl    $.Ltmp1, %ecx
+; CHECK-NEXT:    cmovneq %rax, %rcx
+; CHECK-NEXT:    jmpq    *%rcx
+; CHECK-NEXT:  .Ltmp0:
+; CHECK-NEXT:  .LBB0_1:
+; CHECK-NEXT:    callq   bar
+; CHECK:         .section        .text.unlikely.foo1
+;; Verify that the cold address-taken block and its labels are moved to .text.unlikely.
+; CHECK-NEXT:    .Ltmp1:
+; CHECK-NEXT:  foo1.cold:
+; CHECK-NEXT:    callq   baz
+entry:
+  %1 = select i1 %0, i8* blockaddress(@foo1, %bb1), i8* blockaddress(@foo1, %bb2) ; <i8*> [#uses=1]
+  indirectbr i8* %1, [label %bb1, label %bb2], !prof !2
+
+bb1:                                                ; preds = %entry
+  %2 = call i32 @bar()
+  br label %bb3
+
+bb2:                                                ; preds = %entry
+  %3 = call i32 @baz()
+  br label %bb3
+
+bb3:                                                ; preds = %bb2, %bb1
+  %4 = tail call i32 @qux()
+  ret void
+}
+
+declare i32 @bar()
+declare i32 @baz()
+declare i32 @qux()
+
+!0 = !{!"function_entry_count", i64 7000}
+!1 = !{!"function_section_prefix", !".hot"}
+!2 = !{!"branch_weights", i32 7000, i32 0}