ThinLTO+CFI: short-circuit direct calls to jump table entries

[dmikulin]

Summary

In ThinLTO+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.

[efriedma]

No testcase?

llvm/lib/IR/Value.cpp

477

This looks like it doesn't handle the case of a function with a function pointer argument correctly.

[vlad.tsyrklevich]

Ditto Eli's comments. Also adding a FDecl->replaceDirectCalls(F) in the else block after line 1016 seems like it would also correctly rewrite direct calls for the other two cases (external functions, and local functions declared in the current TU.)

llvm/include/llvm/IR/Value.h

303

s/make/replace/

llvm/lib/Transforms/IPO/LowerTypeTests.cpp

988

LLVM style nit: same as line as the brace above.

[vlad.tsyrklevich]

Also, definitions in the merged module are set to have internal linkage which will not work if we plan to call them directly in other modules. We'll need to change their linkage.

[dmikulin]

Yeah, need a test case.

In D46326#1084792, @vlad.tsyrklevich wrote:

Also, definitions in the merged module are set to have internal linkage which will not work if we plan to call them directly in other modules. We'll need to change their linkage.

Can you give an example? I'm not sure I completely follow...

llvm/lib/IR/Value.cpp

477

I thought getCalledFunction on a CallInst takes care of it...

[dmikulin]

In D46326#1084593, @vlad.tsyrklevich wrote:

Also adding a FDecl->replaceDirectCalls(F) in the else block after line 1016 seems like it would also correctly rewrite direct calls for the other two cases (external functions, and local functions declared in the current TU.)

Aren't the other cases captured by F->replaceUsesExceptBlockAddr(FDecl) at line 1016?

[vlad.tsyrklevich]

In D46326#1085464, @dmikulin wrote:

In D46326#1084593, @vlad.tsyrklevich wrote:

Also adding a FDecl->replaceDirectCalls(F) in the else block after line 1016 seems like it would also correctly rewrite direct calls for the other two cases (external functions, and local functions declared in the current TU.)

Aren't the other cases captured by F->replaceUsesExceptBlockAddr(FDecl) at line 1016?

