[llvm] [StackSafetyAnalysis] Don't call getTruncateOrZeroExtend for pointers of different sizes (PR #79804)

[Pierre van Houtryve via llvm-commits]

https://github.com/Pierre-vh updated https://github.com/llvm/llvm-project/pull/79804

>From f3d5e4cc7580bbc683af49c78945e0eb6455ce7f Mon Sep 17 00:00:00 2001
From: pvanhout <pierre.vanhoutryve at amd.com>
Date: Mon, 29 Jan 2024 11:29:14 +0100
Subject: [PATCH 1/2] [StackSafetyAnalysis] Don't call getTruncateOrZeroExtend
 for differently-sized pointers

Otherwise SCEV asserts.

Fixes SWDEV-442670
---
 llvm/lib/Analysis/StackSafetyAnalysis.cpp      | 11 +++++++++++
 .../Analysis/StackSafetyAnalysis/extend-ptr.ll | 18 ++++++++++++++++++
 2 files changed, 29 insertions(+)
 create mode 100644 llvm/test/Analysis/StackSafetyAnalysis/extend-ptr.ll

diff --git a/llvm/lib/Analysis/StackSafetyAnalysis.cpp b/llvm/lib/Analysis/StackSafetyAnalysis.cpp
index 19991c1a7baee60..b06c31f1b1c3310 100644
--- a/llvm/lib/Analysis/StackSafetyAnalysis.cpp
+++ b/llvm/lib/Analysis/StackSafetyAnalysis.cpp
@@ -273,6 +273,17 @@ ConstantRange StackSafetyLocalAnalysis::offsetFrom(Value *Addr, Value *Base) {
     return UnknownRange;
 
   auto *PtrTy = PointerType::getUnqual(SE.getContext());
+
+  // SCEV will crash if we try to extend/truncate between differently-sized
+  // pointer.
+  const auto CanExtend = [&](Type *Ty) {
+    return !Ty->isPointerTy() ||
+           DL.getTypeSizeInBits(Ty) == DL.getTypeSizeInBits(PtrTy);
+  };
+
+  if (!CanExtend(Addr->getType()) || !CanExtend(Base->getType()))
+    return UnknownRange;
+
   const SCEV *AddrExp = SE.getTruncateOrZeroExtend(SE.getSCEV(Addr), PtrTy);
   const SCEV *BaseExp = SE.getTruncateOrZeroExtend(SE.getSCEV(Base), PtrTy);
   const SCEV *Diff = SE.getMinusSCEV(AddrExp, BaseExp);
diff --git a/llvm/test/Analysis/StackSafetyAnalysis/extend-ptr.ll b/llvm/test/Analysis/StackSafetyAnalysis/extend-ptr.ll
new file mode 100644
index 000000000000000..bb9ae92c24fdda3
--- /dev/null
+++ b/llvm/test/Analysis/StackSafetyAnalysis/extend-ptr.ll
@@ -0,0 +1,18 @@
+; RUN: opt %s -disable-output -S -passes="print<stack-safety-local>" 2>&1 | FileCheck %s
+
+; Datalayout from AMDGPU, p5 is 32 bits and p is 64.
+; We used to call SCEV getTruncateOrZeroExtend on %x.ascast/%x which caused an assertion failure.
+
+; CHECK:      @a dso_preemptable
+; CHECK-NEXT:   args uses:
+; CHECK-NEXT:     x[]: full-set
+; CHECK-NEXT:   allocas uses:
+
+target datalayout = "e-p:64:64-p1:64:64-p2:32:32-p3:32:32-p4:64:64-p5:32:32-p6:32:32-p7:160:256:256:32-p8:128:128-p9:192:256:256:32-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-v2048:2048-n32:64-S32-A5-G1-ni:7:8:9"
+
+define void @a(ptr addrspace(5) %x) {
+entry:
+  %x.ascast = addrspacecast ptr addrspace(5) %x to ptr
+  store i64 0, ptr %x.ascast, align 8
+  ret void
+}

>From d91dae99d3773a9db646a7f249120098018a69be Mon Sep 17 00:00:00 2001
From: pvanhout <pierre.vanhoutryve at amd.com>
Date: Mon, 29 Jan 2024 12:16:00 +0100
Subject: [PATCH 2/2] Fix both cases

---
 llvm/lib/Analysis/StackSafetyAnalysis.cpp     | 20 +++++++++++--------
 .../StackSafetyAnalysis/extend-ptr.ll         |  3 ++-
 2 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/llvm/lib/Analysis/StackSafetyAnalysis.cpp b/llvm/lib/Analysis/StackSafetyAnalysis.cpp
index b06c31f1b1c3310..d78ef086c84e79a 100644
--- a/llvm/lib/Analysis/StackSafetyAnalysis.cpp
+++ b/llvm/lib/Analysis/StackSafetyAnalysis.cpp
@@ -274,15 +274,12 @@ ConstantRange StackSafetyLocalAnalysis::offsetFrom(Value *Addr, Value *Base) {
 
   auto *PtrTy = PointerType::getUnqual(SE.getContext());
 
-  // SCEV will crash if we try to extend/truncate between differently-sized
-  // pointer.
-  const auto CanExtend = [&](Type *Ty) {
-    return !Ty->isPointerTy() ||
-           DL.getTypeSizeInBits(Ty) == DL.getTypeSizeInBits(PtrTy);
-  };
-
-  if (!CanExtend(Addr->getType()) || !CanExtend(Base->getType()))
+  // FIXME: Pass does not deal with pointers from different address spaces that
+  // don't have the same size.
+  if (DL.getTypeSizeInBits(Addr->getType()) != PointerSize ||
+      DL.getTypeSizeInBits(Base->getType()) != PointerSize) {
     return UnknownRange;
+  }
 
   const SCEV *AddrExp = SE.getTruncateOrZeroExtend(SE.getSCEV(Addr), PtrTy);
   const SCEV *BaseExp = SE.getTruncateOrZeroExtend(SE.getSCEV(Base), PtrTy);
@@ -371,6 +368,13 @@ bool StackSafetyLocalAnalysis::isSafeAccess(const Use &U, AllocaInst *AI,
   if (isa<SCEVCouldNotCompute>(AccessSize))
     return false;
 
+  // FIXME: Pass does not deal with pointers from different address spaces that
+  // don't have the same size.
+  if (DL.getTypeSizeInBits(U->getType()) != PointerSize ||
+      DL.getTypeSizeInBits(AI->getType()) != PointerSize) {
+    return false;
+  }
+
   const auto *I = cast<Instruction>(U.getUser());
 
   auto ToCharPtr = [&](const SCEV *V) {
diff --git a/llvm/test/Analysis/StackSafetyAnalysis/extend-ptr.ll b/llvm/test/Analysis/StackSafetyAnalysis/extend-ptr.ll
index bb9ae92c24fdda3..2bfe32c654ff1d1 100644
--- a/llvm/test/Analysis/StackSafetyAnalysis/extend-ptr.ll
+++ b/llvm/test/Analysis/StackSafetyAnalysis/extend-ptr.ll
@@ -13,6 +13,7 @@ target datalayout = "e-p:64:64-p1:64:64-p2:32:32-p3:32:32-p4:64:64-p5:32:32-p6:3
 define void @a(ptr addrspace(5) %x) {
 entry:
   %x.ascast = addrspacecast ptr addrspace(5) %x to ptr
-  store i64 0, ptr %x.ascast, align 8
+  %tmp = load i64, ptr %x.ascast
+  store i64 %tmp, ptr %x.ascast, align 8
   ret void
 }