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

[via llvm-commits]

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

None

>From 83fb51a148e7e934bab177d594af2030815c0dc1 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 | 37 ++++++++++++++++---
 .../Target/NVPTX/NVPTXTargetTransformInfo.h   |  1 +
 2 files changed, 32 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/Target/NVPTX/NVPTXTargetTransformInfo.cpp b/llvm/lib/Target/NVPTX/NVPTXTargetTransformInfo.cpp
index a491d0ed4a912..582f0bab315e7 100644
--- a/llvm/lib/Target/NVPTX/NVPTXTargetTransformInfo.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXTargetTransformInfo.cpp
@@ -31,11 +31,12 @@ using namespace llvm;
 // Whether the given intrinsic reads threadIdx.x/y/z.
 static bool readsThreadIndex(const IntrinsicInst *II) {
   switch (II->getIntrinsicID()) {
-    default: return false;
-    case Intrinsic::nvvm_read_ptx_sreg_tid_x:
-    case Intrinsic::nvvm_read_ptx_sreg_tid_y:
-    case Intrinsic::nvvm_read_ptx_sreg_tid_z:
-      return true;
+  default:
+    return false;
+  case Intrinsic::nvvm_read_ptx_sreg_tid_x:
+  case Intrinsic::nvvm_read_ptx_sreg_tid_y:
+  case Intrinsic::nvvm_read_ptx_sreg_tid_z:
+    return true;
   }
 }
 
@@ -652,7 +653,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,