triton-windows 3.3.0.post19__cp311-cp311-win_amd64.whl → 3.4.0.post20__cp311-cp311-win_amd64.whl

triton/experimental/gluon/language/_semantic.py

@@ -0,0 +1,287 @@
+from typing import Sequence, List, TypeVar, Tuple, Callable
+from triton.language.semantic import TritonSemantic
+from . import _core as ttgl
+from ._layouts import SliceLayout
+from triton._C.libtriton.gluon_ir import GluonOpBuilder
+from triton.compiler.code_generator import flatten_values_to_ir, unflatten_ir_values
+
+TensorTy = TypeVar("TensorTy")
+
+
+def _check(cond: bool, msg_fn: Callable[[], str], category=ValueError):
+    if not cond:
+        raise category(msg_fn())
+
+
+class GluonSemantic(TritonSemantic[TensorTy]):
+    tensor = ttgl.tensor
+    lang = ttgl
+
+    builder: GluonOpBuilder
+
+    def __init__(self, builder: GluonOpBuilder):
+        self.builder = builder
+
+    def _wrap_tensor_infer_layout(self, tensor):
+        ty = ttgl.distributed_type(tensor.type.scalar, tensor.shape,
+                                   self.builder.get_gluon_layout_from_tensor(tensor.handle))
+        return self.tensor(tensor.handle, ty)
+
+    def _broadcast_shapes(self, lhs_shape: List[int], rhs_shape: List[int]):
+        if len(lhs_shape) != len(rhs_shape):
+            raise ValueError(f"Cannot broadcast, rank mismatch: {lhs_shape}, {rhs_shape}")
+
+        ret_shape = []
+        for i, left in enumerate(lhs_shape):
+            right = rhs_shape[i]
+            if left == 1:
+                ret_shape.append(right)
+            elif (right == 1) or (right == left):
+                ret_shape.append(left)
+            else:
+                raise ValueError("Cannot make_shape_compatible: incompatible dimensions "
+                                 "at index " + str(i) + ": " + str(left) + " and " + str(right))
+        return ret_shape
+
+    def expand_dims(self, input: TensorTy, axis: int) -> TensorTy:
+        dst_shape = [ttgl._unwrap_if_constexpr(x) for x in input.shape]
+        dst_shape.insert(axis, 1)
+
+        if axis < 0:
+            axis += len(input.shape)
+
+        _check(isinstance(input.type, ttgl.distributed_type),
+               lambda: f"expected expand_dims input to be a distributed_type but got: {input.type!r}")
+        layout = input.type.layout
+        _check(isinstance(layout, SliceLayout),
+               lambda: f"expected expand_dims input to have a SliceLayout, but got: {layout}")
+        _check(layout.dim == axis,
+               lambda: f"expected expand_dims input layout to be sliced in axis {axis} but got {layout.dim}")
+
+        ret_ty = ttgl.distributed_type(input.type.scalar, dst_shape, layout.parent)
+        handle = self.builder.create_expand_dims(input.handle, axis, ret_ty.to_ir(self.builder))
+        return self.tensor(handle, ret_ty)
+
+    def join(self, a: TensorTy, b: TensorTy) -> TensorTy:
+        a, b = self.broadcast_impl_value(a, b)
+        _check(a.shape != [], "Cannot join scalars in gluon")
+        value = super().join(a, b)
+        return self._wrap_tensor_infer_layout(value)
+
+    def split(self, a: TensorTy) -> Tuple[TensorTy, TensorTy]:
+        lhs, rhs = super().split(a)
+        return self._wrap_tensor_infer_layout(lhs), self._wrap_tensor_infer_layout(rhs)
+
+    def permute(self, input: TensorTy, dims: Tuple[int]) -> TensorTy:
+        value = super().permute(input, dims)
+        return self._wrap_tensor_infer_layout(value)
+
+    def broadcast_impl_shape(self, input: TensorTy, shape: Tuple[int]) -> TensorTy:
+        _check(isinstance(input.type, ttgl.distributed_type),
+               lambda: f"expected expand_dims input to be a distributed_type but got: {input.type!r}")
+        src_shape = input.type.get_block_shapes()
+        _check(len(src_shape) == len(shape), lambda: f"Cannot broadcast, rank mismatch: {src_shape}, {shape}")
+        if shape == src_shape:
+            return input
+        for i, item in enumerate(src_shape):
+            if shape[i] != item and item != 1:
+                raise ValueError(f"Cannot broadcast, the expanded size of the tensor ({shape[i]})"
+                                 f" must match the existing size ({item}) at non-singleton dimension"
+                                 f" {i}: {src_shape}, {shape}")
+        ret_ty = ttgl.distributed_type(input.type.scalar, shape, input.type.layout)
+        handle = self.builder.create_broadcast(input.handle, ret_ty.to_ir(self.builder))
+        return self.tensor(handle, ret_ty)
+
+    def broadcast_impl_value(self, lhs: TensorTy, rhs: TensorTy) -> TensorTy:
+        lhs_ty = lhs.type
+        rhs_ty = rhs.type
+
+        if not lhs_ty.is_block() or not rhs_ty.is_block():
+            return super().broadcast_impl_value(lhs, rhs)
+
+        _check(isinstance(lhs_ty, ttgl.distributed_type),
+               lambda: f"expected broadcast left input to be a distributed_type but got: {lhs_ty!r}")
+        _check(isinstance(rhs_ty, ttgl.distributed_type),
+               lambda: f"expected broadcast right input to be a distributed_type but got: {rhs_ty!r}")
+
+        lhs_shape = lhs_ty.get_block_shapes()
+        rhs_shape = rhs_ty.get_block_shapes()
+        ret_shape = self._broadcast_shapes(lhs_shape, rhs_shape)
+        if lhs_ty.layout != rhs_ty.layout:
+            raise ValueError(f"Layout mismatch in broadcast: {lhs_ty.layout} vs {rhs_ty.layout}")
+
+        lhs = self.broadcast_impl_shape(lhs, ret_shape)
+        rhs = self.broadcast_impl_shape(rhs, ret_shape)
+        return lhs, rhs
+
+    def arange(self, start, end, layout):
+        shape = [end - start]
+        ret_ty = ttgl.distributed_type(ttgl.int32, shape, layout)
+        return super().arange(start, end, ret_ty=ret_ty)
+
+    def reshape(self, input: TensorTy, dst_shape: List[int], can_reorder: bool):
+        _check(not can_reorder, "can_reorder is not supported in gluon")
+        value = super().reshape(input, dst_shape, can_reorder)
+        return self._wrap_tensor_infer_layout(value)
+
+    def splat(self, value, shape, layout):
+        ret_ty = ttgl.distributed_type(value.dtype, shape, layout)
+        handle = self.builder.create_splat(ret_ty.to_ir(self.builder), value.handle)
+        return ttgl.tensor(handle, ret_ty)
+
+    def full(self, shape, value, dtype, layout):
+        scalar = self.make_scalar(value, dtype)
+        return self.splat(scalar, shape, layout)
+
+    def convert_layout(self, value, layout):
+        ty = value.type
+        _check(isinstance(ty, ttgl.distributed_type),
+               lambda: f"expected convert_layout input to be a distributed_type but got: {ty!r}")
+        ret_ty = ttgl.distributed_type(ty.element_ty, ty.shape, layout)
+        handle = self.builder.create_convert_layout(ret_ty.to_ir(self.builder), value.handle)
+        return ttgl.tensor(handle, ret_ty)
+
+    def allocate_shared(self, element_ty, shape, layout, value):
+        ty = ttgl.shared_memory_descriptor_type(element_ty, shape, layout, shape)
+        if value is not None:
+            handle = self.builder.create_local_alloc(ty.to_ir(self.builder), value.handle)
+        else:
+            handle = self.builder.create_local_alloc(ty.to_ir(self.builder))
+        return ttgl.shared_memory_descriptor(handle, element_ty, shape, layout, shape)
+
+    def shared_load(self, mem_desc, layout):
+        ret_ty = ttgl.distributed_type(mem_desc.dtype, mem_desc.shape, layout)
+        handle = self.builder.create_local_load(ret_ty.to_ir(self.builder), mem_desc.handle)
+        return ttgl.tensor(handle, ret_ty)
+
+    def shared_store(self, mem_desc, value):
+        self.builder.create_local_store(mem_desc.handle, value.handle)
+
+    def shared_dealloc(self, mem_desc):
+        self.builder.create_local_dealloc(mem_desc.handle)
+
+    def _memdesc_subview(self, mem_desc, offsets, shape):
+        layout = mem_desc.layout
+        ty = ttgl.shared_memory_descriptor_type(mem_desc.dtype, shape, layout, mem_desc.type.alloc_shape)
+        builder = self.builder
+        handle = builder.create_memdesc_subview(ty.to_ir(builder), mem_desc.handle, offsets)
+        return ttgl.shared_memory_descriptor(handle, **ty.__dict__)
+
+    def memdesc_slice(self, mem_desc, start, length, dim):
+        offsets = [self.builder.get_int32(0)] * mem_desc.rank
+        offsets[dim] = self.to_tensor(start).handle
+        shape = list(mem_desc.shape)
+        shape[dim] = length
+        return self._memdesc_subview(mem_desc, offsets, shape)
+
+    def memdesc_index(self, mem_desc, index):
+        shape = mem_desc.shape[1:]
+        offsets = [self.builder.get_int32(0)] * mem_desc.rank
+        offsets[0] = self.to_tensor(index).handle
+        return self._memdesc_subview(mem_desc, offsets, shape)
+
+    def memdesc_trans(self, mem_desc, order):
+        assert len(order) == len(
+            mem_desc.shape), f"source rank ({mem_desc.rank}) and order length ({len(order)}) must match"
+
+        shape = [mem_desc.shape[i] for i in order]
+        alloc_shape = mem_desc.type.alloc_shape
+        new_alloc_shape = alloc_shape[:len(alloc_shape) - mem_desc.rank]
+        new_alloc_shape += [alloc_shape[len(alloc_shape) - mem_desc.rank:][i] for i in order]
+
+        handle = self.builder.create_memdesc_trans(mem_desc.handle, order)
+        layout = self.builder.get_gluon_layout_from_memdesc(handle)
+        return ttgl.shared_memory_descriptor(handle, element_ty=mem_desc.dtype, shape=shape,
+                                             alloc_shape=new_alloc_shape, layout=layout)
+
+    def memdesc_reshape(self, mem_desc, shape, layout):
+        ty = ttgl.shared_memory_descriptor_type(mem_desc.dtype, shape, layout, mem_desc.type.alloc_shape)
+        handle = self.builder.create_memdesc_reshape(ty.to_ir(self.builder), mem_desc.handle)
+        return ttgl.shared_memory_descriptor(handle, **ty.__dict__)
+
+    def memdesc_reinterpret(self, mem_desc, dtype, shape, layout):
+        ty = ttgl.shared_memory_descriptor_type(dtype, shape, layout, shape)
+        handle = self.builder.create_memdesc_reinterpret(ty.to_ir(self.builder), mem_desc.handle)
+        return ttgl.shared_memory_descriptor(handle, **ty.__dict__)
+
+    def wrap_tensor(self, x, scalar_ty, ret_shape, layout):
+        if ret_shape:
+            res_ty = ttgl.distributed_type(scalar_ty, ret_shape, layout)
+        else:
+            res_ty = scalar_ty
+        return self.tensor(x, res_ty)
+
+    @staticmethod
+    def _check_same_layout(xs):
+        for x in xs:
+            _check(isinstance(x.type, ttgl.distributed_type), lambda: f"expected distributed_type but got: {x.type!r}")
+        layouts = [x.type.layout for x in xs]
+        l0 = layouts[0]
+        _check(all(l == l0 for l in layouts[1:]),
+               lambda: f"Expected inputs to have matching layouts, but got: {layouts}")
+
+    def reduction(self, inputs: Sequence[TensorTy], axis: int, region_builder_fn) -> Tuple[TensorTy, ...]:
+        _check(axis is not None, lambda: "All-reduce is not yet implemented in gluon")
+        # get result shape
+        shape = inputs[0].type.shape
+        rank = len(shape)
+        _check(0 <= axis < rank, lambda: f"expected reduction axis to be in the range [0, {rank}) but got {axis}")
+        self._check_same_layout(inputs)
+        ret_shape = [s for i, s in enumerate(shape) if i != axis]
+        ret_layout = SliceLayout(axis, inputs[0].type.layout)
+        assert all(t.type.shape == shape for t in inputs), "all reduction inputs must have the same shape"
+
+        reduce_op = self.builder.create_reduce([t.handle for t in inputs], axis)
+        region_builder_fn(reduce_op)
+        assert reduce_op.verify()
+
+        return tuple(
+            self.wrap_tensor(reduce_op.get_result(i), inputs[i].type.scalar, ret_shape, ret_layout)
+            for i in range(len(inputs)))
+
+    def warp_specialize(self, args, default_partition, worker_partitions, worker_num_warps: Sequence[int],
+                        worker_num_regs: Sequence[int], generator):
+        num_partitions = len(worker_partitions)
+        assert num_partitions == len(
+            worker_num_warps
+        ), f"warp specialize got {num_partitions} partitions but {len(worker_num_warps)} warp counts"
+        assert num_partitions == len(
+            worker_num_regs
+        ), f"warp specialize got {num_partitions} partitions but {len(worker_num_regs)} register counts"
+
+        builder = self.builder
+        insert_pt = builder.get_insertion_point()
+
+        # Emit the default partition to get the result types.
+        default_block = builder.new_block()
+        builder.set_insertion_point_to_start(default_block)
+        default_results = generator.call_JitFunction(default_partition, args, kwargs={})
+        mlir_results = []
+        if default_results is not None:
+            mlir_results = flatten_values_to_ir(default_results)
+        builder.create_warp_yield(mlir_results)
+        result_types = [r.get_type() for r in mlir_results]
+
+        # Create the warp specialize op.
+        builder.restore_insertion_point(insert_pt)
+        mlir_args = flatten_values_to_ir(args)
+        ws_op = builder.create_warp_specialize(result_types, mlir_args, worker_num_warps)
+        ws_op.get_default_region().push_back(default_block)
+        ws_op.set_requested_registers(worker_num_regs)
+
+        # Emit the partition regions.
+        builder.create_block_with_parent(ws_op.get_partition_op_holder(), [])
+        partitions_op = builder.create_warp_specialize_partitions(num_partitions)
+        arg_types = [arg.get_type() for arg in mlir_args]
+        for i in range(num_partitions):
+            block = builder.create_block_with_parent(partitions_op.get_region(i), arg_types)
+            block_args = [block.get_argument(j) for j in range(len(mlir_args))]
+            block_args = unflatten_ir_values(block_args, [arg.type for arg in args])
+            generator.call_JitFunction(worker_partitions[i], block_args, kwargs={})
+            builder.create_warp_return()
+
+        builder.set_insertion_point_after(ws_op.get_operation())
+        mlir_results = [ws_op.get_result(i) for i in range(len(result_types))]
+        if default_results is None:
+            return
+        return tuple(unflatten_ir_values(mlir_results, [r.type for r in default_results]))

