from __future__ import annotations
from typing import Tuple, List, TYPE_CHECKING

from dataclasses import dataclass, field
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._semantic import _check, _compute_tmem_reg_layout

from . import tma
from . import clc
from ..hopper import async_store, cluster, fence_async_shared, mbarrier
from ..ampere import async_copy, mma_v2

from triton._C.libtriton import ir
import triton._C.libtriton.gluon_ir as gluon_ir
if TYPE_CHECKING:
    from triton._C.libtriton.gluon_ir import GluonOpBuilder
    from ..._semantic import GluonSemantic

__all__ = [
    "add2",
    "allocate_tensor_memory",
    "async_copy",
    "async_store",
    "clc",
    "cluster",
    "fence_async_shared",
    "fma2",
    "mbarrier",
    "max2",
    "min2",
    "mma_v2",
    "mul2",
    "sub2",
    "tensor_memory_descriptor",
    "tensor_memory_descriptor_type",
    "TensorMemoryLayout",
    "TensorMemoryScalesLayout",
    "tma",
    "_TensorMemoryLinearLayout",
]


def _packed_arith(operation, operands, dtype, semantic):
    """Build packed arithmetic with inferred result types and FP4 layouts."""
    operands = tuple(semantic.to_tensor(operand) for operand in operands)
    _check(any(isinstance(operand.type, ttgl.distributed_type) for operand in operands),
           lambda: "packed arithmetic requires at least one distributed tensor operand")
    floating = [
        operand for operand in operands
        if isinstance(operand.type, ttgl.distributed_type) and operand.dtype.is_floating()
    ]
    reference = max(floating, key=lambda operand: operand.numel.value, default=None)
    dtype = _unwrap_if_constexpr(dtype)

    if reference is None:
        _check(dtype is not None, lambda: "packed FP4 operands require an explicit result dtype")
        shape = list(operands[0].type.shape)
        shape[-1] *= 2
    else:
        shape = reference.type.shape
        if dtype is None:
            addend = operands[-1]
            dtype = addend.dtype if operation in ("add", "sub", "fma") and addend.dtype.is_fp8() else reference.dtype
            for operand in operands:
                if not operand.dtype.is_floating() or (dtype.is_fp8() and operand.dtype.is_fp8()):
                    continue
                dtype = semantic.computation_type_impl(dtype, False, operand.dtype, not operand.type.is_block(), False)

    _check(isinstance(dtype, ttgl.dtype), lambda: f"expected 'dtype' to be a dtype but got {dtype}")
    normalized = []
    for operand in operands:
        is_packed_fp4 = (isinstance(operand.type, ttgl.distributed_type) and operand.dtype.is_int()
                         and operand.dtype.primitive_bitwidth == 8 and operand.type.shape != shape)
        if not is_packed_fp4:
            _check(reference is not None, lambda: "packed FP4 operands require a floating-point tensor operand")
            if not operand.type.is_block():
                operand = semantic.cast(operand, dtype)
            _, operand = semantic.broadcast_impl_value(reference, operand)
        normalized.append(operand.handle)

    builder = semantic.builder
    result_type = dtype if reference is None else reference.type.with_element_ty(dtype)
    result_ir_type = result_type.to_ir(builder)
    handle = builder.create_packed_arith(result_ir_type, operation, normalized)
    return semantic._wrap_handle_infer_layout(handle, dtype, shape)


@builtin
def add2(lhs, rhs, dtype=None, _semantic=None):
    """Add two tensors using a native two-lane packed instruction."""
    return _packed_arith("add", (lhs, rhs), dtype, _semantic)


@builtin
def sub2(lhs, rhs, dtype=None, _semantic=None):
    """Subtract two tensors using a native two-lane packed instruction."""
    return _packed_arith("sub", (lhs, rhs), dtype, _semantic)


@builtin
def mul2(lhs, rhs, dtype=None, _semantic=None):
    """Multiply two tensors using a native two-lane packed instruction."""
    return _packed_arith("mul", (lhs, rhs), dtype, _semantic)


@builtin
def fma2(lhs, rhs, acc, dtype=None, _semantic=None):
    """Perform a native two-lane packed fused multiply-add."""
    return _packed_arith("fma", (lhs, rhs, acc), dtype, _semantic)


@builtin
def min2(lhs, rhs, dtype=None, _semantic=None):
    """Select the minimum with a native packed half-precision instruction."""
    return _packed_arith("min", (lhs, rhs), dtype, _semantic)


