Support non-power-of-two TLE pipe capacities

Author: Prophet-Hongyi

python/test/tle/unit/test_tle.py

@@ -275,6 +275,40 @@ def test_buffered_tensor_type_attributes(self):
         assert hasattr(tle.gpu.buffered_tensor, 'make_permute')
         assert hasattr(tle.gpu.buffered_tensor, 'slot')
 
+    def test_memory_descriptor_shape_is_not_a_register_block_shape(self):
+        """Memory descriptors accept pipeline capacities that register blocks reject."""
+        buffer, _ = self._make_buffer([6, 1, 64, 128])
+
+        assert buffer.type.shape == (6, 1, 64, 128)
+        assert buffer.type.numel == 6 * 64 * 128
+        assert not buffer.type.is_block()
+        with pytest.raises(ValueError):
+            tl.block_type(tl.float16, [6, 1, 64, 128])
+
+    def test_memory_descriptor_shape_still_rejects_invalid_dimensions(self):
+        with pytest.raises(ValueError, match="must be positive"):
+            self._make_buffer([0, 16])
+        with pytest.raises(TypeError, match="must be an integer"):
+            self._make_buffer(["6", 16])
+
+    def test_buffered_tensor_type_equality_includes_dtype_and_storage(self):
+        buffer, semantic = self._make_buffer([6, 16])
+        layout = buffer.type.layout
+        fp32_type = tle.gpu.buffered_tensor_type(tl.float32, [6, 16], tle.gpu.smem, layout, semantic)
+        tmem_type = tle.gpu.buffered_tensor_type(tl.float16, [6, 16], tle.gpu.tmem, layout, semantic)
+
+        assert buffer.type != fp32_type
+        assert buffer.type != tmem_type
+
+    def test_barrier_descriptor_accepts_non_power_of_two_capacity(self):
+        semantic = self._FakeSemantic()
+        layout = tle.gpu.swizzled_shared_layout.make_default(2)
+        barrier_type = tle.gpu.barrier_type(6, 1, "all", None, layout, semantic)
+
+        assert barrier_type.shape == (6, 1)
+        assert barrier_type.numel == 6
+        assert not barrier_type.is_block()
+
     def test_buffered_tensor_slot_indexes_leading_dimension(self):
         """slot(stage) returns a typed view with the leading stage dimension removed."""
         buffer, semantic = self._make_buffer([4, 16, 32])

python/test/tle/unit/test_tle_gpu_slot.py

@@ -41,6 +41,16 @@ def _slot_local_ptr_store_kernel(out_ptr, BLOCK: tl.constexpr):
     tl.store(out_ptr + idx, vals)
 
 
+@triton.jit
+def _non_power_of_two_stage_slot_kernel(out_ptr, BLOCK: tl.constexpr):
+    idx = tl.arange(0, BLOCK)
+    smem = tle.gpu.alloc([6, BLOCK], dtype=tl.int32, layout=None, scope=tle.gpu.smem, nv_mma_shared_layout=False)
+    slot = smem.slot(5)
+    ptrs = tle.gpu.local_ptr(slot, (idx, ))
+    tl.store(ptrs, idx + 11)
+    tl.store(out_ptr + idx, tl.load(ptrs))
+
+
 def test_buffered_tensor_slot_lowers_to_memdesc_index_and_executes():
     block = 64
     out = torch.empty((block, ), device="cuda", dtype=torch.int32)
@@ -54,3 +64,17 @@ def test_buffered_tensor_slot_lowers_to_memdesc_index_and_executes():
     _slot_local_ptr_store_kernel[(1, )](out, BLOCK=block, num_warps=4)
     expected = torch.arange(0, block, device="cuda", dtype=torch.int32) + 7
     torch.testing.assert_close(out, expected, atol=0, rtol=0)
+
+
+def test_non_power_of_two_stage_descriptor_lowers_and_executes():
+    block = 64
+    out = torch.empty((block, ), device="cuda", dtype=torch.int32)
+
+    compiled = _non_power_of_two_stage_slot_kernel.warmup(out, BLOCK=block, grid=(1, ), num_warps=4)
+    ttgir = compiled.asm["ttgir"]
+    assert "ttg.memdesc_index" in ttgir
+    assert "!ttg.memdesc<6x64xi32" in ttgir
+
+    _non_power_of_two_stage_slot_kernel[(1, )](out, BLOCK=block, num_warps=4)
+    expected = torch.arange(0, block, device="cuda", dtype=torch.int32) + 11
+    torch.testing.assert_close(out, expected, atol=0, rtol=0)

python/triton/experimental/tle/language/gpu/types.py

@@ -49,6 +49,50 @@ def to_ir(self, builder: ir.builder) -> None:
 READY = "ready"
 
 
