Index: llvm/lib/CodeGen/SafeStack.cpp
===================================================================
--- llvm/lib/CodeGen/SafeStack.cpp
+++ llvm/lib/CodeGen/SafeStack.cpp
@@ -101,24 +101,6 @@
 
 namespace {
 
-/// Rewrite an SCEV expression for a memory access address to an expression that
-/// represents offset from the given alloca.
-///
-/// The implementation simply replaces all mentions of the alloca with zero.
-class AllocaOffsetRewriter : public SCEVRewriteVisitor<AllocaOffsetRewriter> {
-  const Value *AllocaPtr;
-
-public:
-  AllocaOffsetRewriter(ScalarEvolution &SE, const Value *AllocaPtr)
-      : SCEVRewriteVisitor(SE), AllocaPtr(AllocaPtr) {}
-
-  const SCEV *visitUnknown(const SCEVUnknown *Expr) {
-    if (Expr->getValue() == AllocaPtr)
-      return SE.getZero(Expr->getType());
-    return Expr;
-  }
-};
-
 /// The SafeStack pass splits the stack of each function into the safe
 /// stack, which is only accessed through memory safe dereferences (as
 /// determined statically), and the unsafe stack, which contains all
@@ -233,9 +215,18 @@
 
 bool SafeStack::IsAccessSafe(Value *Addr, uint64_t AccessSize,
                              const Value *AllocaPtr, uint64_t AllocaSize) {
-  AllocaOffsetRewriter Rewriter(SE, AllocaPtr);
-  const SCEV *Expr = Rewriter.visit(SE.getSCEV(Addr));
+  const SCEV *AddrExpr = SE.getSCEV(Addr);
+  const auto *Base = dyn_cast<SCEVUnknown>(SE.getPointerBase(AddrExpr));
+  if (!Base || Base->getValue() != AllocaPtr) {
+    LLVM_DEBUG(
+        dbgs() << "[SafeStack] "
+               << (isa<AllocaInst>(AllocaPtr) ? "Alloca " : "ByValArgument ")
+               << *AllocaPtr << "\n"
+               << "SCEV " << *AddrExpr << " not directly based on alloca\n");
+    return false;
+  }
 
+  const SCEV *Expr = SE.removePointerBase(AddrExpr);
   uint64_t BitWidth = SE.getTypeSizeInBits(Expr->getType());
   ConstantRange AccessStartRange = SE.getUnsignedRange(Expr);
   ConstantRange SizeRange =
Index: llvm/test/Transforms/SafeStack/pr54784.ll
===================================================================
--- /dev/null
+++ llvm/test/Transforms/SafeStack/pr54784.ll
@@ -0,0 +1,26 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -safe-stack < %s | FileCheck %s
+
+target triple = "x86_64-unknown-unknown"
+
+define void @min_of_pointers(i32* %p) safestack {
+; CHECK-LABEL: @min_of_pointers(
+; CHECK-NEXT:    [[UNSAFE_STACK_PTR:%.*]] = load i8*, i8** @__safestack_unsafe_stack_ptr, align 8
+; CHECK-NEXT:    [[UNSAFE_STACK_STATIC_TOP:%.*]] = getelementptr i8, i8* [[UNSAFE_STACK_PTR]], i32 -16
+; CHECK-NEXT:    store i8* [[UNSAFE_STACK_STATIC_TOP]], i8** @__safestack_unsafe_stack_ptr, align 8
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr i8, i8* [[UNSAFE_STACK_PTR]], i32 -4
+; CHECK-NEXT:    [[A_UNSAFE1:%.*]] = bitcast i8* [[TMP1]] to i32*
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i32* [[A_UNSAFE1]], [[P:%.*]]
+; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr i8, i8* [[UNSAFE_STACK_PTR]], i32 -4
+; CHECK-NEXT:    [[A_UNSAFE:%.*]] = bitcast i8* [[TMP2]] to i32*
+; CHECK-NEXT:    [[S:%.*]] = select i1 [[CMP]], i32* [[A_UNSAFE]], i32* [[P]]
+; CHECK-NEXT:    store i32 0, i32* [[S]], align 4
+; CHECK-NEXT:    store i8* [[UNSAFE_STACK_PTR]], i8** @__safestack_unsafe_stack_ptr, align 8
+; CHECK-NEXT:    ret void
+;
+  %a = alloca i32
+  %cmp = icmp ult i32* %a, %p
+  %s = select i1 %cmp, i32* %a, i32* %p
+  store i32 0, i32* %s
+  ret void
+}