[TLE][MTHREADS] Support barrier/wgmma/pipe/ws primitives (#944)

* [TLE][MTHREADS] Add memdesc slot binding with constant bounds checks

* [TLE][MTHREADS] Add backend-local TLE barrier support

* [TLE][MTHREADS] Support completion barriers for TME copies

* [TLE][MTHREADS] Add explicit warp-specialize container support

* [TLE][MTHREADS] Add warp-specialize frontend diagnostics

* [TLE][MTHREADS] Add explicit warp-specialize integration coverage

* [TLE][MTHREADS] Lower barrier allocations to hardware IDs

* [TLE][MTHREADS] Synchronize barrier initialization for all kernels

* [TLE][MTHREADS] Lower explicit TME completion transactions

* [TLE][MTHREADS] Complete barrier consumer protocol

* [TLE][MTHREADS] Lower WGMMA operations to SQMMA

* [TLE][MTHREADS] Add single-field pipe lowering

* [TLE][MTHREADS] Add static warp-specialize lowering

* [TLE][MTHREADS] Support three-stage buffered pipelines

* [TLE][MTHREADS] Enable automatic SQMMA shared layouts

Author: Kylin1207

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

@@ -24,7 +24,11 @@
 from typing import Optional, Sequence
 from enum import Enum
 from . import types as tle
+from .mthreads import common as mthreads_common
+from .mthreads import buffer as mthreads_buffer
 from .mthreads import copy as mthreads_copy
+from .mthreads import warp_specialize as mthreads_warp_specialize
+from .mthreads import wgmma as mthreads_wgmma
 from .iluvatar import copy as iluvatar_copy
 from triton.compiler.code_generator import flatten_values_to_ir, unflatten_ir_values
 
@@ -45,6 +49,9 @@
 def _mark_wgmma_user_promise(_semantic, _generator):
     if _generator is None or _semantic is None:
         return
+    # This module attribute is NVIDIA-specific and must not leak into mthreads IR.
+    if mthreads_common.enabled():
+        return
     _generator.module.set_attr(
         _WGMMA_PIPELINE_MODE_ATTR,
         _semantic.builder.get_string_attr(_WGMMA_PIPELINE_MODE_USER_PROMISE),
@@ -186,8 +193,12 @@ def warp_specialize(functions_and_args, worker_num_warps, worker_num_regs, _sema
     if _generator is None:
         raise ValueError("warp_specialize requires a Triton code generator")
     functions_and_args = tl._unwrap_if_constexpr(functions_and_args)
-    worker_num_warps = [tl._unwrap_if_constexpr(w) for w in tl._unwrap_if_constexpr(worker_num_warps)]
-    worker_num_regs = [tl._unwrap_if_constexpr(r) for r in tl._unwrap_if_constexpr(worker_num_regs)]
+    mthreads_enabled = mthreads_common.enabled()
+    if mthreads_enabled:
+        worker_num_warps, worker_num_regs = mthreads_warp_specialize.normalize_config(worker_num_warps, worker_num_regs)
+    else:
+        worker_num_warps = [tl._unwrap_if_constexpr(w) for w in tl._unwrap_if_constexpr(worker_num_warps)]
+        worker_num_regs = [tl._unwrap_if_constexpr(r) for r in tl._unwrap_if_constexpr(worker_num_regs)]
     if len(functions_and_args) < 1:
         raise ValueError("warp_specialize requires at least a default partition function")
     num_partitions = len(functions_and_args) - 1
@@ -200,8 +211,9 @@ def warp_specialize(functions_and_args, worker_num_warps, worker_num_regs, _sema
 
     builder = _semantic.builder
     insert_pt = builder.get_insertion_point()
-    inline_user_promise = _is_wgmma_user_promise_marked(_generator)
-    if inline_user_promise:
+    if mthreads_enabled:
+        call_jit_function = mthreads_warp_specialize.partition_function_caller(_generator)
+    elif _is_wgmma_user_promise_marked(_generator):
         call_jit_function = _generator.inline_JitFunction
     else:
         call_jit_function = _generator.call_JitFunction
@@ -225,7 +237,10 @@ def warp_specialize(functions_and_args, worker_num_warps, worker_num_regs, _sema
     worker_arg_handles, worker_items = _deduplicate_warp_specialize_captures(worker_items)
 
     builder.restore_insertion_point(insert_pt)
-    ws_op = builder.create_warp_specialize(result_types, worker_arg_handles, worker_num_warps)
+    if mthreads_enabled:
+        ws_op = mthreads_warp_specialize.create_op(builder, result_types, worker_num_warps)
+    else:
+        ws_op = builder.create_warp_specialize(result_types, worker_arg_handles, worker_num_warps)
     real_default_block = builder.create_block_with_parent(ws_op.get_default_region(), [])
     default_block.merge_block_before(real_default_block)
     default_block = real_default_block
@@ -234,7 +249,10 @@ def warp_specialize(functions_and_args, worker_num_warps, worker_num_regs, _sema
     ws_op.set_requested_registers(worker_num_regs)
 
     builder.create_block_with_parent(ws_op.get_partition_op_holder(), [])
-    partitions_op = builder.create_warp_specialize_partitions(num_partitions)
+    if mthreads_enabled:
+        partitions_op = mthreads_warp_specialize.create_partitions(builder, worker_arg_handles, num_partitions)
+    else:
+        partitions_op = builder.create_warp_specialize_partitions(num_partitions)
     partition_arg_types = [arg.get_type() for arg in worker_arg_handles]
     for idx, (worker_fn, worker_args, flattened, remapped) in enumerate(worker_items):
         block = builder.create_block_with_parent(partitions_op.get_region(idx), partition_arg_types)
@@ -288,6 +306,9 @@ def alloc(
         dtype: Data type
         layout: Memory layout encoding (optional)
         scope: Storage type (default to shared memory)
+        nv_mma_shared_layout: Select an MMA-consumer-defined shared layout when
+            ``layout`` is None. On mthreads this is materialized by the SQMMA
+            lowering rather than as an NVIDIA encoding.
         _semantic: Semantic analyzer (internal use)
 
     Returns:
@@ -331,16 +352,22 @@ def alloc(
 
     # Map scope to storage (backward compatibility)
     storage = scope
+    mthreads_auto_sqmma_shared_layout = (mthreads_common.enabled() and storage == tle.smem
+                                         and mthreads_wgmma.use_auto_shared_layout(layout, nv_mma_shared_layout))
 
     try:
         unwrapped_shape = [tl._unwrap_if_constexpr(dim) for dim in shape]
         full_shape = unwrapped_shape
+        use_mthreads_buffer = (mthreads_common.enabled() and mthreads_buffer.needs_non_power_of_two_leading_dim(
+            _semantic.builder, unwrapped_shape))
+        if use_mthreads_buffer:
+            mthreads_buffer.validate_shape(unwrapped_shape)
         dtype = tl._unwrap_if_constexpr(dtype)
         elem_type = dtype.to_ir(_semantic.builder)
 
         if layout is None:
             if storage == tle.smem:
-                if not nv_mma_shared_layout:
+                if mthreads_auto_sqmma_shared_layout or not nv_mma_shared_layout:
                     layout = tle.swizzled_shared_layout.make_default(rank=len(shape))
                     layout_handle = _semantic.builder.make_swizzled_shared_encoding_attr(
                         layout.vectorSize,
@@ -383,9 +410,20 @@ def alloc(
                 tensor_handle = _semantic.builder.create_local_alloc(mutable_ty, init_value.handle)
             else:
                 tensor_handle = _semantic.builder.create_local_alloc(full_shape, elem_type, layout_handle)
+            if mthreads_auto_sqmma_shared_layout:
+                mthreads_wgmma.mark_auto_shared_layout(_semantic.builder, tensor_handle)
         else:
             raise ValueError(f"Storage type {storage} not yet supported")
 
+        if use_mthreads_buffer:
+            return mthreads_buffer.create_buffered_tensor(
+                tensor_handle,
+                dtype,
+                unwrapped_shape,
+                storage,
+                layout,
+                _semantic,
+            )
         return tle.buffered_tensor(tensor_handle, dtype, unwrapped_shape, storage, layout, _semantic)
 
     except Exception as e:
@@ -773,16 +811,23 @@ def wgmma(
     """
     trans_a = _require_wgmma_bool(trans_a, "trans_a")
     trans_b = _require_wgmma_bool(trans_b, "trans_b")
-    a, b = _canonicalize_wgmma_operands(a, b, trans_a, trans_b, _semantic)
+    mthreads_enabled = mthreads_common.enabled()
+    if mthreads_enabled:
+        mthreads_wgmma.validate_operands(a, b, acc, trans_a, trans_b)
+    else:
+        a, b = _canonicalize_wgmma_operands(a, b, trans_a, trans_b, _semantic)
 
     m, k = [int(tl._unwrap_if_constexpr(dim)) for dim in a.type.shape]
     k_b, n = [int(tl._unwrap_if_constexpr(dim)) for dim in b.type.shape]
     if k != k_b:
         raise ValueError(f"wgmma shape mismatch: a is {a.type.shape}, b is {b.type.shape}")
-    if m < 64 or m % 64 != 0:
-        raise ValueError("wgmma result M dimension must be divisible by 64")
-    if n < 8 or n % 8 != 0:
-        raise ValueError("wgmma result N dimension must be divisible by 8")
+    if mthreads_enabled:
+        mthreads_wgmma.validate_dimensions(m, n)
+    else:
+        if m < 64 or m % 64 != 0:
+            raise ValueError("wgmma result M dimension must be divisible by 64")
+        if n < 8 or n % 8 != 0:
+            raise ValueError("wgmma result N dimension must be divisible by 8")
     if k < 16:
         raise ValueError("wgmma K dimension must be at least 16")
 
@@ -813,6 +858,9 @@ def wgmma(
         if max_num_imprecise_acc < 0:
             raise ValueError("max_num_imprecise_acc must be non-negative")
 
+    if mthreads_enabled:
+        mthreads_wgmma.validate_options(max_num_imprecise_acc, out_dtype)
+
     ret_scalar_ty = _wgmma_ret_scalar_ty(a.dtype, out_dtype)
     ret_ty = tl.block_type(ret_scalar_ty, [m, n])
     builder = _semantic.builder
@@ -851,6 +899,8 @@ def wgmma_wait(pendings, acc=None, _semantic=None, _generator=None) -> tl.tensor
     pendings = _require_wgmma_int(pendings, "pendings")
     if pendings < 0:
         raise ValueError("wgmma_wait pendings must be non-negative")
+    if mthreads_common.enabled():
+        mthreads_wgmma.validate_wait_pendings(pendings)
     if not isinstance(acc, tl.tensor):
         raise ValueError(f"wgmma_wait acc must be a tl.tensor, got {type(acc).__name__}")
     result = _semantic.builder.create_tle_wgmma_wait(acc.handle, pendings)
@@ -921,7 +971,7 @@ def copy(
             tle.copy(tma_desc, local_buf, [64, 64], [x_offset, y_offset], barrier=bar)
             tle.gpu.barrier_wait(bar, phaseIdx=0)
     """
-    mthreads_enabled = mthreads_copy.enabled()
+    mthreads_enabled = mthreads_common.enabled()
     iluvatar_enabled = iluvatar_copy.enabled()
 
     def normcopy(
@@ -1078,9 +1128,12 @@ def tmacopy(
             raise ValueError("copy barrier is only supported for TMA global-to-shared copy")
         return normcopy(src, dst, shape, direction, _semantic)
     if mthreads_enabled:
+        barrier_slot = None
         if barrier is not None:
-            raise ValueError("TMA copy barrier is only supported on NVIDIA backend")
-        return mthreads_copy.tmacopy(src, dst, direction, shape, offsets, _semantic)
+            if direction != CopyDirection.GM_TO_LOCAL:
+                raise ValueError("TMA copy barrier is only supported for global-to-shared TMA copy")
+            barrier_slot = _tma_completion_barrier_slot(barrier, _semantic)
+        return mthreads_copy.tmacopy(src, dst, direction, shape, offsets, barrier_slot, _semantic)
     else:
         return tmacopy(src, dst, direction, shape, offsets, barrier, _semantic)
 

python/triton/experimental/tle/language/gpu/mthreads/__init__.py

@@ -18,6 +18,6 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
-from . import copy
+from . import buffer, common, copy, pipe, warp_specialize, wgmma
 
-__all__ = ["copy"]
+__all__ = ["buffer", "common", "copy", "pipe", "warp_specialize", "wgmma"]

python/triton/experimental/tle/language/gpu/mthreads/buffer.py

@@ -0,0 +1,124 @@
+# Copyright 2025-     FlagOS Contributors
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+from math import prod
+
+from triton._utils import TRITON_MAX_TENSOR_NUMEL, is_power_of_two
+
+from .. import types as gpu_types
+
+
+def is_backend_builder(builder) -> bool:
+    return hasattr(builder, "mark_musa_tle_auto_shared_layout")
+
+
+def needs_non_power_of_two_leading_dim(builder, shape) -> bool:
+    if not is_backend_builder(builder) or not shape:
+        return False
+    return isinstance(shape[0], int) and shape[0] > 0 and not is_power_of_two(shape[0])
+
+
+def validate_shape(shape) -> tuple[int, ...]:
+    shape = tuple(shape)
+    if len(shape) < 2:
+        raise ValueError("mthreads TLE non-power-of-two buffered tensor requires rank >= 2")
+    for index, dim in enumerate(shape):
+        if not isinstance(dim, int):
+            raise TypeError(f"Shape element {index} must have type `constexpr[int]`, got `constexpr[{type(dim)}]")
+        if dim <= 0:
+            raise ValueError(f"Shape element {index} must be positive")
+        if index > 0 and not is_power_of_two(dim):
+            raise ValueError(f"Shape element {index} must be a power of 2")
+    numel = prod(shape)
+    if numel > TRITON_MAX_TENSOR_NUMEL:
+        raise ValueError(f"numel ({numel}) exceeds triton maximum tensor numel ({TRITON_MAX_TENSOR_NUMEL})")
+    return shape
+
+
+class buffered_tensor_type(gpu_types.buffered_tensor_type):
+
+    def __init__(self, element_ty, shape, storage, layout=None, semantic=None, alloc_shape=None):
+        shape = validate_shape(shape)
+        self.element_ty = element_ty
+        self.shape = shape
+        self.numel = prod(shape)
+        self.name = f"<{self.shape}, {self.element_ty}>"
+        self.storage = storage
+        self.layout = layout
+        self.alloc_shape = list(shape if alloc_shape is None else alloc_shape)
+        assert semantic, "buffered_tensor array must be created with a builder"
+        self.semantic = semantic
+
+    def _unflatten_ir(self, handles, cursor):
+        value = buffered_tensor(
+            handles[cursor],
+            self.scalar,
+            self.shape,
+            self.storage,
+            self.layout,
+            self.semantic,
+            alloc_shape=self.alloc_shape,
+        )
+        if hasattr(self, "_tle_remote_shard_id"):
+            shard_id = getattr(self, "_tle_remote_shard_id")
+            scope = getattr(self, "_tle_remote_scope", None)
+            setattr(value, "_tle_remote_shard_id", shard_id)
+            setattr(value, "_tle_remote_scope", scope)
+            setattr(value.type, "_tle_remote_shard_id", shard_id)
+            setattr(value.type, "_tle_remote_scope", scope)
+        return value, cursor + 1
+
+    def with_element_ty(self, scalar_ty):
+        return buffered_tensor_type(
+            scalar_ty,
+            self.shape,
+            self.storage,
+            self.layout,
+            self.semantic,
+            alloc_shape=self.alloc_shape,
+        )
+
+
+class buffered_tensor(gpu_types.buffered_tensor):
+
+    def __init__(self, handle, element_ty, shape, storage, layout=None, semantic=None, alloc_shape=None):
+        self.handle = handle
+        self.shape = list(shape)
+        self.type = buffered_tensor_type(
+            element_ty,
+            shape,
+            storage,
+            layout,
+            semantic,
+            alloc_shape=alloc_shape,
+        )
+        self.dtype = element_ty
+
+
+def create_buffered_tensor(handle, element_ty, shape, storage, layout, semantic, alloc_shape=None):
+    return buffered_tensor(
+        handle,
+        element_ty,
+        shape,
+        storage,
+        layout,
+        semantic,
+        alloc_shape=alloc_shape,
+    )

python/triton/experimental/tle/language/gpu/mthreads/common.py

@@ -0,0 +1,38 @@
+# Copyright 2025-     FlagOS Contributors
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+import os
+
+try:
+    from triton._flagtree_backend import FLAGTREE_BACKEND
+except ModuleNotFoundError:
+    FLAGTREE_BACKEND = os.environ.get("FLAGTREE_BACKEND", "")
+
+
+def _has_mthreads_libtriton() -> bool:
+    try:
+        from triton._C import libtriton
+    except ImportError:
+        return False
+    return hasattr(libtriton, "mthreads")
+
+
+def enabled() -> bool:
+    return FLAGTREE_BACKEND == "mthreads" or _has_mthreads_libtriton()