triton-windows 3.3.1.post19__cp310-cp310-win_amd64.whl → 3.5.0.post21__cp310-cp310-win_amd64.whl

triton/experimental/gluon/language/_semantic.py

@@ -0,0 +1,380 @@
+from typing import Sequence, List, TypeVar, Tuple, Callable
+import math
+from triton.language.semantic import TritonSemantic
+from . import _core as ttgl
+from ._layouts import AutoLayout, DistributedLayout, 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 GluonCallerContext:
+
+    def __init__(self, num_warps: int):
+        self.num_warps = num_warps
+
+    def mangle(self):
+        return f"_NW{self.num_warps}"
+
+    def initialize_callee(self, fn, builder):
+        fn.set_attr("ttg.num-warps", builder.get_int32_attr(self.num_warps))
+
+
+class GluonSemantic(TritonSemantic[TensorTy]):
+    tensor = ttgl.tensor
+    lang = ttgl
+
+    builder: GluonOpBuilder
+
+    def __init__(self, builder: GluonOpBuilder):
+        self.builder = builder
+
+    def _wrap_handle_infer_layout(self, handle, scalar_ty, shape):
+        if shape == []:
+            ty = scalar_ty
+        else:
+            ty = ttgl.distributed_type(scalar_ty, shape, self.builder.get_gluon_layout_from_tensor(handle))
+        return self.tensor(handle, ty)
+
+    def _wrap_tensor_infer_layout(self, tensor):
+        return self._wrap_handle_infer_layout(tensor.handle, tensor.type.scalar, tensor.shape)
+
+    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, AutoLayout)),
+               lambda: f"expected expand_dims input to have a SliceLayout, but got: {layout}")
+        _check(
+            isinstance(layout, AutoLayout) or layout.dim == axis,
+            lambda: f"expected expand_dims input layout to be sliced in axis {axis} but got {layout.dim}")
+
+        handle = self.builder.create_expand_dims(input.handle, axis)
+        return self._wrap_handle_infer_layout(handle, input.type.scalar, dst_shape)
+
+    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)
+
+        is_lhs_auto = isinstance(lhs_ty.layout, AutoLayout)
+        is_rhs_auto = isinstance(rhs_ty.layout, AutoLayout)
+        if is_lhs_auto and not is_rhs_auto:
+            lhs = self.set_auto_layout(lhs, rhs_ty.layout)
+        elif is_rhs_auto and not is_lhs_auto:
+            rhs = self.set_auto_layout(rhs, lhs_ty.layout)
+        elif 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]
+        if layout is None:
+            layout = AutoLayout()
+        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)
+        if layout is None:
+            layout = AutoLayout()
+        return self.splat(scalar, shape, layout)
+
+    def convert_layout(self, value, layout, assert_trivial=False):
+        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)
+        ret_ty_ir = ret_ty.to_ir(self.builder)
+        if assert_trivial and not self.builder.is_convert_layout_trivial(ret_ty_ir, value.handle):
+            raise TypeError(f"layout conversion from {ty.layout} to {layout} is not trivial")
+        handle = self.builder.create_convert_layout(ret_ty_ir, 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):
+        assert value.shape == mem_desc.shape, f"source shape {value.shape} and destination shape {mem_desc.shape} must match"
+        assert value.dtype == mem_desc.dtype, f"source dtype {value.dtype} and destination dtype {mem_desc.dtype} must match"
+        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 set_auto_layout(self, value, layout):
+        src_ty = value.type
+        assert isinstance(layout,
+                          DistributedLayout), f"set_auto_layout must set to a distributed layout but got {layout}"
+        assert isinstance(src_ty.layout,
+                          AutoLayout), f"set_auto_layout input must have auto layout but got {value.type.layout}"
+        handle = self.builder.create_set_auto_layout(layout._to_ir(self.builder), value.handle)
+        res_ty = ttgl.distributed_type(src_ty.element_ty, src_ty.shape, layout)
+        return self.tensor(handle, res_ty)
+
+    def memdesc_slice(self, mem_desc, start, length, dim):
+        offsets = [0] * mem_desc.rank
+        offsets[dim] = start
+        shape = list(mem_desc.shape)
+        shape[dim] = length
+        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_subslice(ty.to_ir(builder), mem_desc.handle, offsets)
+        return ttgl.shared_memory_descriptor(handle, **ty.__dict__)
+
+    def memdesc_index(self, mem_desc, index):
+        shape = mem_desc.shape[1:]
+        index = self.to_tensor(index).handle
+        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_index(ty.to_ir(builder), mem_desc.handle, index)
+        return ttgl.shared_memory_descriptor(handle, **ty.__dict__)
+
+    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):
+        _check(
+            math.prod(shape) == math.prod(mem_desc.shape),
+            lambda: (f"memdesc_reshape total elements mismatch: "
+                     f"{mem_desc.shape} -> {shape}"),
+        )
+
+        handle = self.builder.create_memdesc_reshape(mem_desc.handle, shape)
+        layout = self.builder.get_gluon_layout_from_memdesc(handle)
+        alloc_shape = mem_desc.type.alloc_shape
+        prefix_len = len(alloc_shape) - mem_desc.rank
+        new_alloc_shape = alloc_shape[:prefix_len] + list(shape)
+
+        return ttgl.shared_memory_descriptor(
+            handle,
+            element_ty=mem_desc.dtype,
+            shape=shape,
+            alloc_shape=new_alloc_shape,
+            layout=layout,
+        )
+
+    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 associative_scan(self, inputs: Sequence[TensorTy], axis: int, region_builder_fn,
+                         reverse: bool) -> Tuple[TensorTy, ...]:
+        shape = inputs[0].type.shape
+        rank = len(shape)
+
+        assert -rank <= axis < rank, f"scan axis {axis} must be < inputs rank ({rank})"
+
+        if axis < 0:
+            axis += rank
+
+        for t in inputs:
+            assert t.type.shape == shape, "all scan inputs must have the same shape"
+
+        scan_op = self.builder.create_scan([t.handle for t in inputs], axis, reverse)
+        region_builder_fn(scan_op)
+        assert scan_op.verify()
+
+        return tuple(
+            self._wrap_handle_infer_layout(scan_op.get_result(i), inputs[i].type.scalar, shape)
+            for i in range(len(inputs)))
+
+    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]
+        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_handle_infer_layout(reduce_op.get_result(i), inputs[i].type.scalar, ret_shape)
+            for i in range(len(inputs)))
+
+    def histogram(self, input: TensorTy, num_bins: int, mask: TensorTy, layout) -> TensorTy:
+        _check(len(input.shape) == 1, lambda: "histogram only supports 1D input")
+        _check(input.dtype.is_int(), lambda: "histogram only supports integer input")
+        _check(layout is not None, lambda: "histogram requires a destination layout")
+        if mask is not None:
+            mask, input = self.broadcast_impl_value(mask, input)
+            _check(mask.type.scalar.is_bool(), lambda: "Mask must have boolean scalar type")
+            mask = mask.handle
+        layout_attr = layout._to_ir(self.builder)
+        handle = self.builder.create_histogram(input.handle, num_bins, mask, layout_attr)
+        return self.wrap_tensor(handle, ttgl.int32, [num_bins], layout)
+
+    def warp_specialize(self, default_args, default_partition, worker_args, 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, default_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(worker_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):
+            caller_context = GluonCallerContext(num_warps=worker_num_warps[i])
+            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 worker_args])
+            generator.call_JitFunction(worker_partitions[i], block_args, kwargs={}, caller_context=caller_context)
+            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,80 @@
+from typing import TypeVar
+from triton.runtime.jit import JITFunction
+import triton.language.standard as tl_standard
+from .._runtime import GluonJITFunction, jit
+from triton import knobs
+from . import _core as ttgl
+
+T = TypeVar("T")
+
+
+def _import_from_triton(fn: JITFunction[T]) -> GluonJITFunction[T]:
+    assert knobs.runtime.interpret or isinstance(fn, JITFunction)
+    # Wrap the function and preserve its original docstring
+    gluon_fn = jit(fn.fn)
+    gluon_fn.__doc__ = fn.__doc__
+    return gluon_fn
+
+
+cdiv = _import_from_triton(tl_standard.cdiv)
+sum = _import_from_triton(tl_standard.sum)
+max = _import_from_triton(tl_standard.max)
+min = _import_from_triton(tl_standard.min)
+reduce_or = _import_from_triton(tl_standard.reduce_or)
+xor_sum = _import_from_triton(tl_standard.xor_sum)
+
+
+@jit
+def zeros(shape, dtype, layout=None):
+    """
+    Create a tensor filled with zeros.
+
+    Args:
+        shape (Sequence[int]): The shape of the tensor.
+        dtype (dtype): The data type for the tensor.
+        layout (Optional[DistributedLayout]): The distributed layout of the tensor, defaults to AutoLayout().
+
+    Returns:
+        tensor: A tensor where every element is zero.
+    """
+    return ttgl.full(shape, 0, dtype, layout)
+
+
+@jit
+def full_like(input, value, shape=None, dtype=None, layout=None):
+    """
+    Create a tensor with the same properties as a given tensor, filled with a specified value.
+
+    Args:
+        input (tensor): Reference tensor to infer default shape, dtype, and layout.
+        value (int or float): The fill value.
+        shape (Sequence[int], optional): Target shape. Defaults to input.shape.
+        dtype (dtype, optional): Target data type. Defaults to input.dtype.
+        layout (DistributedLayout, optional): Target layout. Defaults to input.layout.
+
+    Returns:
+        tensor: A tensor where every element equals value.
+    """
+    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):
+    """
+    Create a tensor with the same properties as a given tensor, filled with zeros.
+
+    Args:
+        input (tensor): Reference tensor to infer default shape, dtype, and layout.
+        shape (Sequence[int], optional): Target shape. Defaults to input.shape.
+        dtype (dtype, optional): Target data type. Defaults to input.dtype.
+        layout (DistributedLayout, optional): Target layout. Defaults to input.layout.
+
+    Returns:
+        tensor: A tensor where every element is zero.
+    """
+    return full_like(input, 0, shape=shape, dtype=dtype, layout=layout)

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

@@ -0,0 +1,4 @@
+from ._layouts import AMDMFMALayout
+from . import cdna3, cdna4
+
+__all__ = ["AMDMFMALayout", "cdna3", "cdna4"]

triton/experimental/gluon/language/amd/_layouts.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import List, Optional
+from triton.language.core import _unwrap_if_constexpr
+
+from triton.experimental.gluon.language._layouts import _realize_cta_layout, DistributedLayout
+from triton.experimental.gluon import language as ttgl
+
+__all__ = [
+    "AMDMFMALayout",
+]
+
+
+@dataclass(frozen=True)
+class AMDMFMALayout(DistributedLayout):
+    """
+    Represents a layout for AMD MFMA (matrix core) operations.
+
+    Args:
+        version (int): Major and minor identifier for the MFMA instruction.
+        instr_shape: (M, N) dimension for the instrinsic shape.
+        transposed (bool): indicates the result tensor is transposed so that each thread holds consecutive elements in the same row instead of column, which is good for chained dot and global write.
+        warps_per_cta (List[int]): Number of warps per CTA.
+        elem_type Optional(ttgl.dtype): Supported types are int32, fp32 and fp64. Default is fp32.
+        tiles_per_warp Optional(List[int]): Number of tiles per WARP. For mfma layout, if missing, use the default where we have unit tile size on all dimensions.
+        ctas_per_cga (Optional[List[int]]): CTAs per CGA grouping.
+        cta_split_num (Optional[List[int]]): Split factors for CTAs.
+        cta_order (Optional[List[int]]): CTA ordering.
+    """
+    version: int
+    instr_shape: List[int]
+    transposed: bool
+    warps_per_cta: List[int]
+    elem_type: ttgl.dtype = ttgl.float32
+    tiles_per_warp: Optional[List[int]] = None
+    ctas_per_cga: Optional[List[int]] = None
+    cta_split_num: Optional[List[int]] = None
+    cta_order: Optional[List[int]] = None
+
+    def __post_init__(self):
+        super().__setattr__("version", _unwrap_if_constexpr(self.version))
+        super().__setattr__("instr_shape", _unwrap_if_constexpr(self.instr_shape))
+        super().__setattr__("transposed", _unwrap_if_constexpr(self.transposed))
+        super().__setattr__("warps_per_cta", _unwrap_if_constexpr(self.warps_per_cta))
+        super().__setattr__("tiles_per_warp", _unwrap_if_constexpr(self.tiles_per_warp))
+        super().__setattr__("elem_type", _unwrap_if_constexpr(self.elem_type))
+        super().__setattr__("ctas_per_cga", _unwrap_if_constexpr(self.ctas_per_cga))
+        super().__setattr__("cta_split_num", _unwrap_if_constexpr(self.cta_split_num))
+        super().__setattr__("cta_order", _unwrap_if_constexpr(self.cta_order))
+
+        if self.tiles_per_warp is None:
+            object.__setattr__(self, "tiles_per_warp", [1] * len(self.warps_per_cta))
+
+        self.verify()
+
+    def _to_ir(self, builder):
+        type = self.elem_type.to_ir(builder)
+        return builder.get_amd_mfma_layout(self.version, self.instr_shape, self.transposed, self.warps_per_cta, type,
+                                           self.tiles_per_warp, self.ctas_per_cga, self.cta_split_num, self.cta_order)
+
+    def mangle(self) -> str:
+
+        def stringify(x):
+            if x is None:
+                return ""
+            return "_".join(map(str, x))
+
+        return f"MFMA_{self.version}_{stringify(self.instr_shape)}_{self.transposed}_{stringify(self.warps_per_cta)}_{stringify(self.tiles_per_warp)}_{self.elem_type}_{stringify(self.ctas_per_cga)}_{stringify(self.cta_split_num)}_{stringify(self.cta_order)}_MFMA"
+
+    def verify(self):
+        assert self.version >= 1 and self.version <= 4, "version must be in the [1, 4] range"
+        valid_shapes = [[32, 32], [16, 16], [64, 4], [4, 64]]
+        assert self.instr_shape in valid_shapes, "invalid intrinsic shape; accepted shapes are " + str(valid_shapes)
+
+        assert self.elem_type.is_fp32() or self.elem_type.is_fp64() \
+          or self.elem_type.is_int32() , "element type must be float32, float64, or int32"
+
+        rank = len(self.warps_per_cta)
+        _realize_cta_layout(self, rank)
+        assert len(self.ctas_per_cga) == rank
+        assert len(self.cta_split_num) == rank
+        assert len(self.cta_order) == rank
+
+    def __hash__(self):
+        return hash((
+            self.version,
+            tuple(self.instr_shape),
+            self.transposed,
+            tuple(self.warps_per_cta),
+            self.elem_type,
+            tuple(self.tiles_per_warp) if self.tiles_per_warp else None,
+            tuple(self.ctas_per_cga) if self.ctas_per_cga else None,
+            tuple(self.cta_split_num) if self.cta_split_num else None,
+            tuple(self.cta_order) if self.cta_order else None,
+        ))

triton/experimental/gluon/language/amd/cdna3/__init__.py

@@ -0,0 +1,100 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING
+
+from triton import knobs
+from triton.experimental.gluon.language import _core as ttgl
+from triton._C.libtriton import ir
+from ..._core import builtin, _unwrap_if_constexpr
+
+if TYPE_CHECKING:
+    from ..._semantic import GluonSemantic
+
+__all__ = ["buffer_load", "buffer_store", "mfma"]
+
+
+def _verify_buffer_ops(ptr, offsets, mask=None, other=None):
+    assert ptr.type.is_ptr(), "ptr must be a scalar pointer type"
+
+    assert isinstance(offsets.type, ttgl.distributed_type), "expected offsets type to be a distributed_type"
+    assert offsets.dtype.is_int32() or offsets.dtype.is_uint32(), "offsets element type must be int32 or uint32"
+
+    element_type = ptr.type.scalar.element_ty
+
+    if other is not None:
+        assert mask is not None, "when other is not None, mask should not be None"
+        assert other.dtype == element_type, "other must have the same data type as ptr scalar type"
+
+
+@builtin
+def buffer_load(ptr, offsets, mask=None, other=None, cache=None, _semantic=None):
+    """
+    AMD buffer load from global memory via a scalar base pointer and a tensor of
+    offsets instead of a tensor of pointers. This operation will load data
+    directly into registers.
+
+    Args:
+        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 "".
+    """
+    _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)
+
+    other = other.handle if other is not None else ir.value()
+    mask = mask.handle if mask is not None else ir.value()
+    cache_modifier = _semantic._str_to_load_cache_modifier(cache) if cache is not None else ir.CACHE_MODIFIER.NONE
+
+    ret_ty = offsets.type.with_element_ty(ptr.type.scalar.element_ty)
+    builder = _semantic.builder
+    handle = builder.create_buffer_load(ret_ty.to_ir(builder), ptr.handle, offsets.handle, mask, other, cache_modifier)
+    return ttgl.tensor(handle, ret_ty)
+
+
+@builtin
+def buffer_store(stored_value, ptr, offsets, mask=None, cache=None, _semantic: GluonSemantic = None):
+    """
+    AMD buffer store a tensor directly to global memory via a scalar base pointer and a tensor of
+    offsets instead of a tensor of pointers.
+    Args:
+        stored_value (tensor to be stored): The tensor to be stored to global memory.
+        ptr (pointer to scalar): Global memory scalar base pointer to store to.
+        offsets (tensor): Offsets tensor for the store operation.
+        mask (tensor, optional): Mask tensor for predicated store. Defaults to None.
+        cache_modifier (str): Cache modifier specifier. Defaults to "".
+    """
+    _verify_buffer_ops(ptr, offsets, mask)
+
+    if mask is not None:
+        offsets, mask = _semantic.broadcast_impl_value(offsets, mask)
+
+    mask = mask.handle if mask is not None else ir.value()
+    cache_modifier = _semantic._str_to_store_cache_modifier(cache) if cache is not None else ir.CACHE_MODIFIER.NONE
+
+    _semantic.builder.create_buffer_store(stored_value.handle, ptr.handle, offsets.handle, mask, cache_modifier)
+
+
+@builtin
+def mfma(a, b, acc, _semantic: GluonSemantic = None):
+    """
+    Computes matrix-multiplication of a * b + acc using AMD native matrix core units.
+    Args:
+        a (tensor): The first operand of mfma.
+        b (tensor): The second operand of mfma.
+        acc (tensor): The accumulator tensor.
+    """
+    assert acc is not None, "acc is required"
+    ret_type = acc.type
+    acc = ttgl._unwrap_if_constexpr(acc)
+
+    handle = _semantic.dot(a, b, acc, input_precision=knobs.language.fp32_default, max_num_imprecise_acc=None,
+                           out_dtype=acc.dtype).handle
+    return ttgl.tensor(handle, ret_type)