[ThinLTO] Disallow importing for functions with indir branch to block address

[wenlei]

Summary

We don't allowing inlining for functions with blockaddress with uses other than strictly callbr. This is because if the blockaddress escapes the function via a global variable, inlining may lead to an invalid cross-function reference.

We check against such cases during inlining, however the check can fail for ThinLTO post-link because CFG simplification can incorrectly removes blocks based on wrong block reachability.

When we import a function with blockaddress taken in a global variable but without importing that variable, we won't go through value mapping to reflect the real address-taken-ness of the cloned blocks. For the imported clone, this leads to blocks reachable from indirect branch through global variable being incorrectly treated as unreachable and removed by SimplifyCFG.

Since inlining for such cases shouldn't be allowed in the first place, I'm marking them as ineligible for importing during pre-link to save the problem of missing address-taken-ness of imported clone as well as bad DCE and inlining.

[wenlei]

llvm/lib/Analysis/ModuleSummaryAnalysis.cpp

275

This logic is from InlineCost.cpp, but I'm not sure in what case do we have an indirect call to a block label?

[wmi]

llvm/lib/Analysis/ModuleSummaryAnalysis.cpp

275

Here is the only usage of CallBrInst in clang: https://github.com/llvm-mirror/clang/blob/master/lib/CodeGen/CGStmt.cpp#L2282, and it is used for gcc inline asm extension: https://gcc.gnu.org/onlinedocs/gcc/Extended-Asm.html#GotoLabels

[wenlei]

llvm/lib/Analysis/ModuleSummaryAnalysis.cpp

275

Ok, thanks for the pointer. In that case, I think it's indeed safe to allow CallBrInst to block address for inlining and importing.

[tejohnson]

LGTM with a test fix below.

llvm/test/ThinLTO/X86/globals-import-blockaddr.ll

11

This only checks that there isn't an imported copy of bar defined after the imported copy of foo. Either add the same CHECK-NOT prior to the CHECK for foo, or make it a positive sense check for the declare of bar. Probably the latter is cleaner.

[wenlei]

Fix test to use positive check.

[wenlei]

llvm/test/ThinLTO/X86/globals-import-blockaddr.ll

11

Good catch, changed to positive check.

llvm/lib/Analysis/ModuleSummaryAnalysis.cpp