@builtin
def max2(lhs, rhs, dtype=None, _semantic=None):
    """Select the maximum with a native packed half-precision instruction."""
    return _packed_arith("max", (lhs, rhs), dtype, _semantic)


@dataclass(frozen=True, eq=True)
class TensorMemoryLayout:
    """
    Describes the layout for tensor memory in Blackwell architecture.

    Args:
        block (Tuple[int, int]): Number of contiguous elements per row / column in a CTA.
        col_stride (int): Number of 32-bit columns to advance between logically
            adjacent columns. Packed layouts use a stride of 1. Unpacked
            layouts use ``32 / bitwidth``.
        cga_layout (Optional[List[List[int]]]): CGA layout bases. Defaults to [].
        two_ctas (bool): Whether the layout is for two-CTA mode. Defaults to False.
        fp4_padded (bool): Whether byte-backed operand A uses the padded MMAv5
            FP4 layout. Its descriptor keeps the packed ``Mx(K/2)xi8`` shape,
            MMAv5 treats logical K as twice descriptor K, and physical TMEM
            reserves one byte per logical FP4 element. Defaults to False.
    """
    block: Tuple[int, int]
    col_stride: int
    cga_layout: List[List[int]] = field(default_factory=list)
    two_ctas: bool = False
    fp4_padded: bool = False

    def __post_init__(self):
        super().__setattr__("block", _unwrap_if_constexpr(self.block))
        super().__setattr__("col_stride", _unwrap_if_constexpr(self.col_stride))
        super().__setattr__("cga_layout", _unwrap_if_constexpr(self.cga_layout))
        super().__setattr__("two_ctas", _unwrap_if_constexpr(self.two_ctas))
        super().__setattr__("fp4_padded", _unwrap_if_constexpr(self.fp4_padded))
        assert len(self.block) == 2
        assert all(len(basis) == 2 for basis in self.cga_layout)
        assert self.col_stride >= 1 and (self.col_stride &
                                         (self.col_stride - 1)) == 0, "tensor memory col_stride must be a power of two"

    def _to_ir(self, builder):
        return builder.get_tensor_memory_layout(
            self.block,
            self.col_stride,
            [list(basis) for basis in self.cga_layout],
            self.two_ctas,
            self.fp4_padded,
        )

    def mangle(self) -> str:
        block_str = f"{self.block[0]}x{self.block[1]}"
        stride_str = f"C{self.col_stride}"
        cga_layout_str = "_".join("~".join(map(str, basis)) for basis in self.cga_layout)
        two_ctas_str = "2CT" if self.two_ctas else ""
        fp4_padded_str = "FP4P" if self.fp4_padded else ""
        return f"TL{block_str}{stride_str}{cga_layout_str}{two_ctas_str}{fp4_padded_str}TL"

    def __hash__(self):
        return hash(
            (tuple(self.block), self.col_stride, tuple(tuple(b)
                                                       for b in self.cga_layout), self.two_ctas, self.fp4_padded))


@dataclass(frozen=True, eq=True)
class TensorMemoryScalesLayout:
    """
    Describes the layout for tensor memory scales in Blackwell architecture.

    Args:
        cga_layout (Optional[List[List[int]]]): CGA layout bases. Defaults to [].
    """
    cga_layout: List[List[int]] = field(default_factory=list)

    def __post_init__(self):
        super().__setattr__("cga_layout", _unwrap_if_constexpr(self.cga_layout))
        assert all(len(basis) == 2 for basis in self.cga_layout)

    def _to_ir(self, builder):
        return builder.get_tensor_memory_scales_layout([list(basis) for basis in self.cga_layout])

    def mangle(self) -> str:
        cga_layout_str = "_".join("~".join(map(str, basis)) for basis in self.cga_layout)
        return f"TLS{cga_layout_str}TLS"

    def __hash__(self):
        return hash(tuple(tuple(b) for b in self.cga_layout))


@dataclass(frozen=True)
class _TensorMemoryLinearLayout:
    """
    Print-only linear layout for TMEM (row/col -> dim0/dim1).
    """
    rows: List[List[int]]
    cols: List[List[int]]
    shape: List[int]

    def _to_ir(self, builder):
        raise RuntimeError("TensorMemoryLinearLayout is print-only; IR materialization is unsupported")

    def mangle(self):
        return f"TMLL_{self.shape}_TMLL"

    def __hash__(self):
        return hash((tuple(map(tuple, self.rows)), tuple(map(tuple, self.cols)), tuple(self.shape)))


