[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

3090

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

3108

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 ↗

(On Diff #295133)

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

3100

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

3090

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.

3100

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.

3108

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

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

3 ↗

(On Diff #295133)

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

[rahmanl]

ping.

[snehasish]

lgtm with some minor comments.

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

3052

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

3052

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 handled
+  // separately.
+  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);
   }
 
+  // Print the main label for the block.
   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 (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/basic-block-sections-blockaddress-taken.ll

+;; This test verifies that basic-block-sections works with address-taken basic blocks.
+; RUN: llc < %s -mtriple=x86_64 -basic-block-sections=all | FileCheck %s
+
+define void @foo(i1 zeroext %0) nounwind {
+entry:
+  %1 = select i1 %0, i8* blockaddress(@foo, %bb1), i8* blockaddress(@foo, %bb2) ; <i8*> [#uses=1]
+  indirectbr i8* %1, [label %bb1, label %bb2]
+
+; CHECK:         .text
+; CHECK-LABEL: foo:
+; CHECK:         movl    $.Ltmp0, %eax
+; CHECK-NEXT:    movl    $.Ltmp1, %ecx
+; CHECK-NEXT:    cmovneq %rax, %rcx
+; CHECK-NEXT:    jmpq    *%rcx
+
+bb1:                                                ; preds = %entry
+  %2 = call i32 @bar()
+  ret void
+; CHECK:         .section .text,"ax",@progbits,unique,1
+; CHECK-NEXT:  .Ltmp0:
+; CHECK-NEXT:  foo.1
+; CHECK-NEXT:    callq   bar
+;
+
+bb2:                                                ; preds = %entry
+  %3 = call i32 @baz()
+  ret void
+; CHECK:         .section .text,"ax",@progbits,unique,2
+; CHECK-NEXT:  .Ltmp1:
+; CHECK-NEXT:  foo.2
+; CHECK-NEXT:    callq   baz
+}
+
+declare i32 @bar()
+declare i32 @baz()