triton/experimental/gluon/language/_standard.py

@@ -0,0 +1,47 @@
+# flake8: noqa
+import triton
+import triton.language.standard as tl_standard
+from .._runtime import jit
+from triton import knobs
+from . import _core as ttgl
+
+_IMPORT_FROM_TRITON = [
+    "sum",
+    "max",
+    "min",
+    "reduce_or",
+    "xor_sum",
+]
+
+__all__ = [
+    "full_like",
+    "zeros",
+    "zeros_like",
+    *_IMPORT_FROM_TRITON,
+]
+
+for name in _IMPORT_FROM_TRITON:
+    # Convert JITFunction -> GluonJitFunction
+    fn = getattr(tl_standard, name)
+    assert knobs.runtime.interpret or isinstance(fn, triton.runtime.JITFunction)
+    globals()[name] = jit(fn.fn)
+
+
+@jit
+def zeros(shape, dtype, layout):
+    return ttgl.full(shape, 0, dtype, layout)
+
+
+@jit
+def full_like(input, value, shape=None, dtype=None, layout=None):
+    return ttgl.full(
+        input.shape if shape is None else shape,
+        value,
+        input.dtype if dtype is None else dtype,
+        input.type.layout if layout is None else layout,
+    )
+
+
+@jit
+def zeros_like(input, shape=None, dtype=None, layout=None):
+    return full_like(input, 0, shape=shape, dtype=dtype, layout=layout)

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

