triton-windows 3.4.0.post20__cp311-cp311-win_amd64.whl → 3.5.0.post21__cp311-cp311-win_amd64.whl

triton/experimental/gluon/language/amd/cdna4/async_copy.py

@@ -0,0 +1,151 @@
+from ..._core import ir, builtin, _unwrap_if_constexpr
+from ..._semantic import _check
+from ..._layouts import BlockedLayout, SliceLayout
+from ..cdna3 import _verify_buffer_ops
+
+__all__ = [
+    "global_load_to_shared",
+    "buffer_load_to_shared",
+    "async_wait",
+    "load_shared_relaxed",
+]
+
+
+@builtin
+def global_load_to_shared(dest, ptr, mask=None, other=None, cache_modifier="", _semantic=None):
+    """
+    AMD global load to shared operation. This operation loads data directly
+    from global memory to shared memory without going through registers. It
+    happens asynchronously and requires a subsequent `async_wait` to ensure the
+    data is available in shared memory.
+    Compared to `buffer_load_to_shared`, it requires a tensor pointer which
+    supports 64-bit indexing range for each thread in a block, which gives more
+    flexibility, but at the cost of higher register pressure and no hardware
+    out-of-bound masking support. Prefer to use `buffer_load_to_shared` when
+    possible for better performance.
+
+    The underlying hardware instruction uses separate registers for global
+    memory address for each thread but the same register for local memory
+    address for the whole warp. Therefore, while using this operation
+    the following conditions must be met or lowering to LLVM will fail:
+
+    - For the `ptr` layout, size per thread * bits per element must be 128 or 32.
+      To get ideal performance, it is recommended to use 128 bits per element.
+    - Writes to `dest` must be coalesced.
+    - If `dest` is swizzled, it only can be swizzled within warp boundary.
+
+    Args:
+        dest (shared_memory_descriptor): Destination shared memory descriptor.
+        ptr (pointer tensor): Tensor of pointers to global memory to load from.
+        mask (tensor, optional): Mask tensor for predicated loads. Defaults to None.
+        other (tensor, optional): Tensor providing default values for masked elements. Defaults to None.
+        cache_modifier (str): Cache modifier specifier. Defaults to "".
+    """
+    _check(ptr.type.is_block(), lambda: "expected ptr to be a tensor")
+    _check(isinstance(ptr.type.layout, (BlockedLayout, SliceLayout)),
+           lambda: "expected ptr type layout to be BlockedLayout or SliceLayout")
+    _check(
+        dest.shape == ptr.shape, lambda:
+        f"expected dest shape to match pointer shape but got dest.shape = {dest.shape}, pointer.shape = {ptr.shape}")
+
+    mask = _unwrap_if_constexpr(mask)
+    if mask is not None:
+        ptr, mask = _semantic.broadcast_impl_value(ptr, mask)
+    other = _unwrap_if_constexpr(other)
+    if other is not None:
+        ptr, other = _semantic.broadcast_impl_value(ptr, other)
+
+    cache_modifier = _semantic._str_to_load_cache_modifier(cache_modifier)
+    mask_handle = mask.handle if mask is not None else ir.value()
+    other_handle = other.handle if other is not None else ir.value()
+    _semantic.builder.create_async_copy_global_to_local(dest.handle, ptr.handle, mask_handle, other_handle,
+                                                        cache_modifier, ir.EVICTION_POLICY.NORMAL, False)
+
+
+@builtin
+def buffer_load_to_shared(dest, ptr, offsets, mask=None, other=None, cache_modifier="", _semantic=None):
+    """
+    AMD buffer load to shared operation. Buffer load is similar to global load
+    but it accesses global memory via a scalar base pointer and a tensor of
+    32-bit offsets instead of a tensor of pointers. This operation loads data
+    directly from global memory to shared memory without going through
+    registers. It happens asynchronously and requires a subsequent `async_wait`
+    to ensure the data is available in shared memory.
+    Compared to `global_load_to_shared`, it has better performance and also
+    supports hardware out-of-bound masking. But it strictly requires a
+    32-bit offset instead of a 64-bit tensor pointer.
+
+    The underlying hardware instruction uses separate registers for global
+    memory address for each thread but the same register for local memory
+    address for the whole warp. Therefore, while using this operation
+    the following conditions must be met or lowering to LLVM will fail:
+
+    - For the `offsets` layout, size per thread * bits per element must be 128 or 32.
+      To get ideal performance, it is recommended to use 128 bits per element.
+    - Writes to `dest` must be coalesced.
+    - If `dest` is swizzled, it only can be swizzled within warp boundary.
+
+    Args:
+        dest (shared_memory_descriptor): Destination shared memory descriptor.
+        ptr (pointer to scalar): Global memory scalar base pointer to load from.
+        offsets (tensor): Offsets tensor for the load operation.
+        mask (tensor, optional): Mask tensor for predicated loads. Defaults to None.
+        other (tensor, optional): Tensor providing default values for masked elements. Defaults to None.
+        cache_modifier (str): Cache modifier specifier. Defaults to "".
+    """
+    _check(isinstance(offsets.type.layout, (BlockedLayout, SliceLayout)),
+           lambda: "expected offsets type layout to be BlockedLayout or SliceLayout")
+    _verify_buffer_ops(ptr, offsets, mask, other)
+
+    mask = _unwrap_if_constexpr(mask)
+    if mask is not None:
+        offsets, mask = _semantic.broadcast_impl_value(offsets, mask)
+    other = _unwrap_if_constexpr(other)
+    if other is not None:
+        offsets, other = _semantic.broadcast_impl_value(offsets, other)
+
+    mask = mask.handle if mask is not None else ir.value()
+    other = other.handle if other is not None else ir.value()
+    stride = ir.value()
+    cache_modifier = _semantic._str_to_load_cache_modifier(cache_modifier)
+
+    _semantic.builder.create_buffer_load_to_local(dest.handle, ptr.handle, offsets.handle, mask, other, stride,
+                                                  cache_modifier)
+
+
+@builtin
+def async_wait(num_outstanding=0, _semantic=None):
+    """
+    Wait for outstanding memory operations, this includes normal load like
+    `load` and `buffer_load`, as well as direct load to shared memory
+    like `global_load_to_shared` and `buffer_load_to_shared`.
+    It will block until the number of outstanding memory operations is less than
+    or equal to `num_outstanding`.
+
+    Args:
+        num_outstanding (int): The number of outstanding operations to wait for. Defaults to 0.
+    """
+    num_outstanding = _unwrap_if_constexpr(num_outstanding)
+    _semantic.builder.create_async_wait_group(num_outstanding)
+
+
+@builtin
+def load_shared_relaxed(smem, layout, _semantic=None):
+    """
+    Load a tensor from shared memory with extra hints for the underlying
+    compiler to avoid emitting unnecessary waits before loading from the target
+    shared memory.
+
+    Args:
+        smem (shared_memory_descriptor): Shared memory descriptor to load from.
+        layout (DistributedLayout): The destination layout of the tensor.
+
+    Returns:
+        tensor: A Gluon tensor containing the loaded data.
+    """
+    SYNCED_VIA_WAIT_ATTR_NAME = "ttg.amdgpu.syncedViaAsyncWait"
+
+    layout = _unwrap_if_constexpr(layout)
+    ret = _semantic.shared_load(smem, layout)
+    ret.handle.set_attr(SYNCED_VIA_WAIT_ATTR_NAME, _semantic.builder.get_bool_attr(True))
+    return ret