No, that statement is doing the reverse of what you're doing. It's replacing directs calls to a function to calls to the jumptable entry. (F/FDecl have different meanings in the 2 other cases than how you're using them.)

llvm/lib/IR/Value.cpp

477

It checks that the call is a direct call, but the use might be an argument of the function not the called function, e.g.

void bar();
foo(&bar);

Iterating through the uses of bar you'll hit the call instruction and replace foo(&bar) with foo(&bar.cfi) though it shouldn't be. You need to check that getCalledFunction is the one you actually expect.

[vlad.tsyrklevich]

In D46326#1084858, @dmikulin wrote:

Yeah, need a test case.

In D46326#1084792, @vlad.tsyrklevich wrote:

Also, definitions in the merged module are set to have internal linkage which will not work if we plan to call them directly in other modules. We'll need to change their linkage.

Can you give an example? I'm not sure I completely follow...

Sorry, that was a pretty terse response--I'd tried a build of Chrome with a variant of this change and was reporting the way the build failed without providing much insight as to how or why. buildBitSetsFromFunctionsNative() will set function definitions in the merged module to have internal linkage; function definitions end up in the merged module when compiled with ThinLTO if they're compiled in a translation unit with no externally visible symbols (like in ex2.cpp below.) If that's the case, the definition of a symbol (X<double>.cfi in this example) will be set to have internal linkage so a direct call to it from another translation unit will fail.

ex1.cpp:

#include <stdio.h>

template<class C> void X();

int main() {
  printf("addr %p\n", &X<double>);
  X<double>();
  ((void(*)())40000)();
  return 0;
}

ex2.cpp:

#include <stdio.h>

template<class C> void X() { printf("X\n"); }

namespace { 
  __attribute__((section(".preinit_array"), used)) void (*__preinit)(void) = X<double>;
}

$ clang++ -fuse-ld=lld -flto=thin -fsanitize=cfi-icall -fvisibility=hidden -o ex ex1.cpp ex2.cpp
lld: error: undefined symbol: void X<double>() (.cfi)
>>> referenced by ex1.cpp
>>>               lto.tmp:(main)

[dmikulin]

No, that statement is doing the reverse of what you're doing. It's replacing directs calls to a function to calls to the jumptable entry. (F/FDecl have different meanings in the 2 other cases than how you're using them.)

Sounds like we need to skip direct calls in replaceUsesExceptBlockAddr(). Make it replaceUsesExceptBlockAddrOrDirectCalls()

llvm/lib/IR/Value.cpp

477

I see... I thought there'd be a more straightforward direct call check, but I can't seem to find anything in the code. Changed to check if the getCalledFunction() is the same as what we're trying to replace, i.e. ''this'.

[dmikulin]

No, that statement is doing the reverse of what you're doing. It's replacing directs calls to a function to calls to the jumptable entry. (F/FDecl have different meanings in the 2 other cases than how you're using them.)

Sounds like we need to skip direct calls in replaceUsesExceptBlockAddr(). Make it replaceUsesExceptBlockAddrOrDirectCalls()

Hmm... This broke some tests, test/Transforms/LowerTypeTests/import-icall.ll to be specific. A direct call to @external, which is a CfiFunctionDecl, is expected to go through a jump table entry. With my change it doesn't.
Playing a bit more with thin and full LTO, it looks like a direct external call to a function with the same type as function pointers used in the program, will go through a jump table entry even in full LTO. Is this by design?

[dmikulin]

$ clang++ -fuse-ld=lld -flto=thin -fsanitize=cfi-icall -fvisibility=hidden -o ex ex1.cpp ex2.cpp
lld: error: undefined symbol: void X<double>() (.cfi)

referenced by ex1.cpp

lto.tmp:(main)

I get the same result with full LTO and without my changes, ToT from a few days ago. It may be a separate problem.

[vlad.tsyrklevich]

I get the same result with full LTO and without my changes, ToT from a few days ago. It may be a separate problem.

ToT LTO and ThinLTO builds both work fine for me. Can you post the error you get? I'm not sure how ToT would produce a direct call to a function with a jumptable entry.

Hmm... This broke some tests, test/Transforms/LowerTypeTests/import-icall.ll to be specific. A direct call to @external, which is a CfiFunctionDecl, is expected to go through a jump table entry. With my change it doesn't.

Presumably it's only expected to because the test just checks the old behavior.

Playing a bit more with thin and full LTO, it looks like a direct external call to a function with the same type as function pointers used in the program, will go through a jump table entry even in full LTO. Is this by design?

It's the same behavior as done for ThinLTO--it can be replaced with direct calls without impacting the correctness of the CFI checking.

[dmikulin]

ToT LTO and ThinLTO builds both work fine for me. Can you post the error you get? I'm not sure how ToT would produce a direct call to a function with a jumptable entry.

Never mind: too many builds, something got mixed up. I need to take another look.

[dmikulin]

Added a test case.
Skip direct calls in addition to blockaddr uses when dealing with CFI.
Don't set linkage to internal for .cfi symbols as they may be needed by direct calls.

[vlad.tsyrklevich]

It would be nice to also add a direct call monolithic LTO test as monolithic LTO has a slightly different code path for renaming functions and based on a quick look above it looks like we might only correctly rename direct calls to local, not external, functions from the code above.

llvm/include/llvm/IR/Value.h

303

Given that this function usefulness is now likely specific to LowerTypeTests it's probably worth just including it there directly instead of polluting all of Value's children.

llvm/lib/IR/Value.cpp

477

I haven't tested it but I think this will not work for foo(&foo), e.g. it will change it to foo.cfi(&foo.cfi) instead of foo.cfi(&foo).

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

2

I'd like to see this test test the correct direct call behavior for 3 types of functions (matching the 3 cases in ::importFunction): local function declared in the current translation unit, local function declared in another TU, and an external function.

7

This variable is named @external but is included in CfiFunctionDefs instead of CfiFunctionsDecls in the summary so it's treated like a local function. e.g. for direct calls it's re-written to @external.cfi. An direct call to an external function would be left as @external (instead of being re-written to @external.cfi_jt to point to the jump table entry.)

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

30–31

For @external and @local_a below we should add an instruction where both address-taken so that we can still test the original behavior (e.g. that non-direct calls point to the jump table aliases as expected.)

[dmikulin]

llvm/lib/IR/Value.cpp

477

I can't seem to reproduce it without at least having the fptr argument wrapped in a bitcast, which makes this check work. I'm thinking I'd need a recursively defined fptr type to avoid the cast.
I don't necessarily like the roundabout way of checking if it's a direct call. Is there a better way?

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

2

I'll work on the test cases.

7

Direct calls to external functions may still go through a jump table entry, if those functions are of the same type as functions called indirectly. I'll change the test case and add more test points.

[efriedma]

llvm/lib/IR/Value.cpp

477

Yes, it's impossible to have a function with itself as an argument... at least for now (it might become possible with the typeless pointer work, but no idea if/when that will happen).

Looking around a bit, it looks like there's a helper CallSite::isCallee which is meant for this sort of check; it takes a Use*.

[vlad.tsyrklevich]

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

7

Shouldn't the logic in this change prevent that from happening? (modulo the thin versus monolithic LTO logic I have already mentioned.)

[dmikulin]

Fixed a few missed cases where direct calls were routed through jump tables.
Added more test points for thin and full LTO.

[vlad.tsyrklevich]

Otherwise looks good to me.

llvm/lib/Transforms/IPO/LowerTypeTests.cpp

1221

This logic is slightly different because it skips block addresses. Is it possible that a local weak function would exist that would cause an exception by not being replaced here?

1451

Ditto the above. Because of !IsDefinition I can't think of a case where a blockaddress would be valid though, so this might be fine.

[dmikulin]

llvm/lib/Transforms/IPO/LowerTypeTests.cpp

1221

We had to special case blockaddr before and I felt that since they don't refer to function entry points (correct?), they should not be routed through jump tables and should be skipped here as well. Or, if this case never occurs, at least it should be safe to skip.

1451

Same here.

[vlad.tsyrklevich]

LGTM but I'd like to see what @pcc has to say.

llvm/lib/Transforms/IPO/LowerTypeTests.cpp

1221

I think you're right, the RAUW would replace the expression with a conditional so blockaddr would not work as is.

[pcc]

LGTM

llvm/lib/Transforms/IPO/LowerTypeTests.cpp

982

It's a little confusing that FDecl is used for a different purpose here than in the rest of the function (where it is used to refer to a jump table entry). Maybe make this declare a separate RealF variable, and then once you take care of the else after return below you can move the declaration of FDecl after the if statement.

987

No else after return.

[thakis]

Various chromium test cases started failing on the first build with this change, and I don't see anything suspicious chromium-side: https://ci.chromium.org/buildbot/chromium.clang/CFI%20Linux%20%28icall%29/9886 Any chance that could be due to this?

[thakis]

On this bot too: https://ci.chromium.org/buildbot/chromium.clang/CFI%20Linux%20ToT/2783

[vlad.tsyrklevich]

Based on a quick look at a couple of stack traces it looks like some direct
calls now call symbols that are no longer overriden by their replacements
in chromium, e.g. calls to realloc() from a DSO previously resolved to the
tcmalloc implementation in chromium instead of realloc.cfi(), the locally
linked tcmalloc implementation in that library.

[pcc]

In D46326#1105694, @vlad.tsyrklevich wrote:

Based on a quick look at a couple of stack traces it looks like some direct
calls now call symbols that are no longer overriden by their replacements
in chromium, e.g. calls to realloc() from a DSO previously resolved to the
tcmalloc implementation in chromium instead of realloc.cfi(), the locally
linked tcmalloc implementation in that library.

I imagine that we will need to limit this optimization to dso_local functions.

[thakis]

Reverted in r332838. Please don't keep trunk broken for days.

[dmikulin]

Sorry, I'm traveling this week, didn't see this until now. I didn't get any bot failure notifications after the commit, so I assumed it was good. Do you have a test case I can use to reproduce it?

[dmikulin]

In D46326#1105694, @vlad.tsyrklevich wrote:

Based on a quick look at a couple of stack traces it looks like some direct
calls now call symbols that are no longer overriden by their replacements
in chromium, e.g. calls to realloc() from a DSO previously resolved to the
tcmalloc implementation in chromium instead of realloc.cfi(), the locally
linked tcmalloc implementation in that library.

If I understand correctly, you have a DSO linked to the main chromium binary. The DSO calls realloc(), which normally resolves at runtime to the realloc() in the main binary. The DSO also defines realloc(). After my changes instead of dynamically resolving realloc(), it's statically resolved to realloc.cfi(). How do you build your DSO? I can't get defined symbols to resolve externally with CFI enabled...

[dmikulin]

Oh, this must be with cross-DSO CFI. I see what the problem is. We need to skip symbols not resolved internally from replacing direct calls to calls to <name>.cfi. An extra test case wound't hurt, so if you have it please send it to me.

Sorry again for breaking top-of-tree, didn't realize it happened. Usually the bots send auto notification to authors of suspected commits. I was keeping an eye on the bots, checked reported breakages, but none of them were mine. Are chromium bots set up the same way or do I need to manually monitor their status?

[pcc]

You don't need cross-DSO CFI to reproduce this, the issue has to do with regular ELF interposition.

Here is a standalone reproducer:

$ cat main.c
void baz();
void foo() {
  puts("main");
}

int main() {
  direct();
  indirect();
}
$ cat dso.c
void foo() {
  puts("dso");
}
typedef void (*fp)();

fp bar() { return foo; }
void indirect() {
  bar()();
}
void direct() {
  foo();
}
$ ~/src2/llvm-project4/ra/bin/clang -fsanitize=cfi-icall -shared -o dso.so dso.c -flto -fPIC
$ ~/src2/llvm-project4/ra/bin/clang -fsanitize=cfi-icall  -o main main.c -flto dso.so
$ LD_LIBRARY_PATH=. ./main
dso
main

If I rebuild dso.so without cfi-icall I see the correct result:

$ ~/src2/llvm-project4/ra/bin/clang  -shared -o dso.so dso.c -flto -fPIC
$ LD_LIBRARY_PATH=. ./main
main
main

[dmikulin]

Changed the code not to short circuit functions that can be overridden at run time and not to expose <func>.cfi names to the runtime.
Added a test point to check for that.

Would it be possible to run this patch through Chromium build/test before I commit?

[vlad.tsyrklevich]

Dmitry, I can run the tests for you locally; however, it's probably worth having @pcc take a look first.

llvm/lib/Transforms/IPO/LowerTypeTests.cpp

421

clang-format

432

Update

1462

Why change the visibility to hidden only for !hasLocalLinkage ?

1646

Did new changes cause this to be required?

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

9

remove?

[dmikulin]

llvm/lib/Transforms/IPO/LowerTypeTests.cpp

1462

setVisibility() asserts if hasLocalLinkage()

1646

Handling of Constants? No, it's been there before my changes.

[pcc]

llvm/lib/Transforms/IPO/LowerTypeTests.cpp

982

Can you check for dso_local instead of default visibility here? I think that should be exactly the property that you want to check for.

999

Is there a reason to set the visibility indirectly via this variable here?

1642

Similarly I think you can check for dso_local here.

[dmikulin]

Can you guys please give it a try with Chromium? Do you get any measurable performance improvements with this change?

llvm/lib/Transforms/IPO/LowerTypeTests.cpp

982

Thanks, that's better.

999

I need to set visibility late, after replaceCfiUses(). If I do it sooner, the symbol becomes dso_local, and it breaks the logic in replaceCfiUses().

[vlad.tsyrklevich]

Did you forget to update the change? I'll run it against the failing chrome tests after. CFI-icall presents a <=1% performance regression for chrome so it would be hard to pick up a difference. From my past look at it, most of the regression appeared to because of cfi-icall checks at every malloc/free (because we override their implementations and require a function pointer call.)

[dmikulin]

Changed the checks to look for dso_local instread of visibility.

[vlad.tsyrklevich]

The browser/media/cronet tests that failed previously now pass.

[pcc]

LGTM

llvm/lib/Transforms/IPO/LowerTypeTests.cpp

999

Okay, please mention that in a comment.

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

16

You might also want to test the case where a default visibility function is dso_local.

llvm/include/llvm/IR/Value.h

public:
 
   /// replaceUsesOutsideBlock - 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 does not support basic block
   /// 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.
   //
   // Materializing a function can introduce new uses, so these methods come in
   // two variants:
   // The methods that start with materialized_ check the uses that are
   // currently known given which functions are materialized. Be very careful
   // when using them since you might not get all uses.
   // The methods that don't start with materialized_ assert that modules is
   // fully materialized.

llvm/lib/IR/Value.cpp

for (; UI != E;) {
     ++UI;
     auto *Usr = dyn_cast<Instruction>(U.getUser());
     if (Usr && Usr->getParent() == BB)
       continue;
     U.set(New);
   }
 }
 
-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 {
   PSK_ZeroIndices,
   PSK_ZeroIndicesAndAliases,
   PSK_ZeroIndicesAndAliasesAndBarriers,
   PSK_InBoundsConstantIndices,
   PSK_InBounds
 };
 
 template <PointerStripKind StripKind>
 static const Value *stripPointerCastsAndOffsets(const Value *V) {
   if (!V->getType()->isPointerTy())
     return V;
 
   // Even though we don't look through PHI nodes, we could be called on an
   // instruction in an unreachable block, which may be on a cycle.
   SmallPtrSet<const Value *, 4> Visited;

llvm/lib/Transforms/IPO/LowerTypeTests.cpp

void buildBitSetsFromFunctionsNative(ArrayRef<Metadata *> TypeIds,
                                        ArrayRef<GlobalTypeMember *> Functions);
   void buildBitSetsFromFunctionsWASM(ArrayRef<Metadata *> TypeIds,
                                      ArrayRef<GlobalTypeMember *> Functions);
   void
   buildBitSetsFromDisjointSet(ArrayRef<Metadata *> TypeIds,
                               ArrayRef<GlobalTypeMember *> Globals,
                               ArrayRef<ICallBranchFunnel *> ICallBranchFunnels);
 
   void replaceWeakDeclarationWithJumpTablePtr(Function *F, Constant *JT);
   void moveInitializerToModuleConstructor(GlobalVariable *GV);
   void findGlobalVariableUsersOf(Constant *C,
                                  SmallSetVector<GlobalVariable *, 8> &Out);
 
   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);
 
   bool lower();
 
   // Lower the module using the action and summary passed as command line
   // arguments. For testing purposes only.
void LowerTypeTestsModule::importTypeTest(CallInst *CI) {
   CI->eraseFromParent();
 }
 
 // ThinLTO backend: the function F has a jump table entry; update this module
 // accordingly. isDefinition describes the type of the jump table entry.
 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) {
+  if (F->isDeclarationForLinker() && isDefinition) {
+    FDecl = Function::Create(F->getFunctionType(), GlobalValue::ExternalLinkage,
+                             Name + ".cfi", &M);
+    FDecl->setVisibility(Visibility);
+    replaceDirectCalls(F, FDecl);
+    return;
+  } else if (F->isDeclarationForLinker() && !isDefinition) {
     // Declaration of an external function.
     FDecl = Function::Create(F->getFunctionType(), GlobalValue::ExternalLinkage,
                              Name + ".cfi_jt", &M);
     FDecl->setVisibility(GlobalValue::HiddenVisibility);
   } else if (isDefinition) {
     F->setName(Name + ".cfi");
     F->setLinkage(GlobalValue::ExternalLinkage);
     F->setVisibility(GlobalValue::HiddenVisibility);
     FDecl = Function::Create(F->getFunctionType(), GlobalValue::ExternalLinkage,
                              Name, &M);
     FDecl->setVisibility(Visibility);
 
     // Delete aliases pointing to this function, they'll be re-created in the
     // merged output
     SmallVector<GlobalAlias*, 4> ToErase;
     for (auto &U : F->uses()) {
       if (auto *A = dyn_cast<GlobalAlias>(U.getUser())) {
         Function *AliasDecl = Function::Create(
             F->getFunctionType(), GlobalValue::ExternalLinkage, "", &M);
         AliasDecl->takeName(A);
if (F->isDeclarationForLinker() && isDefinition) {
     // so that it is treated the same way as a function defined outside of the
     // LTO unit.
     return;
   }
 
   if (F->isWeakForLinker())
     replaceWeakDeclarationWithJumpTablePtr(F, FDecl);
   else
-    F->replaceUsesExceptBlockAddr(FDecl);
+    replaceCfiUses(F, FDecl);
 }
 
 void LowerTypeTestsModule::lowerTypeTestCalls(
     ArrayRef<Metadata *> TypeIds, Constant *CombinedGlobalAddr,
     const DenseMap<GlobalTypeMember *, uint64_t> &GlobalLayout) {
   CombinedGlobalAddr = ConstantExpr::getBitCast(CombinedGlobalAddr, Int8PtrTy);
 
   // For each type identifier in this disjoint set...
void LowerTypeTestsModule::replaceWeakDeclarationWithJumpTablePtr(
   for (auto GV : GlobalVarUsers)
     moveInitializerToModuleConstructor(GV);
 
   // Can not RAUW F with an expression that uses F. Replace with a temporary
   // placeholder first.
   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,
                             Constant::getNullValue(F->getType())),
       JT, Constant::getNullValue(F->getType()));
   PlaceholderFn->replaceAllUsesWith(Target);
   PlaceholderFn->eraseFromParent();
 }
if (Functions[I]->isExported()) {
         JtAlias->setVisibility(GlobalValue::HiddenVisibility);
         ExportSummary->cfiFunctionDecls().insert(F->getName());
       }
     }
     if (!IsDefinition) {
       if (F->isWeakForLinker())
         replaceWeakDeclarationWithJumpTablePtr(F, CombinedGlobalElemPtr);
       else
-        F->replaceAllUsesWith(CombinedGlobalElemPtr);
+        replaceCfiUses(F, CombinedGlobalElemPtr);
     } else {
       assert(F->getType()->getAddressSpace() == 0);
 
       GlobalAlias *FAlias = GlobalAlias::create(
           F->getValueType(), 0, F->getLinkage(), "", CombinedGlobalElemPtr, &M);
       FAlias->setVisibility(F->getVisibility());
       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);
 }
 
 /// Assign a dummy layout using an incrementing counter, tag each function
 /// with its index represented as metadata, and lower each type test to an
 /// integer range comparison. During generation of the indirect function call
if (!ClWriteSummary.empty()) {
 
     yaml::Output Out(OS);
     Out << Summary;
   }
 
   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));
   Function *ICallBranchFunnelFunc =
       M.getFunction(Intrinsic::getName(Intrinsic::icall_branch_funnel));
   if ((!TypeTestFunc || TypeTestFunc->use_empty()) &&
       (!ICallBranchFunnelFunc || ICallBranchFunnelFunc->use_empty()) &&
       !ExportSummary && !ImportSummary)

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