@@ -0,0 +1,4 @@
+from . import blackwell
+from . import hopper
+
+__all__ = ["blackwell", "hopper"]

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

@@ -0,0 +1,202 @@
+from __future__ import annotations
+from typing import Optional, Tuple, List, TYPE_CHECKING
+
+from dataclasses import dataclass
+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._semantic import _check
+
+from . import tma
+from ..hopper import mbarrier, fence_async_shared
+
+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",
+    "fence_async_shared",
+    "mbarrier",
+    "tensor_memory_descriptor",
+    "TensorMemoryLayout",
+    "tma",
+]
+
+
+@dataclass(frozen=True, eq=True)
+class TensorMemoryLayout:
+    block: Tuple[int, int]
+    unpacked: bool
+    cta_split_num: Optional[Tuple[int, int]] = None
+
+    def __post_init__(self):
+        assert len(self.block) == 2
+        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_layout(
+            self.block,
+            self.unpacked,
+            cta_split_num,
+        )
+
+    def mangle(self) -> str:
+        block_str = f"{self.block[0]}x{self.block[1]}"
+        unpacked_str = "U" if self.unpacked else "P"
+        cta_split_str = f"CS{self.cta_split_num[0]}x{self.cta_split_num[1]}" if self.cta_split_num else ""
+        return f"TL{block_str}{unpacked_str}{cta_split_str}TL"
+
+
+class tensor_memory_descriptor_type(base_type):
+
+    def __init__(self, element_ty, shape, layout, alloc_shape):
+        self.element_ty = element_ty
+        self.shape = shape
+        self.layout = layout
+        self.alloc_shape = alloc_shape
+        assert isinstance(layout, TensorMemoryLayout)
+
+    def to_ir(self, builder: GluonOpBuilder) -> None:
+        return builder.get_tensor_mem_desc_ty(
+            self.element_ty.to_ir(builder),
+            self.shape,
+            self.layout._to_ir(builder),
+            self.alloc_shape,
+        )
+
+    def _unflatten_ir(self, handles: List[ir.Value], cursor: int) -> Tuple[tensor_memory_descriptor, int]:
+        value = tensor_memory_descriptor(handles[cursor], self.element_ty, self.shape, self.layout, self.alloc_shape)
+        return value, cursor + 1
+
+    def _flatten_ir_types(self, builder: GluonOpBuilder, out: List[ir.type]) -> None:
+        out.append(self.to_ir(builder))
+
+    def __str__(self) -> str:
+        return f"tensor_memory_descriptor<{self.element_ty}, {self.shape}, {self.layout}>"
+
+    def __eq__(self, other) -> bool:
+        return (type(self) is type(other) and self.shape == other.shape and self.layout == other.layout
+                and self.alloc_shape == other.alloc_shape)
+
+    def __neq__(self, other) -> bool:
+        return not (self == other)
+
+    def mangle(self) -> str:
+        shape_str = "_".join([str(s) for s in self.shape])
+        return f"MD{self.element_ty.mangle()}S{shape_str}SL{self.layout.mangle()}LAS{self.alloc_shape}ASMD"
+
+
+class tensor_memory_descriptor(base_value):
+
+    def __init__(self, handle, element_ty, shape, layout, alloc_shape):
+        self.handle = handle
+        self.type = tensor_memory_descriptor_type(element_ty, shape, layout, alloc_shape)
+
+    def _flatten_ir(self, handles: List[ir.value]) -> None:
+        handles.append(self.handle)
+
+    @property
+    def dtype(self):
+        return self.type.element_ty
+
+    @property
+    def shape(self):
+        return self.type.shape
+
+    @property
+    def rank(self):
+        return len(self.shape)
+
+    @property
+    def layout(self):
+        return self.type.layout
+
+    def __str__(self) -> str:
+        return str(self.type)
+
+    @builtin
+    def load(self, layout, _semantic: GluonSemantic) -> ttgl.tensor:
+        layout = _unwrap_if_constexpr(layout)
+        ret_ty = ttgl.distributed_type(self.dtype, self.shape, layout)
+        builder = _semantic.builder
+        handle = builder.create_tmem_load(ret_ty.to_ir(builder), self.handle)
+        return ttgl.tensor(handle, ret_ty)
+
+    @builtin
+    def store(self, value, pred=True, _semantic: GluonSemantic = None) -> None:
+        pred = _unwrap_if_constexpr(pred)
+        pred = _semantic.to_tensor(pred)
+        _semantic.builder.create_tmem_store(self.handle, value.handle, pred.handle)
+
+    @builtin
+    def slice(self, start, length, _semantic: GluonSemantic) -> None:
+        start = _unwrap_if_constexpr(start)
+        length = _unwrap_if_constexpr(length)
+        _check(isinstance(start, int), lambda: "start must be a constant int")
+        _check(isinstance(length, int), lambda: "length must be a constant int")
+        shape = self.shape[:-1] + [length]
+        layout = self.type.layout
+        layout = TensorMemoryLayout((layout.block[0], min(layout.block[1], length)), layout.unpacked,
+                                    layout.cta_split_num)
+        ret = tensor_memory_descriptor(None, self.dtype, shape, layout, self.type.alloc_shape)
+        builder = _semantic.builder
+        ret.handle = builder.create_tmem_subslice(ret.type.to_ir(builder), self.handle, start)
+        return ret
+
+    @builtin
+    def index(self, index, _semantic: GluonSemantic = None) -> tensor_memory_descriptor:
+        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)
+        return ret
+
+    @builtin
+    def _reinterpret(self, dtype, shape, layout, _semantic: GluonSemantic = None) -> tensor_memory_descriptor:
+        dtype = _unwrap_if_constexpr(dtype)
+        shape = [_unwrap_if_constexpr(s) for s in shape]
+        layout = _unwrap_if_constexpr(layout)
+
+        ty = tensor_memory_descriptor_type(dtype, shape, layout, shape)
+        handle = _semantic.builder.create_memdesc_reinterpret(ty.to_ir(_semantic.builder), self.handle)
+        return tensor_memory_descriptor(handle, **ty.__dict__)
+
+
+@builtin
+def allocate_tensor_memory(element_ty, shape, layout, value=None, _semantic=None):
+    element_ty = _unwrap_if_constexpr(element_ty)
+    shape = _unwrap_if_constexpr(shape)
+    layout = _unwrap_if_constexpr(layout)
+    value = value.handle if value is not None else None
+
+    ty = tensor_memory_descriptor_type(element_ty, shape, layout, shape)
+    builder = _semantic.builder
+    handle = builder.create_tmem_alloc(ty.to_ir(builder), value)
+    return tensor_memory_descriptor(handle, element_ty, shape, layout, shape)
+
+
+@builtin
+def tcgen05_mma(a, b, acc, *, use_acc=True, pred=True, mbarriers=None, mbarrier_preds=None, _semantic=None):
+    use_acc = _semantic.to_tensor(use_acc)
+    pred = _semantic.to_tensor(pred)
+
+    if mbarriers is None:
+        assert mbarrier_preds is None
+        mbarriers = []
+        mbarrier_preds = []
+    else:
+        mbarriers = [bar.handle for bar in mbarriers]
+        if mbarrier_preds is None:
+            true = _semantic.to_tensor(True)
+            mbarrier_preds = [true] * 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)

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

@@ -0,0 +1,32 @@
+from triton.experimental.gluon.language._core import builtin
+from triton.experimental.gluon.language.nvidia.hopper.tma import (
+    async_copy_global_to_shared,
+    async_copy_shared_to_global,
+    store_wait,
+    tensor_descriptor,
+    tensor_descriptor_type,
+)
+
+__all__ = [
+    "async_gather",
+    "async_scatter",
+    "async_copy_global_to_shared",
+    "async_copy_shared_to_global",
+    "store_wait",
+    "tensor_descriptor",
+    "tensor_descriptor_type",
+]
+
+
+@builtin
+def async_gather(tensor_desc, x_offsets, y_offset, barrier, result, pred=True, _semantic=None):
+    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,
+                                              result.handle, pred.handle)
+
+
+@builtin
+def async_scatter(tensor_desc, x_offsets, y_offset, src, _semantic=None):
+    y_offset = _semantic.to_tensor(y_offset)
+    _semantic.builder.create_async_tma_scatter(tensor_desc.handle, x_offsets.handle, y_offset.handle, src.handle)

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

@@ -0,0 +1,11 @@
+from . import mbarrier
+from . import tma
+from ... import _core
+
+__all__ = ["fence_async_shared", "mbarrier", "tma"]
+
+
+@_core.builtin
+def fence_async_shared(cluster=False, _semantic=None):
+    cluster = _core._unwrap_if_constexpr(cluster)
+    _semantic.builder.create_fence_async_shared(cluster)

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

