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

Otherwise SCEV asserts `Can't extend pointer!`

Fixes SWDEV-442670

Author: Pierre-vh

llvm/lib/Analysis/StackSafetyAnalysis.cpp

@@ -243,6 +243,14 @@ class StackSafetyLocalAnalysis {
 
   const ConstantRange UnknownRange;
 
+  /// FIXME: This function is a bandaid, it's only needed
+  /// because this pass doesn't handle address spaces of different pointer
+  /// sizes.
+  ///
+  /// \returns \p Val's SCEV as a pointer of AS zero, or nullptr if it can't be
+  /// converted to AS 0.
+  const SCEV *getSCEVAsPointer(Value *Val);
+
   ConstantRange offsetFrom(Value *Addr, Value *Base);
   ConstantRange getAccessRange(Value *Addr, Value *Base,
                                const ConstantRange &SizeRange);
@@ -268,13 +276,29 @@ class StackSafetyLocalAnalysis {
   FunctionInfo<GlobalValue> run();
 };
 
+const SCEV *StackSafetyLocalAnalysis::getSCEVAsPointer(Value *Val) {
+  Type *ValTy = Val->getType();
+
+  // We don't handle targets with multiple address spaces.
+  if (!ValTy->isPointerTy()) {
+    auto *PtrTy = PointerType::getUnqual(SE.getContext());
+    return SE.getTruncateOrZeroExtend(SE.getSCEV(Val), PtrTy);
+  }
+
+  if (ValTy->getPointerAddressSpace() != 0)
+    return nullptr;
+  return SE.getSCEV(Val);
+}
+
 ConstantRange StackSafetyLocalAnalysis::offsetFrom(Value *Addr, Value *Base) {
   if (!SE.isSCEVable(Addr->getType()) || !SE.isSCEVable(Base->getType()))
     return UnknownRange;
 
-  auto *PtrTy = PointerType::getUnqual(SE.getContext());
-  const SCEV *AddrExp = SE.getTruncateOrZeroExtend(SE.getSCEV(Addr), PtrTy);
-  const SCEV *BaseExp = SE.getTruncateOrZeroExtend(SE.getSCEV(Base), PtrTy);
+  const SCEV *AddrExp = getSCEVAsPointer(Addr);
+  const SCEV *BaseExp = getSCEVAsPointer(Base);
+  if (!AddrExp || !BaseExp)
+    return UnknownRange;
+
   const SCEV *Diff = SE.getMinusSCEV(AddrExp, BaseExp);
   if (isa<SCEVCouldNotCompute>(Diff))
     return UnknownRange;
@@ -362,13 +386,11 @@ bool StackSafetyLocalAnalysis::isSafeAccess(const Use &U, AllocaInst *AI,
 
   const auto *I = cast<Instruction>(U.getUser());
 
-  auto ToCharPtr = [&](const SCEV *V) {
-    auto *PtrTy = PointerType::getUnqual(SE.getContext());
-    return SE.getTruncateOrZeroExtend(V, PtrTy);
-  };
+  const SCEV *AddrExp = getSCEVAsPointer(U.get());
+  const SCEV *BaseExp = getSCEVAsPointer(AI);
+  if (!AddrExp || !BaseExp)
+    return false;
 
-  const SCEV *AddrExp = ToCharPtr(SE.getSCEV(U.get()));
-  const SCEV *BaseExp = ToCharPtr(SE.getSCEV(AI));
   const SCEV *Diff = SE.getMinusSCEV(AddrExp, BaseExp);
   if (isa<SCEVCouldNotCompute>(Diff))
     return false;

llvm/test/Analysis/StackSafetyAnalysis/extend-ptr.ll

@@ -0,0 +1,19 @@
+; 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
+  %tmp = load i64, ptr %x.ascast
+  store i64 %tmp, ptr %x.ascast, align 8
+  ret void
+}