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

[Pierre van Houtryve via llvm-commits]

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

Otherwise SCEV asserts `Can't extend pointer!`
I assume the proper fix would be to teach SCEV about addrspacecast but this is a blocking issue downstream so I'm making a quick fix.

Fixes SWDEV-442670

>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] [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
+}