From 0ebbf3435ff4c3c141549aaf0f791485c28f06f0 Mon Sep 17 00:00:00 2001
From: Nikita Popov <npopov@redhat.com>
Date: Fri, 28 Jan 2022 16:07:40 +0100
Subject: [PATCH] [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.
---
 llvm/lib/Transforms/IPO/ArgumentPromotion.cpp                 | 11 ++++++++---
 .../ArgumentPromotion/load-after-non-willreturn-call.ll       | 11 +++++++----
 2 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/llvm/lib/Transforms/IPO/ArgumentPromotion.cpp b/llvm/lib/Transforms/IPO/ArgumentPromotion.cpp
index ea7baeb..4d5fe9a 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 04a0cbc..6af4a35 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)
-- 
2.7.4