[llvm] 0ebbf34 - [ArgPromotion] Don't assume all entry block instrs are executed

[Nikita Popov via llvm-commits]

Author: Nikita Popov
Date: 2022-01-28T16:08:42+01:00
New Revision: 0ebbf3435ff4c3c141549aaf0f791485c28f06f0

URL: https://github.com/llvm/llvm-project/commit/0ebbf3435ff4c3c141549aaf0f791485c28f06f0
DIFF: https://github.com/llvm/llvm-project/commit/0ebbf3435ff4c3c141549aaf0f791485c28f06f0.diff

LOG: [ArgPromotion] Don't assume all entry block instrs are executed

We should abort this walk if we hit any instruction that is not
guaranteed to transfer.

Added: 
    

Modified: 
    llvm/lib/Transforms/IPO/ArgumentPromotion.cpp
    llvm/test/Transforms/ArgumentPromotion/load-after-non-willreturn-call.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/IPO/ArgumentPromotion.cpp b/llvm/lib/Transforms/IPO/ArgumentPromotion.cpp
index ea7baeb3813d..4d5fe9a6ed9a 100644
--- a/llvm/lib/Transforms/IPO/ArgumentPromotion.cpp
+++ b/llvm/lib/Transforms/IPO/ArgumentPromotion.cpp
@@ -46,6 +46,7 @@
 #include "llvm/Analysis/LazyCallGraph.h"
 #include "llvm/Analysis/Loads.h"
 #include "llvm/Analysis/MemoryLocation.h"
+#include "llvm/Analysis/ValueTracking.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/IR/Argument.h"
@@ -610,12 +611,12 @@ static bool isSafeToPromoteArgument(Argument *Arg, Type *ByValTy, AAResults &AAR
     return true;
   };
 
-  // First, iterate the entry block and mark loads of (geps of) arguments as
-  // safe.
+  // First, iterate functions that are guaranteed to execution on function
+  // entry and mark loads of (geps of) arguments as safe.
   BasicBlock &EntryBlock = Arg->getParent()->front();
   // Declare this here so we can reuse it
   IndicesVector Indices;
-  for (Instruction &I : EntryBlock)
+  for (Instruction &I : EntryBlock) {
     if (LoadInst *LI = dyn_cast<LoadInst>(&I)) {
       Value *V = LI->getPointerOperand();
       if (GetElementPtrInst *GEP = dyn_cast<GetElementPtrInst>(V)) {
@@ -649,6 +650,10 @@ static bool isSafeToPromoteArgument(Argument *Arg, Type *ByValTy, AAResults &AAR
       }
     }
 
+    if (!isGuaranteedToTransferExecutionToSuccessor(&I))
+      break;
+  }
+
   // Now, iterate all uses of the argument to see if there are any uses that are
   // not (GEP+)loads, or any (GEP+)loads that are not safe to promote.
   SmallVector<LoadInst *, 16> Loads;

diff  --git a/llvm/test/Transforms/ArgumentPromotion/load-after-non-willreturn-call.ll b/llvm/test/Transforms/ArgumentPromotion/load-after-non-willreturn-call.ll
index 04a0cbcb2774..6af4a358e0d5 100644
--- a/llvm/test/Transforms/ArgumentPromotion/load-after-non-willreturn-call.ll
+++ b/llvm/test/Transforms/ArgumentPromotion/load-after-non-willreturn-call.ll
@@ -3,11 +3,15 @@
 
 declare void @may_not_return()
 
+; The argument cannot be promoted, as we do not know whether the load can be
+; speculatively executed.
+
 define internal i32 @callee(i32* %p) {
 ; CHECK-LABEL: define {{[^@]+}}@callee
-; CHECK-SAME: (i32 [[P_VAL:%.*]]) {
+; CHECK-SAME: (i32* [[P:%.*]]) {
 ; CHECK-NEXT:    call void @may_not_return()
-; CHECK-NEXT:    ret i32 [[P_VAL]]
+; CHECK-NEXT:    [[X:%.*]] = load i32, i32* [[P]], align 4
+; CHECK-NEXT:    ret i32 [[X]]
 ;
   call void @may_not_return() readnone
   %x = load i32, i32* %p
@@ -17,8 +21,7 @@ define internal i32 @callee(i32* %p) {
 define void @caller(i32* %p) {
 ; CHECK-LABEL: define {{[^@]+}}@caller
 ; CHECK-SAME: (i32* [[P:%.*]]) {
-; CHECK-NEXT:    [[P_VAL:%.*]] = load i32, i32* [[P]], align 4
-; CHECK-NEXT:    [[TMP1:%.*]] = call i32 @callee(i32 [[P_VAL]])
+; CHECK-NEXT:    [[TMP1:%.*]] = call i32 @callee(i32* [[P]])
 ; CHECK-NEXT:    ret void
 ;
   call i32 @callee(i32* %p)