triton/experimental/gluon/language/extra/__init__.py

@@ -0,0 +1,3 @@
+from triton.language.extra import libdevice
+
+__all__ = ["libdevice"]

triton/experimental/gluon/language/nvidia/ampere/__init__.py

@@ -0,0 +1,3 @@
+from . import async_copy, mbarrier
+
+__all__ = ["async_copy", "mbarrier"]

triton/experimental/gluon/language/nvidia/ampere/async_copy.py

@@ -0,0 +1,74 @@
+from ..._semantic import _check
+from ..._core import _unwrap_if_constexpr, builtin
+from triton._C.libtriton import ir
+
+__all__ = [
+    "async_copy_global_to_shared",
+    "mbarrier_arrive",
+    "commit_group",
+    "wait_group",
+]
+
+
+@builtin
+def async_copy_global_to_shared(smem, pointer, mask=None, cache_modifier="", eviction_policy="", volatile=False,
+                                _semantic=None):
+    """
+    Asynchronously copy elements from global memory to shared memory.
+
+    Args:
+        smem (shared_memory_descriptor): Destination shared memory descriptor.
+        pointer (tensor): Source pointer tensor.
+        mask (tensor, optional): Mask tensor for predicated loads. Defaults to None.
+        cache_modifier (str): Cache modifier specifier. Defaults to "".
+        eviction_policy (str): Eviction policy specifier. Defaults to "".
+        volatile (bool): Whether the load is volatile. Defaults to False.
+    """
+    mask = _unwrap_if_constexpr(mask)
+    cache_modifier = _semantic._str_to_load_cache_modifier(cache_modifier)
+    eviction_policy = _semantic._str_to_eviction_policy(eviction_policy)
+    volatile = _unwrap_if_constexpr(volatile)
+    if mask is not None:
+        pointer, mask = _semantic.broadcast_impl_value(pointer, mask)
+    _check(
+        smem.shape == pointer.shape, lambda:
+        f"expected smem shape to match pointer shape but got smem.shape = {smem.shape}, pointer.shape = {pointer.shape}"
+    )
+    mask_handle = mask.handle if mask is not None else ir.value()
+    _semantic.builder.create_async_copy_global_to_local(smem.handle, pointer.handle, mask_handle, ir.value(),
+                                                        cache_modifier, eviction_policy, volatile)
+
+
+@builtin
+def mbarrier_arrive(mbarrier, increment_count=True, _semantic=None):
+    """
+    Arrive on the mbarrier once all outstanding async copies are complete.
+
+    Args:
+        mbarrier (shared_memory_descriptor): Barrier object to arrive on.
+        increment_count (bool): Whether to increment the arrival count. Defaults to True.
+    """
+    increment_count = _unwrap_if_constexpr(increment_count)
+    _semantic.builder.create_async_copy_mbarrier_arrive(mbarrier.handle, increment_count)
+
+
+@builtin
+def commit_group(_semantic=None):
+    """
+    Commit the current asynchronous copy group.
+
+    This finalizes a set of asynchronous copy operations.
+    """
+    _semantic.builder.create_async_commit_group()
+
+
+@builtin
+def wait_group(num_outstanding=0, _semantic=None):
+    """
+    Wait for outstanding asynchronous copy group operations.
+
+    Args:
+        num_outstanding (int): Wait until `num_outstanding` or less async copy groups in-flight. Defaults to 0.
+    """
+    num_outstanding = _unwrap_if_constexpr(num_outstanding)
+    _semantic.builder.create_async_wait_group(num_outstanding)