static void computeFunctionSummary(
 
   // Add personality function, prefix data and prologue data to function's ref
   // list.
   findRefEdges(Index, &F, RefEdges, Visited);
   std::vector<const Instruction *> NonVolatileLoads;
   std::vector<const Instruction *> NonVolatileStores;
 
   bool HasInlineAsmMaybeReferencingInternal = false;
-  for (const BasicBlock &BB : F)
+  bool HasIndirBranchToBlockAddress = false;
+  for (const BasicBlock &BB : F) {
+    // We don't allow inlining of function with indirect branch to blockaddress.
+    // If the blockaddress escapes the function, e.g., via a global variable,
+    // inlining may lead to an invalid cross-function reference. So we shouldn't
+    // import such function either.
+    if (BB.hasAddressTaken()) {
+      for (User *U : BlockAddress::get(const_cast<BasicBlock *>(&BB))->users())
+        if (!isa<CallBrInst>(*U)) {
+          HasIndirBranchToBlockAddress = true;
+          break;
+        }
+    }
+
     for (const Instruction &I : BB) {
       if (isa<DbgInfoIntrinsic>(I))
         continue;
       ++NumInsts;
       // Regular LTO module doesn't participate in ThinLTO import,
       // so no reference from it can be read/writeonly, since this
       // would require importing variable as local copy
       if (IsThinLTO) {
for (const Instruction &I : BB) {
         auto CandidateProfileData =
             ICallAnalysis.getPromotionCandidatesForInstruction(
                 &I, NumVals, TotalCount, NumCandidates);
         for (auto &Candidate : CandidateProfileData)
           CallGraphEdges[Index.getOrInsertValueInfo(Candidate.Value)]
               .updateHotness(getHotness(Candidate.Count, PSI));
       }
     }
+  }
   Index.addBlockCount(F.size());
 
   std::vector<ValueInfo> Refs;
   if (IsThinLTO) {
     auto AddRefEdges = [&](const std::vector<const Instruction *> &Instrs,
                            SetVector<ValueInfo> &Edges,
                            SmallPtrSet<const User *, 8> &Cache) {
       for (const auto *I : Instrs) {
static void computeFunctionSummary(
   // sample PGO, to enable the same inlines as the profiled optimized binary.
   for (auto &I : F.getImportGUIDs())
     CallGraphEdges[Index.getOrInsertValueInfo(I)].updateHotness(
         ForceSummaryEdgesCold == FunctionSummary::FSHT_All
             ? CalleeInfo::HotnessType::Cold
             : CalleeInfo::HotnessType::Critical);
 
   bool NonRenamableLocal = isNonRenamableLocal(F);
-  bool NotEligibleForImport =
-      NonRenamableLocal || HasInlineAsmMaybeReferencingInternal;
+  bool NotEligibleForImport = NonRenamableLocal ||
+                              HasInlineAsmMaybeReferencingInternal ||
+                              HasIndirBranchToBlockAddress;
   GlobalValueSummary::GVFlags Flags(
       F.getLinkage(), F.getVisibility(), NotEligibleForImport,
       /* Live = */ false, F.isDSOLocal(),
       F.hasLinkOnceODRLinkage() && F.hasGlobalUnnamedAddr());
   FunctionSummary::FFlags FunFlags{
       F.hasFnAttribute(Attribute::ReadNone),
       F.hasFnAttribute(Attribute::ReadOnly),
       F.hasFnAttribute(Attribute::NoRecurse), F.returnDoesNotAlias(),

llvm/test/ThinLTO/X86/Inputs/globals-import-blockaddr.ll

 target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "x86_64-unknown-linux-gnu"
 
-@label_addr = internal constant [1 x i8*] [i8* blockaddress(@foo, %lb)], align 8
+@label_addr = internal constant [1 x i8*] [i8* blockaddress(@bar, %lb)], align 8
 
 ; Function Attrs: noinline norecurse nounwind optnone uwtable
 define dso_local [1 x i8*]* @foo() {
+  ret [1 x i8*]* @label_addr
+}
+
+; Function Attrs: noinline norecurse nounwind optnone uwtable
+define dso_local [1 x i8*]* @bar() {
   br label %lb
 
 lb:
   ret [1 x i8*]* @label_addr
 }

llvm/test/ThinLTO/X86/globals-import-blockaddr.ll

 ; RUN: opt -module-summary %s -o %t1.bc
 ; RUN: opt -module-summary %p/Inputs/globals-import-blockaddr.ll -o %t2.bc
-; RUN: llvm-lto2 run -save-temps %t1.bc -r=%t1.bc,foo,l -r=%t1.bc,main,pl %t2.bc -r=%t2.bc,foo,pl -o %t3
+; RUN: llvm-lto2 run -save-temps %t1.bc -r=%t1.bc,foo,l -r=%t1.bc,bar,l -r=%t1.bc,main,pl %t2.bc -r=%t2.bc,foo,pl -r=%t2.bc,bar,pl -o %t3
 ; RUN: llvm-dis %t3.1.3.import.bc -o - | FileCheck %s
 
 ; Verify that we haven't imported GV containing blockaddress
 ; CHECK: @label_addr.llvm.0 = external hidden constant
+; Verify that bar is not imported since it has address-taken block that is target of indirect branch
+; CHECK: declare [1 x i8*]* @bar()
+; Verify that foo is imported
+; CHECK: available_externally [1 x i8*]* @foo
 
 target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "x86_64-unknown-linux-gnu"
 
 declare dso_local [1 x i8*]* @foo();
+declare dso_local [1 x i8*]* @bar();
 
 define dso_local i32 @main() {
-  %p = call [1 x i8*]* @foo()
-  %v = ptrtoint [1 x i8*]* %p to i32
-  ret i32 %v
+  %p1 = call [1 x i8*]* @foo()
+  %p2 = call [1 x i8*]* @bar()
+  %v1 = ptrtoint [1 x i8*]* %p1 to i32
+  %v2 = ptrtoint [1 x i8*]* %p2 to i32
+  %v3 = add i32 %v1, %v2
+  ret i32 %v3
 }