[llvm] [NVPTX] Allow mixed address-space pointer arguments in kernel lowering (PR #193701)

[via llvm-commits]

https://github.com/Kalomidin updated https://github.com/llvm/llvm-project/pull/193701

>From a84ae76ea320776386a831cbdac8bc8ae08b0c18 Mon Sep 17 00:00:00 2001
From: Kalomidin <kalomidin at Kalomidins-MacBook-Pro.local>
Date: Thu, 23 Apr 2026 17:25:43 +0800
Subject: [PATCH] [NVPTX] Allow mixed address-space pointer arguments in kernel
 lowering

---
 .../Target/NVPTX/NVPTXTargetTransformInfo.cpp | 26 ++++++++++++++++++-
 .../Target/NVPTX/NVPTXTargetTransformInfo.h   |  1 +
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Target/NVPTX/NVPTXTargetTransformInfo.cpp b/llvm/lib/Target/NVPTX/NVPTXTargetTransformInfo.cpp
index a491d0ed4a912..66b282fb57b4a 100644
--- a/llvm/lib/Target/NVPTX/NVPTXTargetTransformInfo.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXTargetTransformInfo.cpp
@@ -652,7 +652,31 @@ unsigned NVPTXTTIImpl::getAssumedAddrSpace(const Value *V) const {
         return ADDRESS_SPACE_LOCAL;
     }
   }
-
+  if (int AS = getPointerLoadAddressSpace(V); AS != -1) {
+    return AS;
+  }
+  return -1;
+}
+/* 
+ * geparg = getelementptr T, ptr %arg, i32 0, i32 0
+ * ptr a = load ptr %geparg
+ * we can expect AS of a to be global
+ */
+int NVPTXTTIImpl::getPointerLoadAddressSpace(const Value *V) const {
+  auto *Load = dyn_cast<LoadInst>(V);
+  if (!Load)
+    return -1;
+  auto *Ptr = Load->getPointerOperand();
+  if (!Ptr)
+    return -1;
+  // if it is argument, return GM AS
+  if (const auto *Arg = dyn_cast<Argument>(Ptr)) {
+    if (isKernelFunction(*Arg->getParent()))
+      return ADDRESS_SPACE_GLOBAL;
+  }
+  if (auto *GEP = dyn_cast<GetElementPtrInst>(Ptr)) {
+    return getAssumedAddrSpace(GEP->getPointerOperand());
+  }
   return -1;
 }
 
diff --git a/llvm/lib/Target/NVPTX/NVPTXTargetTransformInfo.h b/llvm/lib/Target/NVPTX/NVPTXTargetTransformInfo.h
index 8bdafd6b905f1..87a816e939f44 100644
--- a/llvm/lib/Target/NVPTX/NVPTXTargetTransformInfo.h
+++ b/llvm/lib/Target/NVPTX/NVPTXTargetTransformInfo.h
@@ -212,6 +212,7 @@ class NVPTXTTIImpl final : public BasicTTIImplBase<NVPTXTTIImpl> {
   Value *rewriteIntrinsicWithAddressSpace(IntrinsicInst *II, Value *OldV,
                                           Value *NewV) const override;
   unsigned getAssumedAddrSpace(const Value *V) const override;
+  int getPointerLoadAddressSpace(const Value *V) const;
 
   void collectKernelLaunchBounds(
       const Function &F,