def _unwrap_tmem_layout_arg(x):
    if isinstance(x, ttgl.constexpr):
        return _unwrap_tmem_layout_arg(x.value)
    if isinstance(x, list):
        return [_unwrap_tmem_layout_arg(i) for i in x]
    if isinstance(x, tuple):
        return tuple(_unwrap_tmem_layout_arg(i) for i in x)
    return x


class tensor_memory_descriptor_type(base_type):

    def __init__(self, element_ty, shape, layout, alloc_shape):
        self.element_ty = _unwrap_if_constexpr(element_ty)
        self.shape = _unwrap_if_constexpr(shape)
        self.layout = _unwrap_if_constexpr(layout)
        self.alloc_shape = _unwrap_if_constexpr(alloc_shape)
        assert isinstance(self.layout, (TensorMemoryLayout, TensorMemoryScalesLayout))

    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"

    @constexpr_function
    def get_reg_layout(self, num_warps=None, instr_variant="32x32b"):
        """
        Return a DistributedLinearLayout compatible with TMEM load/store
        instructions for this descriptor type.

        Args:
            num_warps (Optional[int]): Number of warps participating in the
                operation. Must be provided when it cannot be inferred by the
                caller.
            instr_variant (str): TMEM instruction variant (e.g. ``"32x32b"``).
        """
        tmem_ty = _unwrap_tmem_layout_arg(self)
        if not isinstance(tmem_ty, tensor_memory_descriptor_type):
            raise TypeError(f"expected a tensor_memory_descriptor_type but got {type(tmem_ty)!r}")
        num_warps = _unwrap_tmem_layout_arg(num_warps)
        if num_warps is None:
            raise ValueError("num_warps could not be inferred; pass a positive power of two")
        if not isinstance(num_warps, int) or num_warps <= 0 or (num_warps & (num_warps - 1)) != 0:
            raise ValueError(f"num_warps must be a positive power of two, got {num_warps!r}")

        return _compute_tmem_reg_layout(
            _unwrap_tmem_layout_arg(tmem_ty.element_ty),
            _unwrap_tmem_layout_arg(tmem_ty.shape),
            _unwrap_tmem_layout_arg(tmem_ty.alloc_shape),
            _unwrap_tmem_layout_arg(tmem_ty.layout),
            num_warps,
            _unwrap_tmem_layout_arg(instr_variant),
        )


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
        self.type = tensor_memory_descriptor_type(element_ty, shape, layout, alloc_shape)

    def _set_name(self, builder: ir.builder, name: str) -> None:
        self.handle.set_loc(builder.create_name_loc(name, self.handle.get_loc()))

    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 get_reg_layout(self, num_warps=None, instr_variant="32x32b", _semantic: GluonSemantic = None, _generator=None):
        """
        Return the register layout used to access this tensor memory descriptor.

        Args:
            num_warps (Optional[int]): Number of warps participating in the
                operation. When omitted, infer it from the caller's context.
            instr_variant (str): TMEM instruction variant. Defaults to
                ``"32x32b"``.

        Returns:
            DistributedLayout: A register layout compatible with TMEM
            load/store instructions for this descriptor.
        """
        if num_warps is None:
            num_warps = ttgl.num_warps(_semantic=_semantic, _generator=_generator)
        return self.type.get_reg_layout(num_warps=num_warps, instr_variant=instr_variant)

    @builtin
    def load(self, layout=None, _semantic: GluonSemantic = None, _generator=None) -> ttgl.tensor:
        """
        Load a tensor from tensor memory.

        Args:
            layout (Optional[DistributedLayout]): Destination layout of the tensor.
                When omitted, infer the default TMEM register layout for this
                descriptor type and the caller's warp count.

        Returns:
            tensor: A distributed tensor containing the loaded data.
        """
        if layout is None:
            layout = self.get_reg_layout(_semantic=_semantic, _generator=_generator)
        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)

    def _load_red(self, layout, red_op, abs, propagate_nan, _semantic: GluonSemantic, _generator=None):
        #   red_op: MIN/MAX reduction operation
        #   abs (bool): If True, reduce absolute values.
        #   propagate_nan (NONE): If ALL, propagate NaN in specified reduction operation.
        if layout is None:
            layout = self.get_reg_layout(_semantic=_semantic, _generator=_generator)
        layout = _unwrap_if_constexpr(layout)
        abs_flag = _unwrap_if_constexpr(abs)
        propagate_nan = _unwrap_if_constexpr(propagate_nan)

        ret_ty = ttgl.distributed_type(self.dtype, self.shape, layout)
        builder = _semantic.builder

        result, reduced, red_layout = builder.create_tmem_load(ret_ty.to_ir(builder), self.handle, red_op, abs_flag,
                                                               propagate_nan)

        red_shape = [self.shape[0]]  # [M] for [M,N] input
        red_ty = ttgl.distributed_type(self.dtype, red_shape, red_layout)

        return (ttgl.tensor(result, ret_ty), ttgl.tensor(reduced, red_ty))

    @builtin
    def load_min(self, layout=None, abs=False, propagate_nan=ir.PROPAGATE_NAN.NONE, _semantic: GluonSemantic = None,
                 _generator=None):
        """
        Load a tensor from tensor memory with MIN reduction along the N-dimension.

        Args:
            layout (Optional[DistributedLayout]): Destination layout of the tensor.
                When omitted, infer the default TMEM register layout for this
                descriptor type and the caller's warp count.
            abs (bool): If True, reduce absolute values. Defaults to False.
            propagate_nan (PROPAGATE_NAN): If ALL, propagate NaN in the reduction operation. Defaults to NONE.

        Returns:
            tuple: A tuple containing (tensor, reduced_tensor) where tensor is the loaded data
                   and reduced_tensor is the result of MIN reduction along the N-dimension of loaded data
        """
        return self._load_red(layout, gluon_ir.TMEM_LOAD_REDUCE_MODIFIER.MIN, abs, propagate_nan, _semantic, _generator)

    @builtin
    def load_max(self, layout=None, abs=False, propagate_nan=ir.PROPAGATE_NAN.NONE, _semantic: GluonSemantic = None,
                 _generator=None):
        """
        Load a tensor from tensor memory with MAX reduction along the N-dimension.

        Args:
            layout (Optional[DistributedLayout]): Destination layout of the tensor.
                When omitted, infer the default TMEM register layout for this
                descriptor type and the caller's warp count.
            abs (bool): If True, reduce absolute values. Defaults to False.
            propagate_nan (PROPAGATE_NAN): If ALL, propagate NaN in the reduction operation. Defaults to NONE.

        Returns:
            tuple: A tuple containing (tensor, reduced_tensor) where tensor is the loaded data
                   and reduced_tensor is the result of MAX reduction along the N-dimension of loaded data.
        """
        return self._load_red(layout, gluon_ir.TMEM_LOAD_REDUCE_MODIFIER.MAX, abs, propagate_nan, _semantic, _generator)

    @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, dim=1, _semantic: GluonSemantic = None) -> None:
        """
        Create a slice of the tensor memory descriptor along a given dimension.

        Args:
            start (int): The starting index for subslice.
            length (int): The length of the subslice.
            dim (int): The dimension to slice, including a leading pipeline-stage dimension (default: 1).

        Returns:
            tensor_memory_descriptor: Descriptor for the subslice.
        """
        start = _unwrap_if_constexpr(start)
        length = _unwrap_if_constexpr(length)
        dim = _unwrap_if_constexpr(dim)
        _check(isinstance(start, int), lambda: "start must be a constant int")
        _check(isinstance(length, int), lambda: "length must be a constant int")
        _check(isinstance(dim, int) and dim in range(len(self.shape)), lambda: "invalid slice dimension")
        shape = list(self.shape)
        shape[dim] = length
        layout = self.type.layout
        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, dim)
        return ret

    @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
        shape = self.shape[1:]
        layout = self.layout
        ret = tensor_memory_descriptor(None, self.dtype, shape, layout, shape)
        ret.handle = builder.create_memdesc_index(ret.type.to_ir(builder), self.handle, index.handle)
        return ret

    @builtin
    def reinterpret(self, dtype=None, shape=None, layout=None,
                    _semantic: GluonSemantic = None) -> tensor_memory_descriptor:
        """
        Reinterpret tensor memory descriptor with a new dtype, shape, and layout.

        Args:
            dtype (dtype): The new data type. Defaults to the descriptor dtype.
            shape (Sequence[int]): The new shape. Defaults to the descriptor shape.
            layout (TensorMemoryLayout): The new layout. Defaults to the descriptor layout.

        Returns:
            tensor_memory_descriptor: Descriptor with updated type and layout.
        """
        dtype = self.dtype if dtype is None else _unwrap_if_constexpr(dtype)
        shape = self.shape if shape is None else [_unwrap_if_constexpr(s) for s in shape]
        layout = self.layout if layout is None else _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):
    """
    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)
    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_copy(src, dst, _semantic=None):
    """
    Start an asynchronous copy from shared memory to tensor memory.

    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, multicast=False, 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 or tensor_memory_descriptor): Left hand side operand in shared or tensor memory.
        b (shared_memory_descriptor): Right hand side operand in shared 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.
        multicast (bool): Whether tcgen05 commit should multicast across a CTA cluster. Defaults to False.
        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)

    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.handle] * len(mbarriers)
        else:
            mbarrier_preds = _semantic._convert_to_ir_values(mbarrier_preds, require_i64=False)

    multicast = _unwrap_if_constexpr(multicast)
    _semantic.builder.create_tcgen05_mma(a.handle, b.handle, acc.handle, use_acc.handle, pred.handle, mbarriers,
                                         mbarrier_preds, acc.layout.two_ctas, multicast)


@builtin
def tcgen05_mma_scaled(a, b, acc, a_scale, b_scale, a_type, b_type, *, use_acc=True, pred=True, multicast=False,
                       mbarriers=None, mbarrier_preds=None, _semantic=None):
    """
    Emit a 5th generation TensorCore MMA scaled instruction.
    acc = (a * a_scale) * (b * b_scale) + (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).
        a_scale (tensor): Scale factor for operand A.
        b_scale (tensor): Scale factor for operand B.
        a_type (str): Type of operand A. One of {"e2m1", "e4m3", "e5m2"}.
        b_type (str): Type of operand B. One of {"e2m1", "e4m3", "e5m2"}.
        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.
        multicast (bool): Whether tcgen05 commit should multicast across a CTA cluster. Defaults to False.
        mbarriers (Sequence[mbarrier], 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)
    assert acc.type.layout.block[0] != 64, "tcgen05_mma_scaled does not support blockM=64"

    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.handle] * len(mbarriers)
        else:
            mbarrier_preds = _semantic._convert_to_ir_values(mbarrier_preds, require_i64=False)

    allowed_formats = {"e2m1", "e4m3", "e5m2"}
    assert a_type.value in allowed_formats, f"Unsupported lhs_format: {a_type.value}"
    assert b_type.value in allowed_formats, f"Unsupported rhs_format: {b_type.value}"
    a_type = _semantic._str_to_fp_type(a_type.value)
    b_type = _semantic._str_to_fp_type(b_type.value)
    multicast = _unwrap_if_constexpr(multicast)
    _semantic.builder.create_tcgen05_mma_scaled(a.handle, b.handle, acc.handle, a_scale.handle, b_scale.handle, a_type,
                                                b_type, use_acc.handle, pred.handle, mbarriers, mbarrier_preds,
                                                acc.layout.two_ctas, multicast)


@constexpr_function
def tcgen05_mma_barrier_count(smems, multicast, two_ctas):
    """
    Calculate the number of CTAs that will commit the tcgen05 MMA instruction.

    Args:
        smems (Sequence[shared_memory_descriptor]): Shared memory descriptors used in the tcgen05 instruction.
        multicast (bool): Whether the tcgen05 instruction is multicast.
        two_ctas (bool): Whether the tcgen05 instruction uses cta_group::2.

    Returns:
        int: The number of CTAs that will commit the tcgen05 MMA instruction.
    """
    assert 0 <= len(smems) <= 4, "tcgen05_mma_barrier_count supports 0 to 4 descriptors"
    if not smems or not multicast:
        return 1

    def basis_is_zero(basis):
        return all(b == 0 for b in basis)

    num_cta_bits = len(smems[0].layout.cga_layout)
    for desc in smems[1:]:
        assert len(desc.layout.cga_layout) == num_cta_bits

    num_cta_commits = 0
    for cta in range(1 << num_cta_bits):
        if two_ctas and cta & 1:
            continue
        for desc in smems:
            if all(basis_is_zero(basis) or not (cta & (1 << i)) for i, basis in enumerate(desc.layout.cga_layout)):
                num_cta_commits += 1
                break
    return num_cta_commits


@builtin
def tcgen05_commit(barrier, pred=True, descs=(), _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.

    If `descs` are provided, the commit will be multicast across the CTA cluster
    based on the shared layouts of those descriptors. This should be used when
    the inputs to the tcgen5 MMA come from TMA descriptors using multicast.

    Args:
        barrier (shared_memory_descriptor): The barrier to track completion of tcgen05 MMA and copy instructions.
        pred (bool): Scalar predicate. Operation is skipped if predicate is False. Defaults to True.
        descs (Sequence[shared_memory_descriptor]): Shared memory descriptors for
            the preceding multiplication inputs. Defaults to ().
    """
    pred = _semantic.to_tensor(pred)
    descs = _unwrap_if_constexpr(descs)
    descs = [d.handle for d in descs]
    _semantic.builder.create_tcgen05_commit(barrier.handle, pred.handle, descs)