@@ -0,0 +1,51 @@
+from triton.experimental.gluon.language._layouts import SwizzledSharedLayout
+from triton.experimental.gluon.language._core import builtin, _unwrap_if_constexpr
+
+__all__ = ["MBarrierLayout", "init", "invalidate", "expect", "wait", "arrive"]
+
+
+class MBarrierLayout(SwizzledSharedLayout):
+
+    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):
+    count = _unwrap_if_constexpr(count)
+    _semantic.builder.create_mbarrier_init(mbarrier.handle, count)
+
+
+@builtin
+def invalidate(mbarrier, _semantic=None):
+    _semantic.builder.create_mbarrier_inval(mbarrier.handle)
+
+
+@builtin
+def expect(mbarrier, bytes, pred=True, _semantic=None):
+    bytes = _unwrap_if_constexpr(bytes)
+    pred = _semantic.to_tensor(pred)
+    _semantic.builder.create_mbarrier_expect(mbarrier.handle, bytes, pred.handle)
+
+
+@builtin
+def wait(mbarrier, phase, pred=True, deps=(), _semantic=None):
+    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, count, pred=True, _semantic=None):
+    count = _unwrap_if_constexpr(count)
+    pred = _semantic.to_tensor(pred)
+    _semantic.builder.create_mbarrier_arrive(mbarrier.handle, count, pred.handle)