[llvm] 8d94d3c - [Attributor][FIX] Disallow function signature rewrite for casted calls

[Johannes Doerfert via llvm-commits]

Author: Johannes Doerfert
Date: 2020-05-11T15:32:47-05:00
New Revision: 8d94d3c3b44c3a27a69b153cef9be4b8e481150e

URL: https://github.com/llvm/llvm-project/commit/8d94d3c3b44c3a27a69b153cef9be4b8e481150e
DIFF: https://github.com/llvm/llvm-project/commit/8d94d3c3b44c3a27a69b153cef9be4b8e481150e.diff

LOG: [Attributor][FIX] Disallow function signature rewrite for casted calls

We will now ensure ensure the return type of called function is the type
of all call sites we are going to rewrite. This avoids a problem
partially fixed by D79680. The part that was not covered is a use of
this "weird" casted call site (see `@func3` in `misc_crash.ll`).

misc_crash.ll checks are auto-generated now.

Added: 
    

Modified: 
    llvm/lib/Transforms/IPO/Attributor.cpp
    llvm/test/Transforms/Attributor/misc_crash.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/IPO/Attributor.cpp b/llvm/lib/Transforms/IPO/Attributor.cpp
index e1189239d3a0..c1a089b73488 100644
--- a/llvm/lib/Transforms/IPO/Attributor.cpp
+++ b/llvm/lib/Transforms/IPO/Attributor.cpp
@@ -1341,6 +1341,13 @@ bool Attributor::isValidFunctionSignatureRewrite(
     Argument &Arg, ArrayRef<Type *> ReplacementTypes) {
 
   auto CallSiteCanBeChanged = [](AbstractCallSite ACS) {
+    // Forbid the call site to cast the function return type. If we need to
+    // rewrite these functions we need to re-create a cast for the new call site
+    // (if the old had uses).
+    if (!ACS.getCalledFunction() ||
+        ACS.getInstruction()->getType() !=
+            ACS.getCalledFunction()->getReturnType())
+      return false;
     // Forbid must-tail calls for now.
     return !ACS.isCallbackCall() && !ACS.getInstruction()->isMustTailCall();
   };
@@ -1594,10 +1601,11 @@ ChangeStatus Attributor::rewriteFunctionSignatures(
     for (auto &CallSitePair : CallSitePairs) {
       CallBase &OldCB = *CallSitePair.first;
       CallBase &NewCB = *CallSitePair.second;
+      assert(OldCB.getType() == NewCB.getType() &&
+             "Cannot handle call sites with 
diff erent types!");
       ModifiedFns.insert(OldCB.getFunction());
       CGUpdater.replaceCallSite(OldCB, NewCB);
-      if (!OldCB.use_empty())
-        OldCB.replaceAllUsesWith(&NewCB);
+      OldCB.replaceAllUsesWith(&NewCB);
       OldCB.eraseFromParent();
     }
 

diff  --git a/llvm/test/Transforms/Attributor/misc_crash.ll b/llvm/test/Transforms/Attributor/misc_crash.ll
index 476ebbee2048..989b2bd8e204 100644
--- a/llvm/test/Transforms/Attributor/misc_crash.ll
+++ b/llvm/test/Transforms/Attributor/misc_crash.ll
@@ -1,39 +1,77 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --function-signature
 ; RUN: opt -attributor -S %s | FileCheck %s
 ; RUN: opt -passes=attributor -S %s | FileCheck %s
 
 @var1 = internal global [1 x i32] undef
 @var2 = internal global i32 0
 
-; CHECK-LABEL: define i32 addrspace(1)* @foo(i32 addrspace(4)* nofree readnone %arg)
 define i32 addrspace(1)* @foo(i32 addrspace(4)* %arg) {
+; CHECK-LABEL: define {{[^@]+}}@foo
+; CHECK-SAME: (i32 addrspace(4)* nofree readnone [[ARG:%.*]]) #0
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = addrspacecast i32 addrspace(4)* [[ARG]] to i32 addrspace(1)*
+; CHECK-NEXT:    ret i32 addrspace(1)* [[TMP0]]
+;
 entry:
   %0 = addrspacecast i32 addrspace(4)* %arg to i32 addrspace(1)*
   ret i32 addrspace(1)* %0
 }
 
 define i32* @func1() {
+; CHECK-LABEL: define {{[^@]+}}@func1() #0
+; CHECK-NEXT:    [[PTR:%.*]] = call i32* @func1a() #0
+; CHECK-NEXT:    ret i32* [[PTR]]
+;
   %ptr = call i32* @func1a([1 x i32]* @var1)
   ret i32* %ptr
 }
 
+; UTC_ARGS: --disable
 ; CHECK-LABEL: define internal nonnull align 4 dereferenceable(4) i32* @func1a()
 ; CHECK-NEXT: ret i32* getelementptr inbounds ([1 x i32], [1 x i32]* @var1, i32 0, i32 0)
 define internal i32* @func1a([1 x i32]* %arg) {
   %ptr = getelementptr inbounds [1 x i32], [1 x i32]* %arg, i64 0, i64 0
   ret i32* %ptr
 }
+; UTC_ARGS: --enable
 
 define internal void @func2a(i32* %0) {
+; CHECK-LABEL: define {{[^@]+}}@func2a
+; CHECK-SAME: (i32* nocapture nofree nonnull writeonly align 4 dereferenceable(4) [[TMP0:%.*]]) #1
+; CHECK-NEXT:    store i32 0, i32* @var2, align 4
+; CHECK-NEXT:    ret void
+;
   store i32 0, i32* %0
   ret void
 }
 
-; CHECK-LABEL: define i32 @func2()
-; CHECK-NEXT:   tail call void @func2a()
-; CHECK-NEXT:   %1 = load i32, i32* @var2, align 4
-; CHECK-NEXT:   ret i32 %1
 define i32 @func2() {
+; CHECK-LABEL: define {{[^@]+}}@func2()
+; CHECK-NEXT:    [[TMP1:%.*]] = tail call i32 (i32*, ...) bitcast (void (i32*)* @func2a to i32 (i32*, ...)*)(i32* nonnull align 4 dereferenceable(4) @var2)
+; CHECK-NEXT:    [[TMP2:%.*]] = load i32, i32* @var2, align 4
+; CHECK-NEXT:    ret i32 [[TMP2]]
+;
   %1 = tail call i32 (i32*, ...) bitcast (void (i32*)* @func2a to i32 (i32*, ...)*)(i32* @var2)
   %2 = load i32, i32* @var2
   ret i32 %2
 }
+
+define i32 @func3(i1 %false) {
+; CHECK-LABEL: define {{[^@]+}}@func3
+; CHECK-SAME: (i1 [[FALSE:%.*]])
+; CHECK-NEXT:    [[TMP1:%.*]] = tail call i32 (i32*, ...) bitcast (void (i32*)* @func2a to i32 (i32*, ...)*)(i32* nonnull align 4 dereferenceable(4) @var2)
+; CHECK-NEXT:    br i1 [[FALSE]], label [[USE_BB:%.*]], label [[RET_BB:%.*]]
+; CHECK:       use_bb:
+; CHECK-NEXT:    ret i32 [[TMP1]]
+; CHECK:       ret_bb:
+; CHECK-NEXT:    [[TMP2:%.*]] = load i32, i32* @var2, align 4
+; CHECK-NEXT:    ret i32 [[TMP2]]
+;
+  %1 = tail call i32 (i32*, ...) bitcast (void (i32*)* @func2a to i32 (i32*, ...)*)(i32* @var2)
+  br i1 %false, label %use_bb, label %ret_bb
+use_bb:
+  ret i32 %1
+ret_bb:
+  %2 = load i32, i32* @var2
+  ret i32 %2
+}