triton/experimental/gluon/language/nvidia/ampere/mbarrier.py

@@ -0,0 +1,80 @@
+from triton.experimental.gluon.language._layouts import SwizzledSharedLayout
+from triton.experimental.gluon.language._core import builtin, _unwrap_if_constexpr
+
+__all__ = ["arrive", "init", "invalidate", "MBarrierLayout", "wait"]
+
+
+class MBarrierLayout(SwizzledSharedLayout):
+    """
+    Layout for mbarrier synchronization in Ampere and later architectures.
+
+    Args:
+        ctas_per_cga (int): CTAs per CGA grouping. Defaults to 1.
+        cta_split_num (int): CTA split factor. Defaults to 1.
+    """
+
+    def __init__(self, ctas_per_cga: int = 1, cta_split_num: int = 1):
+        super().__init__(
+            vec=1,
+            per_phase=1,
+            max_phase=1,
+            order=[0],
+            ctas_per_cga=[ctas_per_cga],
+            cta_split_num=[cta_split_num],
+            cta_order=[0],
+        )
+
+
+@builtin
+def init(mbarrier, count, _semantic=None):
+    """
+    Initialize an mbarrier with a specified count.
+
+    Args:
+        mbarrier (shared_memory_descriptor): The barrier object to initialize.
+        count (int): The initial count for the barrier.
+    """
+    count = _unwrap_if_constexpr(count)
+    _semantic.builder.create_mbarrier_init(mbarrier.handle, count)
+
+
+@builtin
+def invalidate(mbarrier, _semantic=None):
+    """
+    Invalidate an mbarrier, resetting its state.
+
+    Args:
+        mbarrier (shared_memory_descriptor): The barrier object to invalidate.
+    """
+    _semantic.builder.create_mbarrier_inval(mbarrier.handle)
+
+
+@builtin
+def wait(mbarrier, phase, pred=True, deps=(), _semantic=None):
+    """
+    Wait until the mbarrier object completes its current phase.
+
+    Args:
+        mbarrier (shared_memory_descriptor): The barrier object to wait on.
+        phase (int): The phase index to wait for.
+        pred (bool): Predicate. Operation is skipped if predicate is False. Defaults to True.
+        deps (Sequence[shared_memory_descriptor]): Dependent allocations barrier is waiting on. Used to track liveness of dependent allocations. Defaults to ().
+    """
+    phase = _semantic.to_tensor(phase)
+    pred = _semantic.to_tensor(pred)
+    deps = [x.handle for x in deps]
+    _semantic.builder.create_mbarrier_wait(mbarrier.handle, phase.handle, pred.handle, deps)
+
+
+@builtin
+def arrive(mbarrier, *, pred=True, _semantic=None):
+    """
+    Arrive on an mbarrier, signaling that a thread has reached the barrier.
+
+    Args:
+        mbarrier (shared_memory_descriptor): The barrier object to arrive on.
+        pred (bool): Predicate. Operation is skipped if predicate is False. Defaults to True.
+    """
+    count = 1
+    pred = _semantic.to_tensor(pred)
+    _semantic.builder.create_mbarrier_arrive(mbarrier.handle, count, pred.handle)