+class memory_descriptor_type(tl.base_type):
+    """Base type for memory-backed descriptors with an arbitrary static shape.
+
+    TLE memory descriptors are not register-resident Triton blocks.  In
+    particular, pipeline capacities such as 3 or 6 are valid leading
+    dimensions, and descriptor allocation sizes are governed by the target
+    memory space rather than ``TRITON_MAX_TENSOR_NUMEL``.
+    """
+
+    def __init__(self, element_ty: tl.dtype, shape: List):
+        if not isinstance(element_ty, tl.dtype):
+            raise TypeError(f"element_ty has type `{type(element_ty).__name__}`; expected `dtype`.")
+        if not isinstance(shape, (list, tuple)) or not shape:
+            raise TypeError("memory descriptor shape must be a non-empty list or tuple")
+
+        normalized_shape = []
+        for index, dim in enumerate(shape):
+            dim = tl._unwrap_if_constexpr(dim)
+            if not isinstance(dim, int):
+                raise TypeError(f"Shape element {index} must be an integer, got {type(dim).__name__}")
+            if dim <= 0:
+                raise ValueError(f"Shape element {index} must be positive, got {dim}")
+            normalized_shape.append(dim)
+
+        self.element_ty = element_ty
+        self.shape = tuple(normalized_shape)
+        self.numel = 1
+        for dim in self.shape:
+            self.numel *= dim
+        self.name = f'<{self.shape}, {self.element_ty}>'
+
+    @staticmethod
+    def is_block():
+        return False
+
+    @property
+    def scalar(self):
+        return self.element_ty
+
+    @property
+    def nbytes(self):
+        return self.numel * (self.element_ty.primitive_bitwidth // 8)
+
+
 def _storage_to_memdesc_space(storage: scope) -> str:
     if storage is smem:
         return "smem"
@@ -356,7 +400,7 @@ def make_permute(self, handle, dims):
         )
 
 
-class buffered_tensor_type(tl.block_type):
+class buffered_tensor_type(memory_descriptor_type):
 
     def __init__(self, element_ty: tl.dtype, shape: List, storage: scope, layout: Optional[shared_layout] = None,
                  semantic: TritonSemantic = None, alloc_shape: List = None):
@@ -401,8 +445,13 @@ def mangle(self) -> str:
     def __str__(self) -> str:
         return f"buffered_tensor_<{self.element_ty}, {self.shape}, {self.layout}, {self.alloc_shape}, >"
 
+    def with_element_ty(self, scalar_ty: tl.dtype):
+        return buffered_tensor_type(scalar_ty, self.shape, self.storage, self.layout, self.semantic,
+                                    alloc_shape=self.alloc_shape)
+
     def __eq__(self, other) -> bool:
-        if not (type(self) is type(other) and self.shape == other.shape and self.layout == other.layout
+        if not (type(self) is type(other) and self.element_ty == other.element_ty and self.shape == other.shape
+                and self.storage is other.storage and self.layout == other.layout
                 and self.alloc_shape == other.alloc_shape):
             return False
         self_shard = getattr(self, "_tle_remote_shard_id", None)
@@ -515,7 +564,7 @@ def __getitem__(self, index, _semantic=None):
                        allocation_key=self.allocation_key)
 
 
-class barrier_type(tl.block_type):
+class barrier_type(memory_descriptor_type):
 
     def __init__(
         self,

third_party/tle/dialect/lib/Transforms/TleLowerPipeToNvws.cpp

@@ -1237,14 +1237,22 @@ static PipeState createPipeState(PipeCreateOp op) {
         ttg::MemDescType::get({1}, builder.getI32Type(), closeTagSlotEncoding,
                               sharedMemorySpace, /*mutableMemory=*/true);
 
-    RankedTensorType closeTagArrayTensorType =
-        getCloseTagTensorType(op, builder, {capacity, 1});
-    Value initialCloseTags =
-        createCloseTagTensor(builder, loc, closeTagArrayTensorType,
-                             /*value=*/false);
-    closeTags = ttg::LocalAllocOp::create(builder, loc, closeTagArrayType,
-                                          initialCloseTags);
     closeTagTensorType = getCloseTagTensorType(op, builder, {1});
+    closeTags =
+        ttg::LocalAllocOp::create(builder, loc, closeTagArrayType, Value());
+    // A pipe capacity describes shared-memory slots, not a distributed
+    // register tensor.  Initializing the whole ring through a
+    // tensor<capacity x 1> would therefore reintroduce Triton's power-of-two
+    // register-block restriction for otherwise valid capacities such as 3 or
+    // 6.  Initialize one scalar memdesc slot at a time instead.
+    for (int64_t stage = 0; stage < capacity; ++stage) {
+      Value stageValue = arith::ConstantIntOp::create(builder, loc, stage, 32);
+      Value slot = ttg::MemDescIndexOp::create(builder, loc, closeTagSlotType,
+                                               closeTags, stageValue);
+      Value tag = createCloseTagTensor(builder, loc, closeTagTensorType,
+                                       /*value=*/false);
+      ttg::LocalStoreOp::create(builder, loc, tag, slot);
+    }
   }
   Value token = ttnvws::CreateTokenOp::create(
       builder, loc, static_cast<uint32_t>(capacity),

third_party/tle/test/GPU/test_tle_lower_pipe_to_nvws.mlir

@@ -25,6 +25,18 @@
 #smem = #ttg.shared_memory
 
 module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32} {
+  // CHECK-LABEL: tt.func @lower_non_power_of_two_pipe_capacity
+  tt.func @lower_non_power_of_two_pipe_capacity(%a: !ttg.memdesc<6x16xf16, #shared, #smem, mutable>) {
+    // CHECK: %[[TAGS:.*]] = ttg.local_alloc
+    // CHECK-SAME: !ttg.memdesc<6x1xi32
+    // CHECK-NOT: tensor<6x1xi32
+    // CHECK-COUNT-6: ttg.local_store
+    // CHECK: nvws.create_token
+    // CHECK-SAME: numBuffers = 6
+    tle.pipe.create %a {capacity = 6 : i32, pipe_name = "six_stage", field_names = ["a"], scope = "cta"} : !ttg.memdesc<6x16xf16, #shared, #smem, mutable>
+    tt.return
+  }
+
   // CHECK-LABEL: tt.func @lower_pipe_to_nvws
   tt.func @lower_pipe_to_nvws(%a: !ttg.memdesc<2x16xf16, #shared, #smem, mutable>) {
     %c0 = arith.constant 0 : i32