+---
+TypeIdMap:
+  typeid1:
+    TTRes:
+      Kind:            AllOnes
+CfiFunctionDefs:
+  - internal_def
+CfiFunctionDecls:
+  - external_decl
+...

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

+---
+TypeIdMap:
+  typeid1:
+    TTRes:
+      Kind:            AllOnes
+CfiFunctionDefs:
+  - local_func1
+  - local_func2
+  - local_func3
+CfiFunctionDecls:
+  - extern_decl
+  - extern_weak
+...

llvm/test/Transforms/LowerTypeTests/blockaddress.ll

 ; 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)
 ; CHECK-NEXT: }
 
 target triple = "x86_64-unknown-linux"
 
 define void @f1() {

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

+; 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

+; 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

 ; CHECK-DAG: @f2                   = alias void (i32), {{.*}}getelementptr {{.*}}void ()* [[JT1]]
 ; CHECK-DAG: @external.cfi_jt      = hidden alias void (), {{.*}}getelementptr {{.*}}void ()* [[JT1]]
 ; CHECK-DAG: @external_weak.cfi_jt = hidden alias void (), {{.*}}getelementptr {{.*}}void ()* [[JT1]]
 
 ; CHECK-DAG: @__typeid_typeid2_global_addr = hidden alias i8, bitcast (void ()* [[JT2:.*]] to i8*)
 
 ; 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)
 ; CHECK-DAG: declare !type !{{.*}} void @f2.cfi(i32)
 ; CHECK-DAG: define void @f3(i32 {{.*}}) !type !3
 ; CHECK-DAG: !3 = !{i64 0, !"typeid3"}
 ; CHECK-DAG: declare !type !{{.*}} void @g.cfi()
 

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

 ; RUN: opt -S -lowertypetests -mtriple=x86_64-unknown-linux-gnu < %s | FileCheck --check-prefix=X64 %s
 ; RUN: opt -S -lowertypetests -mtriple=wasm32-unknown-unknown < %s | FileCheck --check-prefix=WASM32 %s
 
 ; Tests that we correctly handle bitsets with disjoint call target sets.
 
 target datalayout = "e-p:64:64"
 
 ; X64: @f = alias void (), void ()* @[[JT0:.*]]
 ; X64: @g = alias void (), void ()* @[[JT1:.*]]
 
 ; 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
 }
 
 !0 = !{i32 0, !"typeid1"}
 !1 = !{i32 0, !"typeid2"}
 

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

 entry:
 ; CHECK: ret i1 icmp ne (void ()* select (i1 icmp ne (void ()* @f, void ()* null), void ()* @[[JT:.*]], void ()* null), void ()* null)
   ret i1 icmp ne (void ()* @f, void ()* null)
 }
 
 ; 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
 }
 
 declare i1 @llvm.type.test(i8* %ptr, metadata %bitset) nounwind readnone
 
 define i1 @foo(i8* %p) {
   %x = call i1 @llvm.type.test(i8* %p, metadata !"typeid1")

llvm/test/Transforms/LowerTypeTests/function.ll

 ; WASM32: private constant [0 x i8] zeroinitializer
 
 ; NATIVE: @f = alias void (), void ()* @[[JT:.*]]
 
 ; X86: @g = internal alias void (), bitcast ([8 x i8]* getelementptr inbounds ([2 x [8 x i8]], [2 x [8 x i8]]* bitcast (void ()* @[[JT]] to [2 x [8 x i8]]*), i64 0, i64 1) to void ()*)
 ; 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
 }
 
 ; NATIVE: define internal void @g.cfi()
 ; WASM32: define internal void @g() !type !{{[0-9]+}} !wasm.index ![[I1:[0-9]+]]
 define internal void @g() !type !0 {

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

 define void @local_decl() {
   call void @local_decl()
   ret void
 }
 
 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()
 
 ; CHECK: declare void @external()
 ; CHECK: declare extern_weak void @external_weak()
 ; CHECK: declare i8 @local_a()
 ; CHECK: declare hidden void @external.cfi_jt()
 ; CHECK: declare hidden void @external_weak.cfi_jt()

llvm/test/Transforms/LowerTypeTests/section.ll

 ; Test that functions with "section" attribute are accepted, and jumptables are
 ; emitted in ".text".
 
 ; RUN: opt -S -lowertypetests < %s | FileCheck %s
 
 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:
   ret void
 }
 
 define i1 @g() {
 entry: