[CodeGen][BasicBlockSections] Fix wrong alignment directive placement in basic block section cases

[sinan]

Summary

MachineBlockPlacement pass sets an alignment attribute to the loop header MBB and this attribute will lead to an alignment directive during emitting asm. In the case of the basic block section, the alignment directive is put before the section label, and thus the alignment is set to the predecessor of the loop header, which is not what we expect and increases the code size (both inserting nop and set section alignment).

e.g. reify.__part.2 is the loop header and .p2align 4, 0x90 should be moved below .section .text.reify.reify.__part.2.

	.section	.text.reify.reify.__part.1,"ax",@progbits
reify.__part.1:                         # %if.end
	.cfi_startproc
	.cfi_def_cfa %rsp, 8
	movq	(%rdi), %rcx
	movq	24(%rcx), %rdx
	leaq	1(%rdx), %rax
.Ltmp1:
	#DEBUG_VALUE: reify:key <- $rax
	cmpq	16(%rcx), %rdx
	jle	reify.__part.3
	jmp	reify.__part.2
.LBB_END0_1:
	.size	reify.__part.1, .LBB_END0_1-reify.__part.1
	.cfi_endproc
	.p2align	4, 0x90
	.section	.text.reify.reify.__part.2,"ax",@progbits
        # .p2align is expected to be put here!
reify.__part.2:                         # %while.body
                                        # =>This Inner Loop Header: Depth=1
	.cfi_startproc
	.cfi_def_cfa %rsp, 8

[Nr] Name                             Type                   Address                       Off         Size      ES   Flg Lk Inf Al
[ 3] .text.reify                        PROGBITS        0000000000000000 000040 00000f 00  AX  0   0 16
[ 5] .text.reify.reify.__part.1 PROGBITS        0000000000000000 000050 000020 00  AX  0   0 16
[ 7] .text.reify.reify.__part.2 PROGBITS        0000000000000000 000070 000024 00  AX  0   0  1

[rahmanl]

Thanks for this fix.

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

1 ↗

