[llvm] a5a272a - [SafeStack] Don't create SCEV min between pointer and integer (PR54784)

[Nikita Popov via llvm-commits]

Author: Nikita Popov
Date: 2022-04-08T09:44:00+02:00
New Revision: a5a272a491406874e5147ba474182d30098ddfd4

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

LOG: [SafeStack] Don't create SCEV min between pointer and integer (PR54784)

Rather than rewriting the alloca pointer to zero, use
removePointerBase() to drop the base pointer. This will simply bail
if the base pointer is not the alloca. We could try doing something
more fancy here (like dropping the sources not based on the alloca
on the premise that they aren't SafeStack-relevant), but I don't
think that's worthwhile.

Fixes https://github.com/llvm/llvm-project/issues/54784.

Differential Revision: https://reviews.llvm.org/D123309

Added: 
    llvm/test/Transforms/SafeStack/pr54784.ll

Modified: 
    llvm/lib/CodeGen/SafeStack.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/CodeGen/SafeStack.cpp b/llvm/lib/CodeGen/SafeStack.cpp
index 868fba3ae662c..1f7eac4a9b479 100644
--- a/llvm/lib/CodeGen/SafeStack.cpp
+++ b/llvm/lib/CodeGen/SafeStack.cpp
@@ -101,24 +101,6 @@ static cl::opt<bool> ClColoring("safe-stack-coloring",
 
 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 @@ uint64_t SafeStack::getStaticAllocaAllocationSize(const AllocaInst* AI) {
 
 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 =

diff  --git a/llvm/test/Transforms/SafeStack/pr54784.ll b/llvm/test/Transforms/SafeStack/pr54784.ll
new file mode 100644
index 0000000000000..d8433f941ec25
--- /dev/null
+++ b/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
+}