Merge commit 'fe18b9b43a67392e082128cdcdc2a85e7dbdf26b'

intel · Aug 1, 2024 · 27dea53 · 27dea53
2 parents 3bee61a + fe18b9b
commit 27dea53
Show file tree

Hide file tree

Showing 5 changed files with 19 additions and 8 deletions.
diff --git a/include/triton/Dialect/Triton/IR/TritonOps.td b/include/triton/Dialect/Triton/IR/TritonOps.td
@@ -859,7 +859,7 @@ def TT_MakeTensorPtrOp : TT_Op<"make_tensor_ptr",
                                 SameVariadicOperandSize,
                                 TypesMatchWith<"infer pointer type from the result type",
                                                "result", "base",
-                                               "getPointerType(getElementTypeOfTensorPointerType($_self))">]> {
+                                               "getPointerType(getElementTypeOfTensorPointerType($_self), getAddressSpace($_self))">]> {
   let summary = "Make a tensor pointer type with meta information of the parent tensor and the block specified";
 
   let description = [{

diff --git a/include/triton/Dialect/Triton/IR/Types.h b/include/triton/Dialect/Triton/IR/Types.h
@@ -22,7 +22,9 @@ unsigned getPointeeBitWidth(Type type);
 
 Type getPointeeType(Type type);
 
-Type getPointerType(Type type);
+Type getPointerType(Type type, int addressSpace = 1);
+
+int getAddressSpace(Type type);
 
 Type getElementTypeOfTensorPointerType(Type type);
 

diff --git a/lib/Dialect/Triton/IR/Types.cpp b/lib/Dialect/Triton/IR/Types.cpp
@@ -157,7 +157,16 @@ Type getPointerTypeSameShape(Type type) {
   }
 }
 
-Type getPointerType(Type type) { return PointerType::get(type, 1); }
+// upstream Triton only uses address space 1 for Pointer Type
+Type getPointerType(Type type, int addressSpace) {
+  return PointerType::get(type, addressSpace);
+}
+
+int getAddressSpace(Type type) {
+  if (auto ptrType = dyn_cast<PointerType>(type))
+    return ptrType.getAddressSpace();
+  return 1;
+}
 
 bool isTensorPointerType(Type type) {
   if (auto ptrType = dyn_cast<PointerType>(type))

diff --git a/test/TritonIntelGPU/distribute-to-warps.mlir b/test/TritonIntelGPU/distribute-to-warps.mlir
@@ -32,13 +32,13 @@ module attributes {"triton_gpu.num-warps" = 4 : i32, "triton_gpu.threads-per-war
     // CHECK:      scf.for {{.*}} iter_args([[ARG10:%.*]] = [[CST]], [[ARG11:%.*]] = [[TPTR1]], [[ARG12:%.*]] = [[TPTR2]])
     // CHECK-DAG:    [[LOAD1:%.*]] = tt.load [[ARG11]] : !tt.ptr<tensor<32x32xf16, [[WARP2]]>>
     // CHECK-DAG:    [[LOAD2:%.*]] = tt.load [[ARG12]] : !tt.ptr<tensor<32x32xf16, [[WARP2]]>>
-    // CHECK:        [[ALLOC1:%.*]] = triton_intel_gpu.alloc : <f16>
+    // CHECK:        [[ALLOC1:%.*]] = triton_intel_gpu.alloc : <f16, 3>
     // CHECK:        [[PTR1:%.*]] = tt.make_tensor_ptr [[ALLOC1]], {{.*}} {order = array<i32: 1, 0>} : <tensor<32x32xf16, [[WARP2]]>, 3>
     // CHECK:        tt.store [[PTR1]], [[LOAD1]] : !tt.ptr<tensor<32x32xf16, [[WARP2]]>, 3>
     // CHECK:        gpu.barrier
     // CHECK:        [[PTR2:%.*]] = tt.make_tensor_ptr [[ALLOC1]], {{.*}} {order = array<i32: 1, 0>} : <tensor<64x32xf16, #triton_gpu.dot_op<{opIdx = 0, parent = [[WARP1]]}>>, 3>
     // CHECK:        [[LOAD3:%.*]] = tt.load [[PTR2]] : !tt.ptr<tensor<64x32xf16, #triton_gpu.dot_op<{opIdx = 0, parent = [[WARP1]]}>>, 3>
-    // CHECK:        [[ALLOC2:%.*]] = triton_intel_gpu.alloc : <f16>
+    // CHECK:        [[ALLOC2:%.*]] = triton_intel_gpu.alloc : <f16, 3>
     // CHECK:        [[PTR3:%.*]] = tt.make_tensor_ptr [[ALLOC2]], {{.*}} {order = array<i32: 1, 0>} : <tensor<32x32xf16, [[WARP2]]>, 3>
     // CHECK:        tt.store [[PTR3]], [[LOAD2]] : !tt.ptr<tensor<32x32xf16, [[WARP2]]>, 3>
     // CHECK:        gpu.barrier

diff --git a/third_party/intel/lib/TritonIntelGPUTransforms/DistributeToWarps.cpp b/third_party/intel/lib/TritonIntelGPUTransforms/DistributeToWarps.cpp
@@ -288,9 +288,9 @@ void distributeConvertLayoutOp(ttg::ConvertLayoutOp op, Value warpId,
 
   // FIXME: allocOp may carry the size info.
   OpBuilder b(op);
-  auto baseType = tt::PointerType::get(
-      oldSrcType.getElementType(),
-      triton::TritonGEN::TritonGENMemorySpace::kCrossWorkgroup);
+  auto baseType =
+      tt::PointerType::get(oldSrcType.getElementType(),
+                           triton::TritonGEN::TritonGENMemorySpace::kWorkgroup);
   auto base = b.create<ttgi::AllocOp>(loc, baseType);
 
   SmallVector<Value> shape;