In thin and full LTO + CFI, direct function calls may go through jump table

entries to reach the target. Since these calls don't require type checks,
we can short-circuit them to their real targets.

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

llvm-svn: 332610

llvm/include/llvm/IR/Value.h

@@ -299,12 +299,6 @@ public:
   /// values or constant users.
   void replaceUsesOutsideBlock(Value *V, BasicBlock *BB);
 
-  /// replaceUsesExceptBlockAddr - Go through the uses list for this definition
-  /// and make each use point to "V" instead of "this" when the use is outside
-  /// the block. 'This's use list is expected to have at least one element.
-  /// Unlike replaceAllUsesWith this function skips blockaddr uses.
-  void replaceUsesExceptBlockAddr(Value *New);
-
   //----------------------------------------------------------------------
   // Methods for handling the chain of uses of this Value.
   //

llvm/lib/IR/Value.cpp

@@ -465,35 +465,6 @@ void Value::replaceUsesOutsideBlock(Value *New, BasicBlock *BB) {
   }
 }
 
-void Value::replaceUsesExceptBlockAddr(Value *New) {
-  SmallSetVector<Constant *, 4> Constants;
-  use_iterator UI = use_begin(), E = use_end();
-  for (; UI != E;) {
-    Use &U = *UI;
-    ++UI;
-
-    if (isa<BlockAddress>(U.getUser()))
-      continue;
-
-    // Must handle Constants specially, we cannot call replaceUsesOfWith on a
-    // constant because they are uniqued.
-    if (auto *C = dyn_cast<Constant>(U.getUser())) {
-      if (!isa<GlobalValue>(C)) {
-        // Save unique users to avoid processing operand replacement
-        // more than once.
-        Constants.insert(C);
-        continue;
-      }
-    }
-
-    U.set(New);
-  }
-
-  // Process operand replacement of saved constants.
-  for (auto *C : Constants)
-    C->handleOperandChange(this, New);
-}
-
 namespace {
 // Various metrics for how much to strip off of pointers.
 enum PointerStripKind {

llvm/lib/Transforms/IPO/LowerTypeTests.cpp

@@ -429,6 +429,17 @@ class LowerTypeTestsModule {
 
   void createJumpTable(Function *F, ArrayRef<GlobalTypeMember *> Functions);
 
+  /// replaceCfiUses - Go through the uses list for this definition
+  /// and make each use point to "V" instead of "this" when the use is outside
+  /// the block. 'This's use list is expected to have at least one element.
+  /// Unlike replaceAllUsesWith this function skips blockaddr and direct call
+  /// uses.
+  void replaceCfiUses(Value *Old, Value *New);
+
+  /// replaceDirectCalls - Go through the uses list for this definition and
+  /// replace each use, which is a direct function call.
+  void replaceDirectCalls(Value *Old, Value *New);
+
 public:
   LowerTypeTestsModule(Module &M, ModuleSummaryIndex *ExportSummary,
                        const ModuleSummaryIndex *ImportSummary);
@@ -967,14 +978,19 @@ void LowerTypeTestsModule::importTypeTest(CallInst *CI) {
 void LowerTypeTestsModule::importFunction(Function *F, bool isDefinition) {
   assert(F->getType()->getAddressSpace() == 0);
 
-  // Declaration of a local function - nothing to do.
-  if (F->isDeclarationForLinker() && isDefinition)
-    return;
-
   GlobalValue::VisibilityTypes Visibility = F->getVisibility();
   std::string Name = F->getName();
-  Function *FDecl;
 
+  if (F->isDeclarationForLinker() && isDefinition) {
+    Function *RealF = Function::Create(F->getFunctionType(),
+                                       GlobalValue::ExternalLinkage,
+                                       Name + ".cfi", &M);
+    RealF->setVisibility(Visibility);
+    replaceDirectCalls(F, RealF);
+    return;
+  }
+
+  Function *FDecl;
   if (F->isDeclarationForLinker() && !isDefinition) {
     // Declaration of an external function.
     FDecl = Function::Create(F->getFunctionType(), GlobalValue::ExternalLinkage,
@@ -1014,7 +1030,7 @@ void LowerTypeTestsModule::importFunction(Function *F, bool isDefinition) {
   if (F->isWeakForLinker())
     replaceWeakDeclarationWithJumpTablePtr(F, FDecl);
   else
-    F->replaceUsesExceptBlockAddr(FDecl);
+    replaceCfiUses(F, FDecl);
 }
 
 void LowerTypeTestsModule::lowerTypeTestCalls(
@@ -1209,7 +1225,7 @@ void LowerTypeTestsModule::replaceWeakDeclarationWithJumpTablePtr(
   Function *PlaceholderFn =
       Function::Create(cast<FunctionType>(F->getValueType()),
                        GlobalValue::ExternalWeakLinkage, "", &M);
-  F->replaceAllUsesWith(PlaceholderFn);
+  replaceCfiUses(F, PlaceholderFn);
 
   Constant *Target = ConstantExpr::getSelect(
       ConstantExpr::getICmp(CmpInst::ICMP_NE, F,
@@ -1433,7 +1449,7 @@ void LowerTypeTestsModule::buildBitSetsFromFunctionsNative(
       if (F->isWeakForLinker())
         replaceWeakDeclarationWithJumpTablePtr(F, CombinedGlobalElemPtr);
       else
-        F->replaceAllUsesWith(CombinedGlobalElemPtr);
+        replaceCfiUses(F, CombinedGlobalElemPtr);
     } else {
       assert(F->getType()->getAddressSpace() == 0);
 
@@ -1443,10 +1459,8 @@ void LowerTypeTestsModule::buildBitSetsFromFunctionsNative(
       FAlias->takeName(F);
       if (FAlias->hasName())
         F->setName(FAlias->getName() + ".cfi");
-      F->replaceUsesExceptBlockAddr(FAlias);
+      replaceCfiUses(F, FAlias);
     }
-    if (!F->isDeclarationForLinker())
-      F->setLinkage(GlobalValue::InternalLinkage);
   }
 
   createJumpTable(JumpTableFn, Functions);
@@ -1602,6 +1616,63 @@ bool LowerTypeTestsModule::runForTesting(Module &M) {
   return Changed;
 }
 
+static bool isDirectCall(Use& U) {
+  auto *Usr = dyn_cast<CallInst>(U.getUser());
+  if (Usr) {
+    CallSite CS(Usr);
+    if (CS.isCallee(&U))
+      return true;
+  }
+  return false;
+}
+
+void LowerTypeTestsModule::replaceCfiUses(Value *Old, Value *New) {
+  SmallSetVector<Constant *, 4> Constants;
+  auto UI = Old->use_begin(), E = Old->use_end();
+  for (; UI != E;) {
+    Use &U = *UI;
+    ++UI;
+
+    // Skip block addresses
+    if (isa<BlockAddress>(U.getUser()))
+      continue;
+
+    // Skip direct calls
+    if (isDirectCall(U))
+      continue;
+
+    // Must handle Constants specially, we cannot call replaceUsesOfWith on a
+    // constant because they are uniqued.
+    if (auto *C = dyn_cast<Constant>(U.getUser())) {
+      if (!isa<GlobalValue>(C)) {
+        // Save unique users to avoid processing operand replacement
+        // more than once.
+        Constants.insert(C);
+        continue;
+      }
+    }
+
+    U.set(New);
+  }
+
+  // Process operand replacement of saved constants.
+  for (auto *C : Constants)
+    C->handleOperandChange(Old, New);
+}
+
+void LowerTypeTestsModule::replaceDirectCalls(Value *Old, Value *New) {
+  auto UI = Old->use_begin(), E = Old->use_end();
+  for (; UI != E;) {
+    Use &U = *UI;
+    ++UI;
+
+    if (!isDirectCall(U))
+      continue;
+
+    U.set(New);
+  }
+}
+
 bool LowerTypeTestsModule::lower() {
   Function *TypeTestFunc =
       M.getFunction(Intrinsic::getName(Intrinsic::type_test));

llvm/test/Transforms/LowerTypeTests/Inputs/cfi-direct-call.yaml

@@ -0,0 +1,10 @@
+---
+TypeIdMap:
+  typeid1:
+    TTRes:
+      Kind:            AllOnes
+CfiFunctionDefs:
+  - internal_def
+CfiFunctionDecls:
+  - external_decl
+...

llvm/test/Transforms/LowerTypeTests/Inputs/cfi-direct-call1.yaml

@@ -0,0 +1,13 @@
+---
+TypeIdMap:
+  typeid1:
+    TTRes:
+      Kind:            AllOnes
+CfiFunctionDefs:
+  - local_func1
+  - local_func2
+  - local_func3
+CfiFunctionDecls:
+  - extern_decl
+  - extern_weak
+...

llvm/test/Transforms/LowerTypeTests/blockaddress.ll

@@ -1,7 +1,7 @@
 ; RUN: opt -S %s -lowertypetests | FileCheck %s
 
 
-; CHECK: define internal i8* @f2.cfi() !type !0 {
+; CHECK: define i8* @f2.cfi() !type !0 {
 ; CHECK-NEXT:  br label %b
 ; CHECK: b:
 ; CHECK-NEXT:  ret i8* blockaddress(@f2.cfi, %b)

llvm/test/Transforms/LowerTypeTests/cfi-direct-call.ll

@@ -0,0 +1,42 @@
+; RUN: opt -lowertypetests -lowertypetests-summary-action=import -lowertypetests-read-summary=%p/Inputs/cfi-direct-call.yaml %s -S | FileCheck %s
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux"
+
+declare void @external_decl()
+declare void @external_nodecl()
+declare void @internal_def()
+
+define i8 @local_a() {
+  call void @external_decl()
+  call void @external_nodecl()
+  call void @internal_def()
+  call void @local_b()
+  ret i8 1
+}
+
+define void @local_b() {
+  ret void
+}
+
+; CHECK: define i8 @local_a() {
+
+; Even though a jump table entry is generated, the call goes directly
+; to the function
+; CHECK-NEXT:   call void @external_decl()
+
+; External call with no CFI decl - no action
+; CHECK-NEXT:   call void @external_nodecl()
+
+; Internal function defined outside the module generates a jump table
+; entry and is renamed to *.cfi: route direct call to actual function,
+; not jump table
+; CHECK-NEXT:   call void @internal_def.cfi()
+
+; Local call - no action
+; CHECK-NEXT:   call void @local_b
+
+; CHECK-NEXT:   ret i8 1
+
+; CHECK: declare void @internal_def.cfi()
+; CHECK: declare hidden void @external_decl.cfi_jt()

llvm/test/Transforms/LowerTypeTests/cfi-direct-call1.ll

@@ -0,0 +1,96 @@
+; RUN: opt -lowertypetests -S %s | FileCheck --check-prefix=FULL %s
+; RUN: opt -lowertypetests -lowertypetests-summary-action=import -lowertypetests-read-summary=%p/Inputs/cfi-direct-call1.yaml -S %s | FileCheck --check-prefix=THIN %s
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux"
+
+define hidden i32 @local_func1() !type !3 !type !4 {
+  ret i32 11
+}
+
+define hidden i32 @local_func2() !type !3 !type !4 {
+  ret i32 22
+}
+
+define hidden i32 @local_func3(i32 %i) local_unnamed_addr !type !5 !type !6 {
+entry:
+  %add = add nsw i32 %i, 44
+  ret i32 %add
+}
+
+declare !type !3 !type !4 extern_weak i32 @extern_weak()
+declare !type !3 !type !4 i32 @extern_decl()
+declare i1 @llvm.type.test(i8*, metadata)
+
+define hidden i32 @main(i32 %argc) {
+entry:
+  %cmp.i = icmp sgt i32 %argc, 1
+  %fptr1 = select i1 %cmp.i, i32 ()* @local_func1, i32 ()* @local_func2
+  %fptr2 = select i1 %cmp.i, i32 ()* @extern_weak, i32 ()* @extern_decl
+  %0 = bitcast i32 ()* %fptr1 to i8*
+  %1 = tail call i1 @llvm.type.test(i8* nonnull %0, metadata !"_ZTSFivE")
+
+  %call2 = tail call i32 %fptr1()
+  %2 = bitcast i32 ()* %fptr2 to i8*
+  %3 = tail call i1 @llvm.type.test(i8* %2, metadata !"_ZTSFivE")
+
+  %call4 = tail call i32 %fptr2()
+  %call5 = tail call i32 @extern_decl()
+  %call7 = tail call i32 @extern_weak()
+  %call6 = tail call i32 @local_func1()
+  %call8 = tail call i32 @local_func3(i32 4)
+  ret i32 12
+}
+
+!3 = !{i64 0, !"_ZTSFivE"}
+!4 = !{i64 0, !"_ZTSFivE.generalized"}
+!5 = !{i64 0, !"_ZTSFiiE"}
+!6 = !{i64 0, !"_ZTSFiiE.generalized"}
+
+; Make sure local_func1 and local_func2 have been renamed to <name>.cfi
+; FULL: define hidden i32 @local_func1.cfi()
+; FULL: define hidden i32 @local_func2.cfi()
+
+; There are no indirect calls of local_func3 type, it should not be renamed
+; FULL: define hidden i32 @local_func3(i32 %i)
+
+; Indirect references to local_func1 and local_func2 must to through jump table
+; FULL: %fptr1 = select i1 %cmp.i, i32 ()* @local_func1, i32 ()* @local_func2
+
+; Indirect references to extern_weak and extern_decl must go through jump table
+; FULL: %fptr2 = select i1 %cmp.i, i32 ()* select (i1 icmp ne (i32 ()* @extern_weak, i32 ()* null), i32 ()* bitcast ([8 x i8]* getelementptr inbounds ([4 x [8 x i8]], [4 x [8 x i8]]* bitcast (void ()* @.cfi.jumptable to [4 x [8 x i8]]*), i64 0, i64 2) to i32 ()*), i32 ()* null), i32 ()* bitcast ([8 x i8]* getelementptr inbounds ([4 x [8 x i8]], [4 x [8 x i8]]* bitcast (void ()* @.cfi.jumptable to [4 x [8 x i8]]*), i64 0, i64 3) to i32 ()*)
+
+; Direct calls to extern_weak and extern_decl should go to original names
+; FULL: %call5 = tail call i32 @extern_decl()
+; FULL: %call7 = tail call i32 @extern_weak()
+
+; Direct call to local_func1 should to the renamed version
+; FULL: %call6 = tail call i32 @local_func1.cfi()
+
+; local_func3 hasn't been renamed, direct call should go to the original name
+; FULL: %call8 = tail call i32 @local_func3(i32 4)
+
+; Check which jump table entries are created
+; FULL: define private void @.cfi.jumptable(){{.*}}
+; FULL-NEXT: entry:
+; FULL-NEXT: call void asm{{.*}}local_func1.cfi{{.*}}local_func2.cfi{{.*}}extern_weak{{.*}}extern_decl
+
+; Make sure all local functions have been renamed to <name>.cfi
+; THIN: define hidden i32 @local_func1.cfi()
+; THIN: define hidden i32 @local_func2.cfi()
+; THIN: define hidden i32 @local_func3.cfi(i32 %i){{.*}}
+
+; Indirect references to local_func1 and local_func2 must to through jump table
+; THIN: %fptr1 = select i1 %cmp.i, i32 ()* @local_func1, i32 ()* @local_func2
+
+; Indirect references to extern_weak and extern_decl must go through jump table
+; THIN: %fptr2 = select i1 %cmp.i, i32 ()* select (i1 icmp ne (i32 ()* @extern_weak, i32 ()* null), i32 ()* @extern_weak.cfi_jt, i32 ()* null), i32 ()* @extern_decl.cfi_jt
+
+; Direct calls to extern_weak and extern_decl should go to original names
+; THIN: %call5 = tail call i32 @extern_decl()
+; THIN: %call7 = tail call i32 @extern_weak()
+
+; Direct call to local_func1 should to the renamed version
+; THIN: %call6 = tail call i32 @local_func1.cfi()
+; THIN: %call8 = tail call i32 @local_func3.cfi(i32 4)
+

llvm/test/Transforms/LowerTypeTests/export-icall.ll

@@ -50,7 +50,7 @@ define void @f3(i32 %x) !type !8 {
 
 ; CHECK-DAG: @g                    = alias void (), void ()* [[JT2]]
 
-; CHECK-DAG: define internal void @h.cfi(i8 {{.*}}) !type !{{.*}}
+; CHECK-DAG: define void @h.cfi(i8 {{.*}}) !type !{{.*}}
 ; CHECK-DAG: declare !type !{{.*}} void @external()
 ; CHECK-DAG: declare !type !{{.*}} void @external_weak()
 ; CHECK-DAG: declare !type !{{.*}} void @f.cfi(i32)

llvm/test/Transforms/LowerTypeTests/function-disjoint.ll

@@ -11,13 +11,13 @@ target datalayout = "e-p:64:64"
 ; WASM32: private constant [0 x i8] zeroinitializer
 @0 = private unnamed_addr constant [2 x void ()*] [void ()* @f, void ()* @g], align 16
 
-; X64: define internal void @f.cfi()
+; X64: define void @f.cfi()
 ; WASM32: define void @f() !type !{{[0-9]+}} !wasm.index ![[I0:[0-9]+]]
 define void @f() !type !0 {
   ret void
 }
 
-; X64: define internal void @g.cfi()
+; X64: define void @g.cfi()
 ; WASM32: define void @g() !type !{{[0-9]+}} !wasm.index ![[I1:[0-9]+]]
 define void @g() !type !1 {
   ret void

llvm/test/Transforms/LowerTypeTests/function-weak.ll

@@ -38,7 +38,7 @@ entry:
 ; CHECK: define void @call_f() {
 define void @call_f() {
 entry:
-; CHECK: call void select (i1 icmp ne (void ()* @f, void ()* null), void ()* @[[JT]], void ()* null)()
+; CHECK: call void @f()
   call void @f()
   ret void
 }

llvm/test/Transforms/LowerTypeTests/function.ll

@@ -24,7 +24,7 @@ target datalayout = "e-p:64:64"
 ; ARM: @g = internal alias void (), bitcast ([4 x i8]* getelementptr inbounds ([2 x [4 x i8]], [2 x [4 x i8]]* bitcast (void ()* @[[JT]] to [2 x [4 x i8]]*), i64 0, i64 1) to void ()*)
 ; THUMB: @g = internal alias void (), bitcast ([4 x i8]* getelementptr inbounds ([2 x [4 x i8]], [2 x [4 x i8]]* bitcast (void ()* @[[JT]] to [2 x [4 x i8]]*), i64 0, i64 1) to void ()*)
 
-; NATIVE: define internal void @f.cfi()
+; NATIVE: define void @f.cfi()
 ; WASM32: define void @f() !type !{{[0-9]+}} !wasm.index ![[I0:[0-9]+]]
 define void @f() !type !0 {
   ret void

llvm/test/Transforms/LowerTypeTests/import-icall.ll

@@ -28,14 +28,14 @@ declare void @external()
 declare extern_weak void @external_weak()
 
 ; CHECK:      define hidden i8 @local_a.cfi() {
-; CHECK-NEXT:   call void @external.cfi_jt()
-; CHECK-NEXT:   call void select (i1 icmp ne (void ()* @external_weak, void ()* null), void ()* @external_weak.cfi_jt, void ()* null)()
+; CHECK-NEXT:   call void @external()
+; CHECK-NEXT:   call void @external_weak()
 ; CHECK-NEXT:   ret i8 1
 ; CHECK-NEXT: }
 
 ; internal @local_b is not the same function as "local_b" in the summary.
 ; CHECK:      define internal i8 @local_b() {
-; CHECK-NEXT:   call i8 @local_a()
+; CHECK-NEXT:   call i8 @local_a.cfi()
 
 ; CHECK:      define void @local_decl()
 ; CHECK-NEXT:   call void @local_decl()

llvm/test/Transforms/LowerTypeTests/section.ll

@@ -6,7 +6,7 @@
 target triple = "x86_64-unknown-linux-gnu"
 
 ; CHECK: @f = alias void (), void ()* @[[JT:.*]]
-; CHECK: define internal void @f.cfi() section "xxx"
+; CHECK: define void @f.cfi() section "xxx"
 
 define void @f() section "xxx" !type !0 {
 entry: