[llvm] 4e7d7ca - [Attributor][FIX] Do not try to rewrite functions with casted call sites

[Johannes Doerfert via llvm-commits]

Author: Johannes Doerfert
Date: 2021-08-12T10:39:53-05:00
New Revision: 4e7d7cae674288aa518faadd63ef36608dbdbc13

URL: https://github.com/llvm/llvm-project/commit/4e7d7cae674288aa518faadd63ef36608dbdbc13
DIFF: https://github.com/llvm/llvm-project/commit/4e7d7cae674288aa518faadd63ef36608dbdbc13.diff

LOG: [Attributor][FIX] Do not try to rewrite functions with casted call sites

If we cast a function at the call site it is hard(er) to get the rewrite
correct, let's not attempt it for now.

Fixes PR51448.

Added: 
    

Modified: 
    llvm/lib/Transforms/IPO/Attributor.cpp
    llvm/test/Transforms/Attributor/IPConstantProp/arg-count-mismatch.ll
    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 91b16ec66ee39..8fb82d60033f3 100644
--- a/llvm/lib/Transforms/IPO/Attributor.cpp
+++ b/llvm/lib/Transforms/IPO/Attributor.cpp
@@ -2015,7 +2015,8 @@ bool Attributor::isValidFunctionSignatureRewrite(
   if (!RewriteSignatures)
     return false;
 
-  auto CallSiteCanBeChanged = [](AbstractCallSite ACS) {
+  Function *Fn = Arg.getParent();
+  auto CallSiteCanBeChanged = [Fn](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).
@@ -2023,11 +2024,12 @@ bool Attributor::isValidFunctionSignatureRewrite(
         ACS.getInstruction()->getType() !=
             ACS.getCalledFunction()->getReturnType())
       return false;
+    if (ACS.getCalledOperand()->getType() != Fn->getType())
+      return false;
     // Forbid must-tail calls for now.
     return !ACS.isCallbackCall() && !ACS.getInstruction()->isMustTailCall();
   };
 
-  Function *Fn = Arg.getParent();
   // Avoid var-arg functions for now.
   if (Fn->isVarArg()) {
     LLVM_DEBUG(dbgs() << "[Attributor] Cannot rewrite var-args functions\n");

diff  --git a/llvm/test/Transforms/Attributor/IPConstantProp/arg-count-mismatch.ll b/llvm/test/Transforms/Attributor/IPConstantProp/arg-count-mismatch.ll
index 159e064de07ab..485cf2865fe6c 100644
--- a/llvm/test/Transforms/Attributor/IPConstantProp/arg-count-mismatch.ll
+++ b/llvm/test/Transforms/Attributor/IPConstantProp/arg-count-mismatch.ll
@@ -36,15 +36,10 @@
 ; FIXME we should recognize this as UB and make it an unreachable.
 
 define dso_local i16 @foo(i16 %a) {
-; NOT_CGSCC_NPM-LABEL: define {{[^@]+}}@foo
-; NOT_CGSCC_NPM-SAME: (i16 [[A:%.*]]) {
-; NOT_CGSCC_NPM-NEXT:    [[CALL:%.*]] = call i16 @bar()
-; NOT_CGSCC_NPM-NEXT:    ret i16 [[CALL]]
-;
-; IS__CGSCC_NPM: Function Attrs: nofree norecurse nosync nounwind readnone willreturn
-; IS__CGSCC_NPM-LABEL: define {{[^@]+}}@foo
-; IS__CGSCC_NPM-SAME: (i16 [[A:%.*]]) #[[ATTR0:[0-9]+]] {
-; IS__CGSCC_NPM-NEXT:    ret i16 0
+; CHECK-LABEL: define {{[^@]+}}@foo
+; CHECK-SAME: (i16 [[A:%.*]]) {
+; CHECK-NEXT:    [[CALL:%.*]] = call i16 bitcast (i16 (i16, i16)* @bar to i16 (i16)*)(i16 [[A]])
+; CHECK-NEXT:    ret i16 [[CALL]]
 ;
   %call = call i16 bitcast (i16 (i16, i16) * @bar to i16 (i16) *)(i16 %a)
   ret i16 %call
@@ -53,18 +48,13 @@ define dso_local i16 @foo(i16 %a) {
 define internal i16 @bar(i16 %p1, i16 %p2) {
 ; IS__TUNIT____: Function Attrs: nofree nosync nounwind readnone willreturn
 ; IS__TUNIT____-LABEL: define {{[^@]+}}@bar
-; IS__TUNIT____-SAME: () #[[ATTR0:[0-9]+]] {
+; IS__TUNIT____-SAME: (i16 [[P1:%.*]], i16 [[P2:%.*]]) #[[ATTR0:[0-9]+]] {
 ; IS__TUNIT____-NEXT:    ret i16 0
 ;
-; IS__CGSCC_OPM: Function Attrs: nofree norecurse nosync nounwind readnone willreturn
-; IS__CGSCC_OPM-LABEL: define {{[^@]+}}@bar
-; IS__CGSCC_OPM-SAME: () #[[ATTR0:[0-9]+]] {
-; IS__CGSCC_OPM-NEXT:    ret i16 0
-;
-; IS__CGSCC_NPM: Function Attrs: nofree norecurse nosync nounwind readnone willreturn
-; IS__CGSCC_NPM-LABEL: define {{[^@]+}}@bar
-; IS__CGSCC_NPM-SAME: () #[[ATTR0]] {
-; IS__CGSCC_NPM-NEXT:    ret i16 undef
+; IS__CGSCC____: Function Attrs: nofree norecurse nosync nounwind readnone willreturn
+; IS__CGSCC____-LABEL: define {{[^@]+}}@bar
+; IS__CGSCC____-SAME: (i16 [[P1:%.*]], i16 [[P2:%.*]]) #[[ATTR0:[0-9]+]] {
+; IS__CGSCC____-NEXT:    ret i16 0
 ;
   ret i16 0
 }
@@ -88,7 +78,7 @@ define internal i16 @bar2(i16 %p1, i16 %p2) {
 ;
 ; IS__CGSCC____: Function Attrs: nofree norecurse nosync nounwind readnone willreturn
 ; IS__CGSCC____-LABEL: define {{[^@]+}}@bar2
-; IS__CGSCC____-SAME: (i16 [[P1:%.*]], i16 [[P2:%.*]]) #[[ATTR0:[0-9]+]] {
+; IS__CGSCC____-SAME: (i16 [[P1:%.*]], i16 [[P2:%.*]]) #[[ATTR0]] {
 ; IS__CGSCC____-NEXT:    [[A:%.*]] = add i16 [[P1]], [[P2]]
 ; IS__CGSCC____-NEXT:    ret i16 [[A]]
 ;

diff  --git a/llvm/test/Transforms/Attributor/misc_crash.ll b/llvm/test/Transforms/Attributor/misc_crash.ll
index cd59c92af6cfd..21a3d9eb75fb0 100644
--- a/llvm/test/Transforms/Attributor/misc_crash.ll
+++ b/llvm/test/Transforms/Attributor/misc_crash.ll
@@ -105,6 +105,23 @@ block:
   ret void
 }
 
+define i16 @foo3() {
+; CHECK-LABEL: define {{[^@]+}}@foo3() {
+; CHECK-NEXT:    [[CALL:%.*]] = call i16 bitcast (i16 (i16*, i16)* @bar3 to i16 ()*)()
+; CHECK-NEXT:    ret i16 [[CALL]]
+;
+  %call = call i16 bitcast (i16 (i16*, i16) * @bar3 to i16 () *)()
+  ret i16 %call
+}
+define internal i16 @bar3(i16* %p1, i16 %p2) {
+; CHECK: Function Attrs: nofree nosync nounwind readnone willreturn
+; CHECK-LABEL: define {{[^@]+}}@bar3
+; CHECK-SAME: (i16* nocapture nofree readnone [[P1:%.*]], i16 returned [[P2:%.*]]) #[[ATTR0]] {
+; CHECK-NEXT:    ret i16 [[P2]]
+;
+  ret i16 %p2
+}
+
 ; CHECK-LABEL: declare {{[^@]+}}@func6
 ; CHECK-SAME: (i8*)
 declare void @func6(i8*)