[PartialInline] Bail out on asm-goto/callbr

[wenlei]

Summary

Fixing ICE when partial inline tries to deal with blockaddress uses of function which is typical for asm-goto/callbr. We ran into this with PGO multi-region partial inline.

[davidxl]

llvm/lib/Transforms/IPO/PartialInlining.cpp

1420

Is it legal to partial inline in case when the function's block address is taken?

[hoy]

llvm/lib/Transforms/IPO/PartialInlining.cpp

974

Moving the check into getSupportedCallBase? The assertion llvm_unreachable("All uses must be calls"); there seems too strong.

1420

I think so. The non-inlined callsites will still call the original version, while the inlined callsites will pull in the new version. The new version has some context-irrelevant code outlined, but both versions should be semantically equal. Please correct me if I'm missing anything?

[davidxl]

llvm/lib/Transforms/IPO/PartialInlining.cpp

1420

what if the labels and references from asm goto are in different regions -- one in inlined region, and the other in the outllined region?

[hoy]

llvm/lib/Transforms/IPO/PartialInlining.cpp

1420

If internal blocks of an inlinee are address taken, the inlining should be disabled during the cost analysis: https://github.com/llvm/llvm-project/blob/main/llvm/lib/Analysis/InlineCost.cpp#L2656

[wenlei]

llvm/lib/Transforms/IPO/PartialInlining.cpp

974

Thought about that but didn't do that for two reasons:

1. The API assume we would get one call for each User, but BlockAddress is a constant, so it doesn't have 1:1 mapping even for CallBr.
2. The contract of assuming valid Call makes sense for other use cases, e.g. getOneCallSiteTo. If we return null for all BlockAddress in getSupportedCallBase, there will be more error handling for getOneCallSiteTo and its caller, which is less clean.

So I choose to make changes on the caller side so that the assumption holds for getSupportedCallBase.

1420

what if the labels and references from asm goto are in different regions -- one in inlined region, and the other in the outllined region?

Is this concerning side entrance into the outlined region through asm goto? The outlining for partial inline requires single entry single exit region, so asm goto should not introduce side entrance.

If internal blocks of an inlinee are address taken, the inlining should be disabled during the cost analysis: https://github.com/llvm/llvm-project/blob/main/llvm/lib/Analysis/InlineCost.cpp#L2656

The code you linked still allow callbr/asm-goto, it disables inlining if address taken blocks are used outside of asm-goto.

[davidxl]

lgtm

llvm/lib/Transforms/IPO/PartialInlining.cpp

1420

In general, the legality check should be done earlier otherwise we may end up with situation where function outlining happens but partial inlining is disabled -- resulting in less efficient code.

However as Wenlei mentioned, cross region jumps may not be an issue.

[hoy]

llvm/lib/Transforms/IPO/PartialInlining.cpp

974

Sounds good.

1420

The code you linked still allow callbr/asm-goto, it disables inlining if address taken blocks are used outside of asm-goto.

Is asm-goto equivalent to callbr? I thought callbr can only deal with function entry but not internal blocks.

[wenlei]

llvm/lib/Transforms/IPO/PartialInlining.cpp

1420

The code you linked still allow callbr/asm-goto, it disables inlining if address taken blocks are used outside of asm-goto.

Is asm-goto equivalent to callbr? I thought callbr can only deal with function entry but not internal blocks.

The only use of callbr is for asm-goto, so yes in practice they're equivalent.

But the code you linked only bails out when the use of block address is *not* callbr. So it doesn't block callbr, which means this statement below isn't true.

If internal blocks of an inlinee are address taken, the inlining should be disabled during the cost analysis

if (BB->hasAddressTaken())
      for (User *U : BlockAddress::get(&*BB)->users())
        if (!isa<CallBrInst>(*U))
          return InlineResult::failure("blockaddress used outside of callbr");

llvm/lib/Transforms/IPO/PartialInlining.cpp

auto ComputeCurrBFI = [&,this](Function *Caller) {
         CurrentCallerBFI = TempBFI.get();
       } else {
         // New pass manager:
         CurrentCallerBFI = &(GetBFI(*Caller));
       }
   };
 
   for (User *User : Users) {
+    // Don't bother with BlockAddress used by CallBr for asm goto.
+    if (isa<BlockAddress>(User))
+      continue;
     CallBase *CB = getSupportedCallBase(User);
     Function *Caller = CB->getCaller();
     if (CurrentCaller != Caller) {
       CurrentCaller = Caller;
       ComputeCurrBFI(Caller);
     } else {
       assert(CurrentCallerBFI && "CallerBFI is not set");
     }
bool PartialInlinerImpl::tryPartialInline(FunctionCloner &Cloner) {
   if (CalleeEntryCount)
     computeCallsiteToProfCountMap(Cloner.ClonedFunc, CallSiteToProfCountMap);
 
   uint64_t CalleeEntryCountV =
       (CalleeEntryCount ? CalleeEntryCount->getCount() : 0);
 
   bool AnyInline = false;
   for (User *User : Users) {
+    // Don't bother with BlockAddress used by CallBr for asm goto.
+    if (isa<BlockAddress>(User))
+      continue;
+
     CallBase *CB = getSupportedCallBase(User);
 
     if (isLimitReached())
       continue;
 
     OptimizationRemarkEmitter CallerORE(CB->getCaller());
     if (!shouldPartialInline(*CB, Cloner, WeightedRcost, CallerORE))
       continue;

llvm/test/Transforms/CodeExtractor/PartialInlinePGOMultiRegion.ll

 
 for.inc2:                                          ; preds = %if.end2
   %19 = load i32, i32* %i, align 4
   %inc2 = add nsw i32 %19, 1
   store i32 %inc2, i32* %i, align 4
   br label %for.cond2
 
 for.end:                                          ; preds = %for.cond2
+  callbr void asm sideeffect "1: nop\0A\09.quad b, ${0:l}, $$5\0A\09", "X,~{dirflag},~{fpsr},~{flags}"(i8* blockaddress(@bar, %l_yes))
+          to label %asm.fallthrough [label %l_yes]
+asm.fallthrough:                                  ; preds = %for.end
+  br label %l_yes
+
+l_yes:
   %20 = load i32, i32* %sum, align 4
   ret i32 %20
 }
 
 define signext i32 @foo(i32 signext %value, i32 signext %ub) #0 !prof !30 {
 ; CHECK-LABEL: @foo
 ; CHECK-NOT: call signext i32 @bar
 ; CHECK: codeRepl1.i: