typeConverter to llvm support addressSpace attribute (#5951)

MemDescType support attribute memorySpace, but TypeConverter to LLVM not
used, this pull request try to fix it.

Co-authored-by: kaixuan.zhang <[email protected]>

Author: zkxvb

include/triton/Conversion/TritonGPUToLLVM/TargetInfoBase.h

@@ -89,6 +89,8 @@ class TargetInfoBase {
 
   virtual int getSharedAddressSpace() const = 0;
 
+  virtual int getAddressSpace(Attribute addressSpace) const = 0;
+
   virtual bool supportVectorizedAtomics() const = 0;
 
   // Helper used by targets to annotate store operations during lowering to

lib/Conversion/TritonGPUToLLVM/TypeConverter.cpp

@@ -54,8 +54,8 @@ Type TritonGPUToLLVMTypeConverter::convertMemDescType(
     MemDescType type, const TargetInfoBase &targetInfo) {
   auto ctx = type.getContext();
   // base ptr
-  auto ptrType =
-      LLVM::LLVMPointerType::get(ctx, targetInfo.getSharedAddressSpace());
+  auto ptrType = LLVM::LLVMPointerType::get(
+      ctx, targetInfo.getAddressSpace(type.getMemorySpace()));
 
   if (isa<triton::nvidia_gpu::TensorMemoryEncodingAttr,
           triton::nvidia_gpu::TensorMemoryScalesEncodingAttr>(

lib/Dialect/Triton/IR/Ops.cpp

@@ -882,7 +882,7 @@ void MakeTensorPtrOp::build(OpBuilder &builder, OperationState &state,
   auto tensorType = RankedTensorType::get(
       SmallVector<int64_t>(tensorShape.begin(), tensorShape.end()),
       pointerType.getPointeeType());
-  auto result = PointerType::get(tensorType, 1);
+  auto result = PointerType::get(tensorType, pointerType.getAddressSpace());
 
   return build(builder, state, result, base, shape, strides, offsets,
                builder.getDenseI32ArrayAttr(order));

third_party/amd/lib/TritonAMDGPUToLLVM/TargetInfo.cpp

@@ -447,6 +447,16 @@ void TargetInfo::assertFail(RewriterBase &rewriter, Location loc,
 
 int TargetInfo::getSharedAddressSpace() const { return 3; }
 
+int TargetInfo::getAddressSpace(Attribute addressSpace) const {
+  int spaceId = 0;
+  if (isa<triton::gpu::SharedMemorySpaceAttr>(addressSpace)) {
+    spaceId = 3;
+  } else {
+    llvm::report_fatal_error("Only support SharedMemorySpace for now");
+  }
+  return spaceId;
+}
+
 bool TargetInfo::supportVectorizedAtomics() const {
   // Note: not currently tested or used, but AMD generally supports vectorized
   // atomics.

third_party/amd/lib/TritonAMDGPUToLLVM/TargetInfo.h

@@ -65,8 +65,11 @@ class TargetInfo : public mlir::triton::TargetInfoBase {
 
   void assertFail(RewriterBase &rewriter, Location loc, StringRef message,
                   StringRef file, StringRef func, int line) const override;
+
   int getSharedAddressSpace() const override;
 
+  int getAddressSpace(Attribute addressSpace) const override;
+
   bool supportVectorizedAtomics() const override;
 
   void storeOpAnnotation(triton::gpu::LocalStoreOp op, size_t localStoreOpCount,

third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TargetInfo.cpp

@@ -617,6 +617,18 @@ void TargetInfo::assertFail(RewriterBase &rewriter, Location loc,
 
 int TargetInfo::getSharedAddressSpace() const { return 3; }
 
+int TargetInfo::getAddressSpace(Attribute addressSpace) const {
+  int spaceId = 0;
+  if (isa<triton::gpu::SharedMemorySpaceAttr,
+          triton::nvidia_gpu::TensorMemorySpaceAttr>(addressSpace)) {
+    spaceId = 3;
+  } else {
+    llvm::report_fatal_error(
+        "Only support SharedMemorySpace, TensorMemorySpace for now");
+  }
+  return spaceId;
+}
+
 bool TargetInfo::supportVectorizedAtomics() const {
   return computeCapability >= 90 && ptxVersion >= 81;
 }

third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TargetInfo.h

@@ -58,8 +58,11 @@ class TargetInfo : public mlir::triton::TargetInfoBase {
 
   void assertFail(RewriterBase &rewriter, Location loc, StringRef message,
                   StringRef file, StringRef func, int line) const override;
+
   int getSharedAddressSpace() const override;
 
+  int getAddressSpace(Attribute addressSpace) const override;
+
   bool supportVectorizedAtomics() const override;
 
   int getPtxVersion() const { return ptxVersion; }