Specify branch probabilities for callbr dests

Summary:
callbr's indirect branches aren't expected to be taken, so reduce their
probabilities to 0 while increasing the default destination to 1. This
allows some code improvements through block placement.

Reviewers: nickdesaulniers

Subscribers: hiraditya, llvm-commits

Tags: #llvm

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

Author: isanbard

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

@@ -2867,10 +2867,10 @@ void SelectionDAGBuilder::visitCallBr(const CallBrInst &I) {
   Return->setInlineAsmBrDefaultTarget();
 
   // Update successor info.
-  addSuccessorWithProb(CallBrMBB, Return);
+  addSuccessorWithProb(CallBrMBB, Return, BranchProbability::getOne());
   for (unsigned i = 0, e = I.getNumIndirectDests(); i < e; ++i) {
     MachineBasicBlock *Target = FuncInfo.MBBMap[I.getIndirectDest(i)];
-    addSuccessorWithProb(CallBrMBB, Target);
+    addSuccessorWithProb(CallBrMBB, Target, BranchProbability::getZero());
     CallBrMBB->addInlineAsmBrIndirectTarget(Target);
   }
   CallBrMBB->normalizeSuccProbs();

llvm/test/CodeGen/AArch64/callbr-asm-label.ll

@@ -6,18 +6,18 @@ define i32 @test1() {
 ; CHECK-LABEL: test1:
 ; CHECK:         .word b
 ; CHECK-NEXT:    .word .Ltmp0
+; CHECK-LABEL: .LBB0_1: // %cleanup
 ; CHECK-LABEL: .Ltmp0:
-; CHECK-LABEL: .LBB0_1: // %l_yes
-; CHECK-LABEL: .LBB0_2: // %cleanup
+; CHECK-LABEL: .LBB0_2: // %indirect
 entry:
-  callbr void asm sideeffect "1:\0A\09.word b, ${0:l}\0A\09", "X"(i8* blockaddress(@test1, %l_yes))
-          to label %cleanup [label %l_yes]
+  callbr void asm sideeffect "1:\0A\09.word b, ${0:l}\0A\09", "X"(i8* blockaddress(@test1, %indirect))
+          to label %cleanup [label %indirect]
 
-l_yes:
+indirect:
   br label %cleanup
 
 cleanup:
-  %retval.0 = phi i32 [ 1, %l_yes ], [ 0, %entry ]
+  %retval.0 = phi i32 [ 1, %indirect ], [ 0, %entry ]
   ret i32 %retval.0
 }
 

llvm/test/CodeGen/AArch64/callbr-asm-obj-file.ll

@@ -41,10 +41,10 @@ declare dso_local i32 @g(...) local_unnamed_addr
 declare dso_local i32 @i(...) local_unnamed_addr
 
 ; CHECK-LABEL: <test2>:
-; CHECK:         bl #0 <test2+0x10>
+; CHECK:         bl #0 <test2+0x18>
 ; CHECK-LABEL: <$d.5>:
 ; CHECK-LABEL: <$x.6>:
-; CHECK-NEXT:    b #16 <$x.8+0x4>
+; CHECK-NEXT:    b #-20 <test2+0x18>
 define hidden i32 @test2() local_unnamed_addr {
   %1 = load i32, i32* @l, align 4
   %2 = icmp eq i32 %1, 0

llvm/test/CodeGen/X86/callbr-asm-branch-folding.ll

@@ -24,17 +24,8 @@ define void @n(i32* %o, i32 %p, i32 %u) nounwind {
 ; CHECK-NEXT:    movq %r15, %rdi
 ; CHECK-NEXT:    callq l
 ; CHECK-NEXT:    testl %eax, %eax
-; CHECK-NEXT:    je .LBB0_1
-; CHECK-NEXT:  .LBB0_10: # %cleanup
-; CHECK-NEXT:    addq $8, %rsp
-; CHECK-NEXT:    popq %rbx
-; CHECK-NEXT:    popq %r12
-; CHECK-NEXT:    popq %r13
-; CHECK-NEXT:    popq %r14
-; CHECK-NEXT:    popq %r15
-; CHECK-NEXT:    popq %rbp
-; CHECK-NEXT:    retq
-; CHECK-NEXT:  .LBB0_1: # %if.end
+; CHECK-NEXT:    jne .LBB0_10
+; CHECK-NEXT:  # %bb.1: # %if.end
 ; CHECK-NEXT:    movl %ebx, {{[-0-9]+}}(%r{{[sb]}}p) # 4-byte Spill
 ; CHECK-NEXT:    cmpl $0, {{.*}}(%rip)
 ; CHECK-NEXT:    # implicit-def: $ebx
@@ -82,7 +73,15 @@ define void @n(i32* %o, i32 %p, i32 %u) nounwind {
 ; CHECK-NEXT:    jmp k # TAILCALL
 ; CHECK-NEXT:  .LBB0_9: # %if.else
 ; CHECK-NEXT:    incq 0
-; CHECK-NEXT:    jmp .LBB0_10
+; CHECK-NEXT:  .LBB0_10: # %cleanup
+; CHECK-NEXT:    addq $8, %rsp
+; CHECK-NEXT:    popq %rbx
+; CHECK-NEXT:    popq %r12
+; CHECK-NEXT:    popq %r13
+; CHECK-NEXT:    popq %r14
+; CHECK-NEXT:    popq %r15
+; CHECK-NEXT:    popq %rbp
+; CHECK-NEXT:    retq
 entry:
   %call = tail call i32 @c()
   %call1 = tail call i32 @l(i32* %o)

llvm/test/CodeGen/X86/callbr-asm.ll

@@ -56,7 +56,7 @@ fail:
 ; Test 3 - asm-goto implements a loop. The loop gets recognized, but many loop
 ; transforms fail due to canonicalization having callbr exceptions. Trivial
 ; blocks at labels 1 and 3 also don't get simplified due to callbr.
-define dso_local i32 @test3(i32 %a) {
+define i32 @test3(i32 %a) {
 ; CHECK-LABEL: test3:
 ; CHECK:       # %bb.0: # %entry
 ; CHECK-NEXT:  .Ltmp1: # Block address taken
@@ -78,7 +78,6 @@ define dso_local i32 @test3(i32 %a) {
 ; CHECK-NEXT:    # Parent Loop BB2_2 Depth=2
 ; CHECK-NEXT:    # => This Loop Header: Depth=3
 ; CHECK-NEXT:    # Child Loop BB2_4 Depth 4
-; CHECK-NEXT:    .p2align 4, 0x90
 ; CHECK-NEXT:  .Ltmp4: # Block address taken
 ; CHECK-NEXT:  .LBB2_4: # %label04
 ; CHECK-NEXT:    # Parent Loop BB2_1 Depth=1
@@ -134,17 +133,17 @@ normal1:                                          ; preds = %normal0
 ; Test 4 - asm-goto referenced with the 'l' (ell) modifier and not.
 define void @test4() {
 ; CHECK-LABEL: test4:
-; CHECK:       # %bb.0:
+; CHECK:       # %bb.0: # %entry
 ; CHECK-NEXT:    #APP
-; CHECK-NOT:     ja .Ltmp50
-; CHECK-NEXT:    ja .Ltmp5
+; CHECK-NEXT:    ja .Ltmp5{{$}}
 ; CHECK-NEXT:    #NO_APP
-; CHECK-NEXT:  .LBB3_1:
+; CHECK-NEXT:  .LBB3_1: # %asm.fallthrough
 ; CHECK-NEXT:    #APP
-; CHECK-NOT:     ja .Ltmp50
-; CHECK-NEXT:    ja .Ltmp5
+; CHECK-NEXT:    ja .Ltmp5{{$}}
 ; CHECK-NEXT:    #NO_APP
-; CHECK-NEXT:    jmp .LBB3_3
+; CHECK-NEXT:  .Ltmp5: # Block address taken
+; CHECK-NEXT:  .LBB3_3: # %quux
+; CHECK-NEXT:    retl
 entry:
   callbr void asm sideeffect "ja $0", "X,~{dirflag},~{fpsr},~{flags}"(i8* blockaddress(@test4, %quux))
           to label %asm.fallthrough [label %quux]