triton/experimental/gluon/language/nvidia/blackwell/__init__.py

@@ -2,21 +2,26 @@ from __future__ import annotations
 from typing import Optional, Tuple, List, TYPE_CHECKING
 
 from dataclasses import dataclass
+from triton.runtime.jit import constexpr_function
 from triton.experimental.gluon.language import _core as ttgl
 from triton.experimental.gluon.language._core import builtin, base_type, base_value, _unwrap_if_constexpr
+from triton.experimental.gluon.language._layouts import BlockedLayout, _get_shape_per_cta
 from triton.experimental.gluon.language._semantic import _check
 
 from . import tma
-from ..hopper import mbarrier, fence_async_shared
+from ..hopper import fence_async_shared, mbarrier
+from ..ampere import async_copy
 
+from triton._C.libtriton import ir
 if TYPE_CHECKING:
     from triton._C.libtriton.gluon_ir import GluonOpBuilder
-    from triton._C.libtriton import gluon_ir as ir
     from ..._semantic import GluonSemantic
 
 __all__ = [
     "allocate_tensor_memory",
+    "async_copy",
     "fence_async_shared",
+    "get_tmem_32x32b_reg_layout",
     "mbarrier",
     "tensor_memory_descriptor",
     "TensorMemoryLayout",
@@ -26,6 +31,14 @@ __all__ = [
 
 @dataclass(frozen=True, eq=True)
 class TensorMemoryLayout:
+    """
+    Describes the layout for tensor memory in Blackwell architecture.
+
+    Args:
+        block (Tuple[int, int]): Tiling block dimensions (M/rows, N/cols).
+        unpacked (bool): For sub-32 bit elements, whether they are unpacked to 32 bits.
+        cta_split_num (Optional[Tuple[int, int]]): CTA split factors. Defaults to None.
+    """
     block: Tuple[int, int]
     unpacked: bool
     cta_split_num: Optional[Tuple[int, int]] = None
@@ -49,6 +62,74 @@ class TensorMemoryLayout:
         return f"TL{block_str}{unpacked_str}{cta_split_str}TL"
 
 
+@dataclass(frozen=True, eq=True)
+class TensorMemoryScalesLayout:
+    """
+    Describes the layout for tensor memory scales in Blackwell architecture.
+
+    Args:
+        cta_split_num (Optional[Tuple[int, int]]): CTA split factors. Defaults to None.
+    """
+    cta_split_num: Optional[Tuple[int, int]] = None
+
+    def __post_init__(self):
+        assert self.cta_split_num is None or len(self.cta_split_num) == 2
+
+    def _to_ir(self, builder):
+        cta_split_num = self.cta_split_num or [1, 1]
+        return builder.get_tensor_memory_scales_layout(cta_split_num, )
+
+    def mangle(self) -> str:
+        cta_split_str = f"CS{self.cta_split_num[0]}x{self.cta_split_num[1]}" if self.cta_split_num else ""
+        return f"TLS{cta_split_str}TLS"
+
+
+@constexpr_function
+def _cdiv(x, div):
+    return (x + div - 1) // div
+
+
+@constexpr_function
+def get_tmem_32x32b_reg_layout(M, N, shape, num_warps, ctas_per_cga=None, cta_split_num=None, cta_order=None):
+    """Returns a BlockedLayout compatible with load/store on tensor memory with the 32x32b instruction variant.
+    """
+    assert len(shape) == 2, "expected a 2D tensor"
+    assert num_warps in [4, 8], "expected 4 or 8 warps"
+
+    shape_per_cta = _get_shape_per_cta(shape, cta_split_num)
+    blocks_per_tile = [shape_per_cta[0] // M, shape_per_cta[1] // N]
+    num_blocks = blocks_per_tile[0] * blocks_per_tile[1]
+
+    num_warp_groups = num_warps // 4
+    if M == 64:
+        threads_per_warp = [16, 2]
+        if num_blocks == 1:
+            size_per_thread = [1, _cdiv(N, num_warp_groups * 2)]
+            warps_per_cta = [4, num_warp_groups]
+        else:
+            size_per_thread = [1, _cdiv(N, 2)]
+            warps_per_cta = [4 * min(blocks_per_tile[0], num_warp_groups)]
+            warps_per_cta.append(_cdiv(num_warp_groups, warps_per_cta[0] // 4))
+    else:
+        if shape[0] > 128:
+            size_per_thread = [1, N]
+            threads_per_warp = [32, 1]
+            warps_per_cta = [4 * num_warp_groups, 1]
+        else:
+            size_per_thread = [1, _cdiv(N, num_warp_groups)]
+            threads_per_warp = [32, 1]
+            warps_per_cta = [4, num_warp_groups]
+    return BlockedLayout(
+        size_per_thread=size_per_thread,
+        threads_per_warp=threads_per_warp,
+        warps_per_cta=warps_per_cta,
+        order=[0, 1],
+        ctas_per_cga=ctas_per_cga,
+        cta_split_num=cta_split_num,
+        cta_order=cta_order,
+    )
+
+
 class tensor_memory_descriptor_type(base_type):
 
     def __init__(self, element_ty, shape, layout, alloc_shape):
@@ -56,7 +137,7 @@ class tensor_memory_descriptor_type(base_type):
         self.shape = shape
         self.layout = layout
         self.alloc_shape = alloc_shape
-        assert isinstance(layout, TensorMemoryLayout)
+        assert isinstance(layout, TensorMemoryLayout) or isinstance(layout, TensorMemoryScalesLayout)
 
     def to_ir(self, builder: GluonOpBuilder) -> None:
         return builder.get_tensor_mem_desc_ty(
@@ -89,6 +170,9 @@ class tensor_memory_descriptor_type(base_type):
 
 
 class tensor_memory_descriptor(base_value):
+    """
+    Represents a tensor memory descriptor handle for Tensor Core Gen5 operations.
+    """
 
     def __init__(self, handle, element_ty, shape, layout, alloc_shape):
         self.handle = handle
@@ -118,6 +202,15 @@ class tensor_memory_descriptor(base_value):
 
     @builtin
     def load(self, layout, _semantic: GluonSemantic) -> ttgl.tensor:
+        """
+        Load a tensor from tensor memory.
+
+        Args:
+            layout (DistributedLayout): Destination layout of the tensor.
+
+        Returns:
+            tensor: A distributed tensor containing the loaded data.
+        """
         layout = _unwrap_if_constexpr(layout)
         ret_ty = ttgl.distributed_type(self.dtype, self.shape, layout)
         builder = _semantic.builder
@@ -126,12 +219,31 @@ class tensor_memory_descriptor(base_value):
 
     @builtin
     def store(self, value, pred=True, _semantic: GluonSemantic = None) -> None:
+        """
+        Store a tensor into tensor memory.
+
+        Args:
+            value (tensor): The tensor to store.
+            pred (bool): Scalar predicate. Operation is skipped if predicate is False. Defaults to True.
+        """
         pred = _unwrap_if_constexpr(pred)
         pred = _semantic.to_tensor(pred)
+        assert value.shape == self.shape, f"source shape {value.shape} does not match destination shape {self.shape}"
+        assert value.dtype == self.dtype, f"source dtype {value.dtype} does not match destination dtype {self.dtype}"
         _semantic.builder.create_tmem_store(self.handle, value.handle, pred.handle)
 
     @builtin
     def slice(self, start, length, _semantic: GluonSemantic) -> None:
+        """
+        Create a slice of the tensor memory descriptor along the last dimension.
+
+        Args:
+            start (int): The starting index for subslice.
+            length (int): The length of the subslice.
+
+        Returns:
+            tensor_memory_descriptor: Descriptor for the subslice.
+        """
         start = _unwrap_if_constexpr(start)
         length = _unwrap_if_constexpr(length)
         _check(isinstance(start, int), lambda: "start must be a constant int")
@@ -147,18 +259,36 @@ class tensor_memory_descriptor(base_value):
 
     @builtin
     def index(self, index, _semantic: GluonSemantic = None) -> tensor_memory_descriptor:
+        """
+        Create a subview of tensor memory by indexing the first dimension.
+
+        Args:
+            index (tensor): The index tensor for the subview.
+
+        Returns:
+            tensor_memory_descriptor: Descriptor for the indexed subview.
+        """
         index = _semantic.to_tensor(index)
         builder = _semantic.builder
-        offsets = [builder.get_int32(0)] * self.rank
-        offsets[0] = index.handle
         shape = self.shape[1:]
         layout = self.layout
         ret = tensor_memory_descriptor(None, self.dtype, shape, layout, self.type.alloc_shape)
-        ret.handle = builder.create_memdesc_subview(ret.type.to_ir(builder), self.handle, offsets)
+        ret.handle = builder.create_memdesc_index(ret.type.to_ir(builder), self.handle, index.handle)
         return ret
 
     @builtin
     def _reinterpret(self, dtype, shape, layout, _semantic: GluonSemantic = None) -> tensor_memory_descriptor:
+        """
+        Reinterpret tensor memory descriptor with a new dtype, shape, and layout.
+
+        Args:
+            dtype (dtype): The new data type.
+            shape (Sequence[int]): The new shape.
+            layout (TensorMemoryLayout): The new layout.
+
+        Returns:
+            tensor_memory_descriptor: Descriptor with updated type and layout.
+        """
         dtype = _unwrap_if_constexpr(dtype)
         shape = [_unwrap_if_constexpr(s) for s in shape]
         layout = _unwrap_if_constexpr(layout)
@@ -170,6 +300,18 @@ class tensor_memory_descriptor(base_value):
 
 @builtin
 def allocate_tensor_memory(element_ty, shape, layout, value=None, _semantic=None):
+    """
+    Allocate tensor memory.
+
+    Args:
+        element_ty (dtype): The element data type.
+        shape (Sequence[int]): The descriptor shape.
+        layout (TensorMemoryLayout): The layout of the tensor memory.
+        value (tensor, optional): Initial tensor to copy. Defaults to None.
+
+    Returns:
+        tensor_memory_descriptor: Descriptor for the allocated memory.
+    """
     element_ty = _unwrap_if_constexpr(element_ty)
     shape = _unwrap_if_constexpr(shape)
     layout = _unwrap_if_constexpr(layout)
@@ -181,8 +323,38 @@ def allocate_tensor_memory(element_ty, shape, layout, value=None, _semantic=None
     return tensor_memory_descriptor(handle, element_ty, shape, layout, shape)
 
 
+@builtin
+def tcgen05_copy(src, dst, _semantic=None):
+    """
+    Start an asynchronous copy from shared memory to tensor memory.
+
+    WARNING: The current semantics of the instruction are not well defined and
+    the API will change in the future. Use at your own risk.
+
+    Args:
+        src (shared_memory_descriptor): Shared memory to copy from.
+        dst (tensor_memory_descriptor): Tensor memory to copy to.
+    """
+    assert isinstance(src, ttgl.shared_memory_descriptor), "source must be a shared memory descriptor"
+    assert isinstance(dst, tensor_memory_descriptor), "destination must be a tensor memory descriptor"
+    _semantic.builder.create_tmem_copy(src.handle, dst.handle)
+
+
 @builtin
 def tcgen05_mma(a, b, acc, *, use_acc=True, pred=True, mbarriers=None, mbarrier_preds=None, _semantic=None):
+    """
+    Emit a 5th generation TensorCore MMA instruction.
+    acc = a * b + (acc if use_acc else 0)
+
+    Args:
+        a (shared_memory_descriptor): Left hand side operand in shared memory.
+        b (shared_memory_descriptor or tensor_memory_descriptor): Right hand side operand in shared or tensor memory.
+        acc (tensor_memory_descriptor): Accumulator value in tensor memory (mutated).
+        use_acc (bool): Whether to use the initial value of the accumulator. Defaults to True.
+        pred (bool): Scalar predicate. Operation is skipped if predicate is False. Defaults to True.
+        mbarriers (Sequence[shared_memory_descriptor], optional): Barriers to signal when the operation is complete. If None, mma is synchronous. Defaults to None.
+        mbarrier_preds (Sequence[bool], optional): Predicates for barriers. Defaults to None.
+    """
     use_acc = _semantic.to_tensor(use_acc)
     pred = _semantic.to_tensor(pred)
 
@@ -194,9 +366,22 @@ def tcgen05_mma(a, b, acc, *, use_acc=True, pred=True, mbarriers=None, mbarrier_
         mbarriers = [bar.handle for bar in mbarriers]
         if mbarrier_preds is None:
             true = _semantic.to_tensor(True)
-            mbarrier_preds = [true] * len(mbarriers)
+            mbarrier_preds = [true.handle] * len(mbarriers)
         else:
             mbarrier_preds = _semantic._convert_to_ir_values(mbarrier_preds, require_i64=False)
 
     _semantic.builder.create_tcgen05_mma(a.handle, b.handle, acc.handle, use_acc.handle, pred.handle, mbarriers,
                                          mbarrier_preds)
+
+
+@builtin
+def tcgen05_commit(barrier, _semantic=None):
+    """
+    This instruction causes the provided mbarrier to be arrived-on with a count
+    of 1 when all async tcgen05 MMA and copy instructions previously issued by
+    the thread are complete.
+
+    Args:
+        barrier (shared_memory_descriptor): The barrier to track completion of tcgen05 MMA and copy instructions.
+    """
+    _semantic.builder.create_tcgen05_commit(barrier.handle)

triton/experimental/gluon/language/nvidia/blackwell/tma.py

@@ -20,6 +20,17 @@ __all__ = [
 
 @builtin
 def async_gather(tensor_desc, x_offsets, y_offset, barrier, result, pred=True, _semantic=None):
+    """
+    Asynchronously gather elements from global memory to shared memory using TMA.
+
+    Args:
+        tensor_desc (tensor_descriptor): The tensor descriptor.
+        x_offsets (tensor): 1D tensor of X offsets.
+        y_offset (int): Scalar Y offset.
+        barrier (shared_memory_descriptor): Barrier that will be signaled when the operation is complete.
+        result (tensor_memory_descriptor): Result shared memory, must have NVMMASharedLayout.
+        pred (bool): Scalar predicate. Operation is skipped if predicate is False. Defaults to True.
+    """
     pred = _semantic.to_tensor(pred)
     y_offset = _semantic.to_tensor(y_offset)
     _semantic.builder.create_async_tma_gather(tensor_desc.handle, x_offsets.handle, y_offset.handle, barrier.handle,
@@ -28,5 +39,14 @@ def async_gather(tensor_desc, x_offsets, y_offset, barrier, result, pred=True, _
 
 @builtin
 def async_scatter(tensor_desc, x_offsets, y_offset, src, _semantic=None):
+    """
+    Asynchronously scatter elements from shared memory to global memory using TMA.
+
+    Args:
+        tensor_desc (tensor_descriptor): The tensor descriptor.
+        x_offsets (tensor): 1D tensor of X offsets.
+        y_offset (int): Scalar Y offset.
+        src (tensor_memory_descriptor): The source data, must be in NVMMASharedLayout.
+    """
     y_offset = _semantic.to_tensor(y_offset)
     _semantic.builder.create_async_tma_scatter(tensor_desc.handle, x_offsets.handle, y_offset.handle, src.handle)