(On Diff #473583)

Can you please simplify the test and check the assembly only to make sure the alignment directive is placed at the right position?
You might want to use -align-all-blocks or -align-all-nofallthru-blocks and then use a simple if-then-else statement to trigger alignment.

[sinan]

update the test case. This test case is generated after basic-block-sections pass and the name mapping between the basic block(bb.2.entry (align 256, bbsections 2)) and the label name(test.__part.2) from the assembly can be clearer in this way.

[rahmanl]

Thank you.

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

void AsmPrinter::emitBasicBlockStart(const MachineBasicBlock &MBB) {
   // End the previous funclet and start a new one.
   if (MBB.isEHFuncletEntry()) {
     for (const HandlerInfo &HI : Handlers) {
       HI.Handler->endFunclet();
       HI.Handler->beginFunclet(MBB);
     }
   }
 
-  // Emit an alignment directive for this block, if needed.
-  const Align Alignment = MBB.getAlignment();
-  if (Alignment != Align(1))
-    emitAlignment(Alignment, nullptr, MBB.getMaxBytesForAlignment());
-
   // 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.isEntryBlock()) {
     OutStreamer->switchSection(
         getObjFileLowering().getSectionForMachineBasicBlock(MF->getFunction(),
                                                             MBB, TM));
     CurrentSectionBeginSym = MBB.getSymbol();
   }
 
+  // Emit an alignment directive for this block, if needed.
+  const Align Alignment = MBB.getAlignment();
+  if (Alignment != Align(1))
+    emitAlignment(Alignment, nullptr, MBB.getMaxBytesForAlignment());
+
   // 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.isIRBlockAddressTaken()) {
     if (isVerbose())
       OutStreamer->AddComment("Block address taken");
 

llvm/test/CodeGen/X86/align-basic-block-sections.mir

+# Check if the alignment directive is put on the correct place when the basic block section option is used.
+# RUN: llc -start-after=bbsections-prepare  %s -o - | FileCheck %s -check-prefix=CHECK
+
+# How to generate the input:
+# foo.c
+# int test(int a) {
+#   switch (a) {
+#     default:
+#       return 10;
+#     case 1:
+#       a += 1;
+#     case 2:
+#       a *= -1;
+#     case 8:
+#       break;
+#   }
+#   return a;
+# }
+#
+# clang -O0 -S -emit-llvm test.c
+# llc < test.ll -stop-after=bbsections-prepare -align-all-nofallthru-blocks=8 -basic-block-sections=all
+
+
+--- |
+  define i32 @test(i32 noundef %a) {
+  entry:
+    switch i32 %a, label %return [
+      i32 1, label %sw.bb1
+      i32 2, label %sw.bb1
+      i32 8, label %sw.epilog
+    ]
+
+  sw.bb1:                                           ; preds = %entry, %entry
+    br label %sw.epilog
+
+  sw.epilog:                                        ; preds = %sw.bb1, %entry
+    %a.addr.1 = phi i32 [ %a, %entry ], [ -2, %sw.bb1 ]
+    br label %return
+
+  return:                                           ; preds = %sw.epilog, %entry
+    %retval.0 = phi i32 [ %a.addr.1, %sw.epilog ], [ 10, %entry ]
+    ret i32 %retval.0
+  }
+
+
+...
+---
+name:            test
+alignment:       16
+exposesReturnsTwice: false
+legalized:       false
+regBankSelected: false
+selected:        false
+failedISel:      false
+tracksRegLiveness: true
+hasWinCFI:       false
+callsEHReturn:   false
+callsUnwindInit: false
+hasEHCatchret:   false
+hasEHScopes:     false
+hasEHFunclets:   false
+failsVerification: false
+tracksDebugUserValues: true
+registers:       []
+liveins:
+  - { reg: '$edi', virtual-reg: '' }
+frameInfo:
+  isFrameAddressTaken: false
+  isReturnAddressTaken: false
+  hasStackMap:     false
+  hasPatchPoint:   false
+  stackSize:       0
+  offsetAdjustment: 0
+  maxAlignment:    1
+  adjustsStack:    false
+  hasCalls:        false
+  stackProtector:  ''
+  functionContext: ''
+  maxCallFrameSize: 0
+  cvBytesOfCalleeSavedRegisters: 0
+  hasOpaqueSPAdjustment: false
+  hasVAStart:      false
+  hasMustTailInVarArgFunc: false
+  hasTailCall:     false
+  localFrameSize:  0
+  savePoint:       ''
+  restorePoint:    ''
+fixedStack:      []
+stack:           []
+callSites:       []
+debugValueSubstitutions: []
+constants:       []
+machineFunctionInfo: {}
+body:             |
+  bb.0.entry:
+    successors: %bb.1(0x40000000), %bb.2(0x40000000)
+    liveins: $edi
+  
+    renamable $edi = KILL $edi, implicit-def $rdi
+    renamable $eax = LEA64_32r renamable $rdi, 1, $noreg, -1, $noreg
+    CMP32ri8 killed renamable $eax, 2, implicit-def $eflags
+    JCC_1 %bb.2, 3, implicit $eflags
+    JMP_1 %bb.1
+  
+  bb.1.sw.bb1 (bbsections 1):
+    successors: %bb.3(0x80000000)
+  
+    renamable $edi = MOV32ri -2, implicit-def $rdi
+    JMP_1 %bb.3
+  
+  bb.2.entry (align 256, bbsections 2):
+    successors: %bb.3(0x40000000), %bb.4(0x40000000)
+    liveins: $rdi
+  
+    renamable $eax = MOV32ri 10
+    CMP32ri8 renamable $edi, 8, implicit-def $eflags
+    JCC_1 %bb.4, 5, implicit $eflags
+    JMP_1 %bb.3
+  
+  bb.3.sw.epilog (bbsections 3):
+    successors: %bb.4(0x80000000)
+    liveins: $rdi
+  
+    $eax = MOV32rr $edi, implicit killed $rdi
+    JMP_1 %bb.4
+  
+  bb.4.return (bbsections 4):
+    liveins: $eax
+  
+    RET64 $eax
+
+...
+
+# CHECK: .section .text.test,"ax",@progbits,unique,2
+# CHECK-NEXT: .p2align  8, 0x90
+# CHECK-NEXT: test.__part.2: # %entry