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

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-nvptx

Author: Kalomidin (Kalomidin)

<details>
<summary>Changes</summary>



---
Full diff: https://github.com/llvm/llvm-project/pull/193701.diff


2 Files Affected:

- (modified) llvm/lib/Target/NVPTX/NVPTXTargetTransformInfo.cpp (+31-6) 
- (modified) llvm/lib/Target/NVPTX/NVPTXTargetTransformInfo.h (+1) 


``````````diff
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,

``````````

</details>


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