"""Builtin operations exposed on the TileLang language surface."""

from __future__ import annotations

import tvm.tirx.script.parser as T
from tilelang._typing import BufferLikeType, BufferLikeTypeTuple, BarrierType, DType
from tilelang import tvm as tvm
from tilelang.language import ptx_arrive_barrier, evaluate
from tilelang.language.eager.builder import macro
from tilelang.language.kernel import get_thread_bindings, get_block_extents
from tvm import DataType, tirx
from tvm.runtime import convert
from tvm.tirx import PrimExpr, Var, Call, BufferLoad, BufferRegion
from tilelang.utils.language import retrieve_ptr, get_buffer_region_from_load, retrieve_buffer_and_offset


def _normalize_index_arg(value: int | PrimExpr | None) -> PrimExpr | None:
    """
    Normalize warp sizing arguments so both Python ints and PrimExpr values
    are accepted uniformly.
    """
    if value is None:
        return None
    if isinstance(value, PrimExpr):
        return value
    if isinstance(value, int):
        return tirx.IntImm("int32", value)
    raise TypeError(f"Expect warp sizing argument to be int or PrimExpr, but got {type(value)}.")


def _mbar_to_buffer_load(mbar: BarrierType) -> BufferLoad:
    """Convert a memory barrier to a buffer load.

    Args:
        mbar: BarrierType
            The memory barrier to convert

    Returns:
        tirx.BufferLoad: A buffer load of the memory barrier
    """
    if isinstance(mbar, tirx.BufferLoad):
        return mbar
    elif isinstance(mbar, tirx.Buffer):
        assert len(mbar.shape) == 1, f"mbarrier must be a single element buffer, but got {mbar.shape}"
        return tirx.BufferLoad(mbar, [0])
    else:
        raise TypeError(f"mbarrier must be an tirx.BufferLoad or a tirx.Buffer, but got {type(mbar)}")


def __ldg(load_or_buf: BufferLoad | tirx.Buffer, index: PrimExpr | int | None = None) -> PrimExpr:
    """Explicitly load via CUDA read-only data cache.

    Prefer calling with a BufferLoad: `T.__ldg(x[i])` emits `__ldg(&x[i])` on CUDA.
    On non-CUDA backends, falls back to a regular load.

    Args:
        load_or_buf: A `BufferLoad` like `x[i]`, or a `Buffer`.
        index: Optional index when passing a `Buffer` directly.

    Returns:
        PrimExpr: The loaded value.
    """
    if isinstance(load_or_buf, BufferLoad):
        dtype = load_or_buf.dtype
        return tirx.call_intrin(str(dtype), tirx.op.Op.get("tl.__ldg"), load_or_buf)
    if isinstance(load_or_buf, tirx.Buffer):
        if index is None:
            raise ValueError("T.__ldg(Buffer, index) requires an index when passing a Buffer.")
        idx = index
        if isinstance(index, (list, tuple)):
            if len(index) != 1:
                raise ValueError("T.__ldg currently supports 1D flattened indices.")
            idx = index[0]
        bl = BufferLoad(load_or_buf, [idx])
        return tirx.call_intrin(str(load_or_buf.dtype), tirx.op.Op.get("tl.__ldg"), bl)
    raise TypeError("T.__ldg expects a BufferLoad or a Buffer.")


def __ffs(value: int | PrimExpr) -> PrimExpr:
    """Find the position of the least significant set bit.

    Lowers to CUDA ``__ffs`` for 32-bit integer inputs and ``__ffsll`` for
    64-bit integer inputs. The return value follows CUDA semantics: one-based
    bit position, or 0 when ``value`` is zero.
    """
    value = tirx.convert(value)
    return tirx.call_intrin("int32", tirx.op.Op.get("tl.__ffs"), value)


def __fns(mask: int | PrimExpr, base: int | PrimExpr, offset: int | PrimExpr) -> PrimExpr:
    """Find the position of the nth set bit in a bitmask.

    Lowers to CUDA ``__fns``. Returns a zero-based bit position (0-31), or
    ``0xFFFFFFFF`` when not found. ``mask`` is a 32-bit bitmask; ``base`` is
    the starting bit (0-31); ``offset`` selects the nth set bit (1-based).
    """
    mask = tirx.convert(mask)
    base = tirx.convert(base)
    offset = tirx.convert(offset)
    return tirx.call_intrin("uint32", tirx.op.Op.get("tl.__fns"), mask, base, offset)


def access_ptr(
    base: BufferLikeType,
    access_type: str | int = "r",
    *extents: PrimExpr | int | tuple[PrimExpr | int, ...] | list[PrimExpr | int],
    offset: PrimExpr | int = 0,
    extent: PrimExpr | int | None = None,
    ignore_last_ndim: int = 0,
) -> PrimExpr:
    """Create a TileLang `tl.access_ptr` from a buffer-like base location.

    This is a frontend convenience wrapper that keeps a `BufferLoad` argument
    in the resulting call so downstream passes can recover the referenced
    `tirx.Buffer` (including strides/storage scope) *and* the `rw_mask`
    (read/write intent) required by synchronization and safety checks.

    The returned `tl.access_ptr` is expected to be lowered to
    `tirx.builtin.tvm_access_ptr` later in the TileLang compilation pipeline.

    Parameters
    ----------
    base : BufferLikeType
        The base location to take the address of. Supported:
        - `tirx.BufferLoad` (e.g. `A[i, j]`): pointer to that element
        - `tirx.BufferRegion`: pointer to the region minima
        - `tirx.Buffer`: pointer to the beginning of the buffer
        - `tirx.Var` with let-binding to one of the above (inside TileLang frame)

    access_type : str | int
        Access mask for the pointer. Common string forms: `"r"`, `"w"`, `"rw"`.
        Integer bitmask is also accepted (1=read, 2=write, 3=read-write).

    *extents : PrimExpr | int
        Optional per-axis extents. When provided and `extent` is not specified,
        the 1D `extent` passed to `tvm_access_ptr` is computed as the product of
        the provided extents (padding leading dimensions with 1 if needed).

        For example:
        - `T.access_ptr(A[i], "r")` -> extent defaults to 1 (element pointer)
        - `T.access_ptr(A[i], "r", 16)` -> extent=16
        - `T.access_ptr(A[i, j], "r", m, n)` -> extent=m*n

    offset : PrimExpr | int
        Additional element offset from the base location.

    extent : PrimExpr | int | None
        Optional explicit 1D extent override (in elements). If provided, it
        takes precedence over `*extents`.

    ignore_last_ndim : int
        If non-zero, the base linear offset is computed only over the leading
        dimensions, ignoring the last `ignore_last_ndim` axes. This is useful
        when treating an N-D buffer as a view of its trailing sub-tensor.

    Returns
    -------
    ptr : PrimExpr
        A handle-typed `tirx.Call` to `tl.access_ptr`.
    """

    from tilelang.language.frame import has_let_value, get_let_value
    from tilelang.language.utils import get_buffer_region_from_load

    if isinstance(base, tirx.Var) and has_let_value(base):
        base = get_let_value(base)

    # Allow passing a single list/tuple as the extents argument.
    if len(extents) == 1 and isinstance(extents[0], (list, tuple)):
        extents = tuple(extents[0])

    def _rw_mask(access_type: str | int) -> int:
        if isinstance(access_type, int):
            return int(access_type)
        if isinstance(access_type, str):
            table = {"r": 1, "w": 2, "rw": 3}
            if access_type not in table:
                raise ValueError(f'Invalid access_type="{access_type}", expected one of {sorted(table.keys())}.')
            return table[access_type]
        raise TypeError(f"T.access_ptr access_type must be str or int, but got {type(access_type)}.")

    def _index_dtype(buf: tirx.Buffer) -> str:
        if len(buf.shape) > 0:
            return str(buf.shape[0].dtype)
        return "int32"

    # Extract underlying buffer and per-axis minima (indices).
    inferred_region_extents: list[PrimExpr] | None = None
    is_buffer_load_base = False
    if isinstance(base, BufferLoad):
        is_buffer_load_base = True
        buf = base.buffer
        region = get_buffer_region_from_load(base)
        if region is not None:
            mins = [r.min for r in region.region]
            inferred_region_extents = [r.extent for r in region.region]
        else:
            mins = list(base.indices)
    elif isinstance(base, BufferRegion):
        buf = base.buffer
        mins = [r.min for r in base.region]
        inferred_region_extents = [r.extent for r in base.region]
    elif isinstance(base, tirx.Buffer):
        buf = base
        idx_dtype = _index_dtype(buf)
        mins = [tirx.IntImm(idx_dtype, 0) for _ in buf.shape]
    else:
        raise TypeError(f"T.access_ptr expects a Buffer, BufferLoad, BufferRegion, or a Var bound to one of them, but got {type(base)}.")

    # Apply ignore_last_ndim by zeroing-out the ignored tail indices.
    idx_dtype = _index_dtype(buf)
    ignore_last_ndim = int(ignore_last_ndim)
    if ignore_last_ndim != 0:
        upto = max(0, len(mins) - ignore_last_ndim)
        mins = list(mins[:upto]) + [tirx.IntImm(idx_dtype, 0) for _ in range(len(mins) - upto)]

    # Support non-zero `offset` only for 1D buffers in the frontend meta-op.
    if isinstance(offset, int):
        if offset != 0:
            if len(mins) != 1:
                raise ValueError(
                    "T.access_ptr(offset!=0) is only supported for 1D buffers when emitting tl.access_ptr. "
                    "Use explicit indexing (e.g. A[i + off]) for N-D buffers."
                )
            mins = [mins[0] + tirx.IntImm(idx_dtype, offset)]
    elif isinstance(offset, PrimExpr):
        if not (isinstance(offset, tirx.IntImm) and int(offset.value) == 0):
            if len(mins) != 1:
                raise ValueError(
                    "T.access_ptr(offset!=0) is only supported for 1D buffers when emitting tl.access_ptr. "
                    "Use explicit indexing (e.g. A[i + off]) for N-D buffers."
                )
            mins = [mins[0] + offset]
    else:
        raise TypeError(f"T.access_ptr offset must be int or PrimExpr, but got {type(offset)}.")

    base_load = BufferLoad(buf, mins)

    # Compute 1D extent (in elements).
    extent_1d: PrimExpr
    if extent is not None:
        extent_1d = convert(extent)
    elif len(extents) > 0:
        exts = [convert(e) for e in extents]
        if len(exts) > len(buf.shape):
            raise ValueError(f"T.access_ptr got {len(exts)} extents for a buffer with ndim={len(buf.shape)}.")
        if len(exts) < len(buf.shape):
            pad = [tirx.IntImm(idx_dtype, 1) for _ in range(len(buf.shape) - len(exts))]
            exts = pad + exts
        extent_1d = tirx.IntImm(idx_dtype, 1)
        for e in exts:
            extent_1d = extent_1d * e
    else:
        # Match `tirx.Buffer.access_ptr` defaults:
        # - BufferLoad base: element pointer (extent=1)
        # - BufferRegion base: product of region extents
        # - Buffer base: full buffer size (product of shape)
        if is_buffer_load_base:
            extent_1d = tirx.IntImm(idx_dtype, 1)
        elif inferred_region_extents is not None:
            extent_1d = tirx.IntImm(idx_dtype, 1)
            for e in inferred_region_extents:
                extent_1d = extent_1d * convert(e)
        else:
            extent_1d = tirx.IntImm(idx_dtype, 1)
            for dim in buf.shape:
                extent_1d = extent_1d * convert(dim)

    return tirx.call_intrin(
        "handle",
        tirx.op.Op.get("tl.access_ptr"),
        base_load,
        extent_1d,
        tirx.IntImm("int32", _rw_mask(access_type)),
    )


def deallocate_tmem(tmem: tirx.Buffer) -> None:
    """Explicitly deallocate a TMEM buffer allocated by ``T.alloc_tmem``.

    By default, TileLang inserts a TMEM deallocation automatically at the end
    of the allocation block. Calling ``T.deallocate_tmem(buf)`` suppresses that
    automatic tail deallocation for ``buf`` and lowers an explicit deallocation
    at the call site instead.

    Notes:
    - The deallocation must obey the hardware TMEM rules: it should be issued by
      the same warp that performed the allocation.
    - Once this API is used, the buffer lifetime is user-managed for the current
      block; deallocating too early or conditionally is the user's responsibility.

    Args:
        tmem: A TMEM buffer previously returned by ``T.alloc_tmem``.
    """

    if not isinstance(tmem, tirx.Buffer):
        raise TypeError(f"T.deallocate_tmem expects a tvm.tirx.Buffer, but got {type(tmem)}.")
    if tmem.scope() != "shared.tmem":
        raise ValueError(f"T.deallocate_tmem expects a shared.tmem buffer, but got scope={tmem.scope()}.")

    return evaluate(tirx.call_intrin("handle", tirx.op.Op.get("tl.deallocate_tmem"), tmem.data))


def create_tma_descriptor(*args):
    """Create a Tensor Memory Access (TMA) descriptor.

    This is an internal API used by copy lowering. The argument list depends on
    the tensor rank and encodes the full TMA descriptor configuration:

        create_tma_descriptor(data_type, rank, global_addr,
            global_shape..., global_stride..., smem_box..., smem_stride...,
            interleave, swizzle, l2_promotion, oob_fill)

    Total arguments: 7 + 4 * rank.

    Returns:
        tirx.Call: A handle to the created TMA descriptor
    """
    return tirx.call_intrin("handle", tirx.op.Op.get("tl.create_tma_descriptor"), *args)


def tma_load(*args):
    """Perform a Tensor Memory Access (TMA) load operation.

    This is an internal API used by copy lowering. Arguments:

        tma_load(descriptor, mbarrier, smem_addr, coord_0, ..., coord_n, eviction_policy)

    Returns:
        tirx.Call: A handle to the TMA load operation
    """
    return tirx.call_intrin("handle", tirx.op.Op.get("tl.tma_load"), *args)


def tma_load_2sm(*args):
    """Perform a TMA load with 2SM (two Streaming Multiprocessors) on Blackwell.

    This is an internal API. Same arguments as :func:`tma_load`, but with
    the ``use_2cta`` annotation enabled for 2-CTA cooperative loading.

    Returns:
        tirx.Call: A handle to the TMA load operation
    """
    return tirx.call_intrin("handle", tirx.op.Op.get("tl.tma_load"), *args, annotations={"use_2cta": 1})


def fence_proxy_async():
    """Issue a shared memory fence for asynchronous proxy operations.

    Ensures that prior asynchronous operations (e.g. TMA stores) are visible
    to subsequent memory accesses. Maps to ``fence.proxy.async.shared::cta``.

    Returns:
        tirx.Call: A handle to the fence operation
    """
    return tirx.call_intrin("handle", tirx.op.Op.get("tl.fence_proxy_async"))


def tma_store_arrive():
    """Signal the arrival of a TMA store operation.

    Commits the current group of outstanding TMA store operations.
    Maps to ``cp.async.bulk.commit_group``.

    Returns:
        tirx.Call: A handle to the store arrive operation
    """
    return tirx.call_intrin("handle", tirx.op.Op.get("tl.tma_store_arrive"))


def tma_store_wait(count: int = 0, read: bool = True):
    """Wait for completion of TMA store operations.

    Waits until the number of outstanding TMA store groups is at most ``count``.
    Maps to ``cp.async.bulk.wait_group.read <count>`` by default, or
    ``cp.async.bulk.wait_group <count>`` when ``read`` is false.

    Args:
        count (int): The maximum number of outstanding store groups allowed
            to remain in flight. Defaults to 0 (wait for all stores to complete).
        read (bool): Whether to use the PTX ``.read`` modifier, which only
            waits for the source reads to complete. Defaults to True.

    Returns:
        tirx.Call: A handle to the store wait operation
    """
    return tirx.call_intrin("handle", tirx.op.Op.get("tl.tma_store_wait"), count, read)


def set_max_nreg(reg_count: int, is_inc: int):
    """Set the maximum number of registers to use.
    Detailed Documentation:
    https://docs.nvidia.com/cuda/parallel-thread-execution/#miscellaneous-instructions-setmaxnreg

    Args:
        reg_count: int
            The number of registers to allocate
        is_inc: int
            Whether to increment or decrement the register count
            0 if decrement, 1 if increment

    Returns:
        tirx.Call: A handle to the register setting operation
    """
    return tirx.call_intrin("handle", tirx.op.Op.get("tl.set_max_nreg"), reg_count, is_inc)


def inc_max_nreg(reg_count: int):
    """Increment the maximum number of registers to use."""
    return set_max_nreg(reg_count, 1)


def dec_max_nreg(reg_count: int):
    """Decrement the maximum number of registers to use."""
    return set_max_nreg(reg_count, 0)


def annotate_producer_reg_dealloc(reg_count: int = 24):
    """Annotate the producer reg dealloc."""
    return tirx.call_intrin(
        "handle",
        tirx.op.Op.get("tl.annotate_producer_reg_dealloc"),
        reg_count,
    )


def annotate_consumer_reg_alloc(reg_count: int = 240):
    """Annotate the consumer reg alloc."""
    return tirx.call_intrin(
        "handle",
        tirx.op.Op.get("tl.annotate_consumer_reg_alloc"),
        reg_count,
    )


def no_set_max_nreg():
    """Disable the maximum register limit setting."""
    return tirx.call_intrin("handle", tirx.op.Op.get("tl.no_set_max_nreg"))


def disable_warp_group_reg_alloc():
    """Disable the warp group reg alloc."""
    return no_set_max_nreg()


def ptx_arrive_cluster_barrier(mbarrier: BarrierType, cta_id: int | Var):
    """Arrive at a shared barrier in cluster.

    Args:
        mbarrier: BarrierType
            The memory barrier to arrive at
        cta_id: int | Var
            The peer CTA rank in cluster to arrive at.
    """
    return tirx.call_intrin("handle", tirx.op.Op.get("tl.ptx_arrive_cluster_barrier"), mbarrier, cta_id)


def mbarrier_wait_parity(mbarrier: BarrierType, parity: int | Var):
    """Wait for memory barrier parity condition.

    Args:
            mbarrier: BarrierType
            The memory barrier to wait on
        parity: int | Var
            The parity value to wait for
    Examples:
        .. code-block:: python

            mbar = T.alloc_barrier(1)
            # Wait for parity 0 on a single mbarrier
            T.mbarrier_wait_parity(mbar, 0)

            mbars = T.alloc_barrier([128] * n)
            # Wait for parity value on one of the mbarriers
            T.mbarrier_wait_parity(mbars[ko], ko)

            # Common usage in pipelined kernels:
            for ko in range(num_stages):
                # Producer waits for consumer to finish previous iteration
                T.mbarrier_wait_parity(mbars[1], ko ^ 1)
                # Producer copies data
                T.copy(A_global, A_shared)
                # Producer signals data ready
                T.mbarrier_arrive(mbars[0])

                # Consumer waits for producer data
                T.mbarrier_wait_parity(mbars[0], ko)
                # Consumer computes
                T.gemm(A_shared, B_shared, C_local)
                # Consumer signals completion
                T.mbarrier_arrive(mbars[1])
    Returns:
        tirx.Call: A handle to the barrier wait operation
    """
    mbarrier = _mbar_to_buffer_load(mbarrier)
    return tirx.call_intrin("handle", tirx.op.Op.get("tl.mbarrier_wait_parity"), mbarrier, parity)


def mbarrier_arrive(mbarrier: BarrierType, cta_id: int | Var | None = None):
    """Arrive at memory barrier.

    Args:
        mbarrier: BarrierType
            The memory barrier to arrive at
        cta_id: int | Var | None
            The peer CTA rank in cluster to arrive at. (Only valid for cluster barriers)
            If not provided, will arrive on current CTA's barrier.
    """
    mbarrier = _mbar_to_buffer_load(mbarrier)
    if cta_id is not None:
        assert mbarrier.buffer.scope() == "shared.cluster_barrier", f"mbarrier must be a cluster barrier, but got {mbarrier.buffer.scope}"
        return ptx_arrive_cluster_barrier(mbarrier, cta_id)
    else:
        return ptx_arrive_barrier(mbarrier)


def mbarrier_expect_tx(mbarrier: BarrierType, tx: int):
    """Set expected transaction count for memory barrier.

    Args:
        mbarrier: BarrierType
            The memory barrier to expect transaction count for
        tx: int
            The expected transaction count

    Returns:
        tirx.Call: A handle to the barrier expectation operation
    """
    mbarrier = _mbar_to_buffer_load(mbarrier)
    return tirx.call_intrin("handle", tirx.op.Op.get("tl.mbarrier_expect_tx"), mbarrier, tx)


def mbarrier_arrive_expect_tx(mbarrier: BarrierType, tx: int):
    """Arrive at a memory barrier and expect completion of async transactions."""
    from tilelang.language.tir.op import ptx_arrive_barrier_expect_tx

    mbarrier = _mbar_to_buffer_load(mbarrier)
    return ptx_arrive_barrier_expect_tx(mbarrier, tx)


def warpgroup_arrive():
    """Signal warpgroup readiness for subsequent WGMMA operations.

    Returns:
        tirx.Call: A handle to the warpgroup arrive operation.
    """
    return tirx.call_intrin("handle", tirx.op.Op.get("tl.warpgroup_arrive"))


def warpgroup_commit_batch():
    """Commit the current warpgroup batch for WGMMA operations.

    Returns:
        tirx.Call: A handle to the warpgroup commit batch operation.
    """
    return tirx.call_intrin("handle", tirx.op.Op.get("tl.warpgroup_commit_batch"))


def warpgroup_wait(num_mma: int):
    """Wait for completion of the specified warpgroup batch.

    Args:
        num_mma: int
            Identifier of the warpgroup MMA batch to wait on.

    Returns:
        tirx.Call: A handle to the warpgroup wait operation.
    """
    return tirx.call_intrin("handle", tirx.op.Op.get("tl.warpgroup_wait"), num_mma)


def get_lane_idx(
    warp_size: int | PrimExpr | None = None,
) -> PrimExpr:
    """Return the logical lane index of the calling thread within a warp.

    Parameters
    ----------
    warp_size : Optional[int, PrimExpr]
        Logical warp (or wavefront) size. Defaults to 32 on NVIDIA and 64 on AMD.

    Example
    -------
    >>> lane = T.get_lane_idx()
    >>> custom_lane = T.get_lane_idx(64)  # override warp size explicitly

    Implementation Notes
    --------------------
    Lowers to the CUDA helper `tl::get_lane_idx(warp_size)` defined in
    `src/tl_templates/cuda/intrin.h`, which computes the lane index from the
    linear thread id using the provided `warp_size`.
    """
    warp_size_expr = _normalize_index_arg(warp_size)
    if warp_size_expr is None:
        return tirx.call_intrin("int32", tirx.op.Op.get("tl.get_lane_idx"))
    return tirx.call_intrin("int32", tirx.op.Op.get("tl.get_lane_idx"), warp_size_expr)


def get_warp_idx_sync(
    warp_size: int | PrimExpr | None = None,
) -> PrimExpr:
    """Return the canonical warp index, assuming the warp's threads are converged.

    Parameters
    ----------
    warp_size : Optional[int, PrimExpr]
        Logical warp size used for the index calculation.

    Example
    -------
    >>> warp = T.get_warp_idx_sync()
    >>> custom_warp = T.get_warp_idx_sync(64)

    Implementation Notes
    --------------------
    Emits `tl::get_warp_idx_sync(warp_size)` which divides the block-linear
    thread id by `warp_size`, matching the semantics of CUTLASS' canonical helpers.
    """
    warp_size_expr = _normalize_index_arg(warp_size)
    if warp_size_expr is None:
        return tirx.call_intrin("int32", tirx.op.Op.get("tl.get_warp_idx_sync"))
    return tirx.call_intrin("int32", tirx.op.Op.get("tl.get_warp_idx_sync"), warp_size_expr)


def get_warp_idx(
    warp_size: int | PrimExpr | None = None,
) -> PrimExpr:
    """Return the canonical warp index without synchronizing the warp.

    Parameters
    ----------
    warp_size : Optional[int, PrimExpr]
        Logical warp size used for the index calculation.

    Example
    -------
    >>> warp = T.get_warp_idx()
    >>> custom_warp = T.get_warp_idx(64)

    Implementation Notes
    --------------------
    Lowers to `tl::get_warp_idx(warp_size)` which divides the block-linear
    thread id by the provided `warp_size` without requiring warp convergence.
    """
    warp_size_expr = _normalize_index_arg(warp_size)
    if warp_size_expr is None:
        return tirx.call_intrin("int32", tirx.op.Op.get("tl.get_warp_idx"))
    return tirx.call_intrin("int32", tirx.op.Op.get("tl.get_warp_idx"), warp_size_expr)


def get_warp_group_idx(
    warp_size: int | PrimExpr | None = None,
    warps_per_group: int | PrimExpr | None = None,
) -> PrimExpr:
    """Return the canonical warp group index for the calling thread.

    Parameters
    ----------
    warp_size : Optional[int, PrimExpr]
        Logical warp size to use (defaults to 32 on NVIDIA / 64 on AMD).
    warps_per_group : Optional[int, PrimExpr]
        Number of warps per warp-group. Defaults to 4 on NVIDIA architectures.

    Example
    -------
    >>> group = T.get_warp_group_idx()
    >>> custom_group = T.get_warp_group_idx(32, 6)  # treat 6 warps as a group

    Implementation Notes
    --------------------
    Generates `tl::get_warp_group_idx(warp_size, warps_per_group)` which
    divides the block-linear thread id by `warp_size * warps_per_group`,
    matching the canonical ordering while allowing architecture-specific overrides.
    """
    warp_size_expr = _normalize_index_arg(warp_size)
    warps_per_group_expr = _normalize_index_arg(warps_per_group)
    args = []
    if warp_size_expr is not None:
        args.append(warp_size_expr)
    if warps_per_group_expr is not None:
        if warp_size_expr is None:
            raise ValueError("get_warp_group_idx expects `warp_size` when specifying `warps_per_group`.")
        args.append(warps_per_group_expr)
    return tirx.call_intrin("int32", tirx.op.Op.get("tl.get_warp_group_idx"), *args)


def shuffle_elect(thread_extent: int) -> PrimExpr:
    """Elect exactly one lane within a logical thread group.

    Parameters
    ----------
    thread_extent : int
        Size (in threads) of the group in which a single lane should be elected.
        Passing 0 elects a single lane in the entire thread block.

    Example
    -------
    >>> is_leader = T.shuffle_elect(64)
    >>> T.if_then_else(is_leader, do_leader_work(), T.evaluate(0))

    Implementation Notes
    --------------------
    Lowered to the CUDA helper `tl::tl_shuffle_elect<thread_extent>()` defined in
    `src/tl_templates/cuda/intrin.h`, which relies on
    `cutlass::canonical_warp_idx()` and `cute::elect_one_sync()`
    to pick one lane per group.
    """
    return tirx.call_intrin("bool", tirx.op.Op.get("tl.tl_shuffle_elect"), thread_extent)


def warpgroup_fence_operand(
    buffer_or_ptr: BufferLikeType | PrimExpr,
    offset: int | PrimExpr = 0,
    num_regs: int | PrimExpr | None = None,
    dtype: DType | None = None,
):
    """Insert a warpgroup fence for the destination accumulator registers.

    This prevents NVCC from sinking uses of accumulator fragments past the corresponding
    WGMMA operations by issuing an empty inline assembly barrier on every register.

    Args:
        buffer_or_ptr: BufferLikeType | PrimExpr
            A buffer representing the accumulator fragment, a buffer load/region
            that identifies a starting element within the fragment, or a pointer expression
            (e.g., tvm_access_ptr/address_of/typed Var).
        offset: int | PrimExpr
            Element offset from the start of the accumulator fragment.
        num_regs: int | PrimExpr | None
            Number of 32-bit registers to fence. If None and a Buffer is provided, it will be
            derived from the buffer shape and dtype.
        dtype: DType | None
            Data type string of the accumulator elements. When passing a buffer or
            buffer-derived expression, dtype is inferred. It is required only when
            passing a raw pointer expression that cannot be inferred.

    Returns:
        tirx.Call: A handle to the warpgroup fence operation.
    """
    if isinstance(buffer_or_ptr, BufferLoad):
        # Treat BufferLoad as a request to fence starting from the loaded element's address
        buf = buffer_or_ptr.buffer
        data_ptr = buf.data
        inferred_dtype = buf.dtype
        if dtype is not None and dtype != inferred_dtype:
            raise ValueError(f"dtype mismatch: provided {dtype}, buffer uses {inferred_dtype}.")
        dtype = inferred_dtype
        # Compute element offset from indices using strides if present, otherwise row-major
        if len(buf.strides) == len(buf.shape) and len(buf.strides) > 0:
            elem_off = 0
            for idx, stride in zip(buffer_or_ptr.indices, buf.strides):
                elem_off = elem_off + idx * stride
        else:
            elem_off = 0
            stride_acc = 1
            for idx, dim in zip(reversed(buffer_or_ptr.indices), reversed(buf.shape)):
                elem_off = elem_off + idx * stride_acc
                stride_acc = stride_acc * dim
        # Combine with user-provided offset
        offset = elem_off + convert(offset)
        if num_regs is None:
            raise ValueError("num_regs must be provided when passing a BufferLoad.")
        return evaluate(
            tirx.call_intrin(
                "handle",
                tirx.op.Op.get("tl.warpgroup_fence_operand"),
                str(dtype),
                data_ptr,
                convert(offset),
                convert(num_regs),
            )
        )

    if isinstance(buffer_or_ptr, tirx.Buffer):
        data_ptr = buffer_or_ptr.data
        inferred_dtype = buffer_or_ptr.dtype
        if dtype is not None and dtype != inferred_dtype:
            raise ValueError(f"dtype mismatch: provided {dtype}, buffer uses {inferred_dtype}.")
        dtype = inferred_dtype
        if num_regs is None:
            total_elems = 1
            for dim in buffer_or_ptr.shape:
                if isinstance(dim, tirx.IntImm):
                    total_elems *= int(dim)
                else:
                    raise ValueError("warpgroup_fence_operand requires num_regs when buffer shape is symbolic.")
            bits_per_elem = DataType(dtype).bits
            num_regs = (total_elems * bits_per_elem + 31) // 32
    elif isinstance(buffer_or_ptr, BufferRegion):
        buf = buffer_or_ptr.buffer
        data_ptr = buf.data
        inferred_dtype = buf.dtype
        if dtype is not None and dtype != inferred_dtype:
            raise ValueError(f"dtype mismatch: provided {dtype}, buffer uses {inferred_dtype}.")
        dtype = inferred_dtype
        # Compute element offset from region min using strides if present, otherwise row-major
        if len(buf.strides) == len(buf.shape) and len(buf.strides) > 0:
            elem_off = 0
            for r, stride in zip(buffer_or_ptr.region, buf.strides):
                elem_off = elem_off + r.min * stride
        else:
            elem_off = 0
            stride_acc = 1
            for r, dim in zip(reversed(buffer_or_ptr.region), reversed(buf.shape)):
                elem_off = elem_off + r.min * stride_acc
                stride_acc = stride_acc * dim
        # Combine with user-provided offset
        offset = elem_off + convert(offset)
        # Try derive num_regs from region extents if fully static; otherwise require user input
        if num_regs is None:
            total_elems = 1
            static = True
            for r in buffer_or_ptr.region:
                if isinstance(r.extent, tirx.IntImm):
                    total_elems *= int(r.extent)
                else:
                    static = False
                    break
            if static:
                bits_per_elem = DataType(dtype).bits
                num_regs = (total_elems * bits_per_elem + 31) // 32
            else:
                raise ValueError("warpgroup_fence_operand requires num_regs when BufferRegion extent is symbolic.")
        return evaluate(
            tirx.call_intrin(
                "handle",
                tirx.op.Op.get("tl.warpgroup_fence_operand"),
                str(dtype),
                data_ptr,
                convert(offset),
                convert(num_regs),
            )
        )
    else:
        data_ptr = buffer_or_ptr
        # Try to infer dtype from common pointer expressions when not provided
        if dtype is None:
            inferred = None
            # Case 1: Pointer from Buffer.access_ptr -> tirx.builtin.tvm_access_ptr
            if isinstance(data_ptr, Call) and data_ptr.op.same_as(tirx.builtin.tvm_access_ptr()):
                # args[0] is a type annotation call; its dtype carries the element dtype
                inferred = str(data_ptr.args[0].dtype)
            # Case 2: Pointer from tirx.address_of(BufferLoad(...))
            elif isinstance(data_ptr, Call) and data_ptr.op.same_as(tirx.builtin.address_of()):
                # args[0] should be a BufferLoad; its dtype is the element dtype
                inferred = str(data_ptr.args[0].dtype)
            # Case 3: Typed pointer Var with PrimType element (typed TIR)
            elif hasattr(data_ptr, "type_annotation") and data_ptr.type_annotation is not None:
                try:
                    elem_ty = getattr(data_ptr.type_annotation, "element_type", None)
                    if elem_ty is not None and hasattr(elem_ty, "dtype"):
                        inferred = str(elem_ty.dtype)
                except Exception:
                    inferred = None
            if inferred is None:
                raise ValueError("dtype must be provided when passing a pointer expression and cannot be inferred.")
            dtype = inferred
        if num_regs is None:
            raise ValueError("num_regs must be provided when passing a pointer expression.")

    return evaluate(
        tirx.call_intrin(
            "handle",
            tirx.op.Op.get("tl.warpgroup_fence_operand"),
            str(dtype),
            data_ptr,
            convert(offset),
            convert(num_regs),
        )
    )


def wait_wgmma(id: int):
    """Wait for WGMMA (Warp Group Matrix Multiply-Accumulate) operations to complete.

    Args:
        id: int
            The id of the WGMMA operation to wait for

    Returns:
        tirx.Call: A handle to the WGMMA wait operation
    """
    return tirx.call_intrin("handle", tirx.op.Op.get("tl.wait_wgmma"), id)


def barrier_wait(mbarrier: BarrierType, parity: int | Var):
    """Wait for a memory barrier to complete.

    Args:
        mbarrier: BarrierType
            The memory barrier to wait on
        parity: int | Var
            The parity value to wait for
    Returns:
        tirx.Call: A handle to the barrier wait operation
    Current implementation is a sugar syntax for mbarrier_wait_parity, as we only support parity 0 and 1.
    """
    return mbarrier_wait_parity(mbarrier, parity)


def barrier_arrive(mbarrier: BarrierType):
    """Arrive at a memory barrier.

    Args:
        mbarrier: BarrierType
            The memory barrier to arrive at
    """
    return mbarrier_arrive(mbarrier)


# Full-warp mask as a proper uint32 TIR constant so the emitted C/C++ source
# prints as `0xFFFFFFFFu` instead of `(int64_t)4294967295` after TIR widening.
_FULL_WARP_MASK = tirx.const(0xFFFFFFFF, "uint32")
_DEFAULT_SHFL_WIDTH = 32


def _as_uint32_mask(mask: int | PrimExpr) -> PrimExpr:
    """Normalize a warp lane mask to a uint32 TIR expression.

    Python literals (e.g. ``0xFFFFFFFF``) would otherwise be widened to int64
    by TIR and printed as ``(int64_t)4294967295`` in the generated source.
    """
    if isinstance(mask, int):
        return tirx.const(mask, "uint32")
    return mask


def shfl_xor(
    value: int | PrimExpr | tirx.Call,
    delta: int | PrimExpr | tirx.Call,
    width: int | PrimExpr = _DEFAULT_SHFL_WIDTH,
    mask: int | PrimExpr = _FULL_WARP_MASK,
):
    """XOR-swap ``value`` across lanes (``__shfl_xor_sync`` on CUDA,
    ``__shfl_xor`` on HIP — mask ignored on HIP).
    """
    return tirx.call_intrin(value.dtype, tirx.op.Op.get("tl.shfl_xor_sync"), _as_uint32_mask(mask), value, delta, width)


def shfl_down(
    value: int | PrimExpr | tirx.Call,
    delta: int | PrimExpr | tirx.Call,
    width: int | PrimExpr = _DEFAULT_SHFL_WIDTH,
    mask: int | PrimExpr = _FULL_WARP_MASK,
):
    """Shift ``value`` down by ``delta`` lanes (``__shfl_down_sync`` on CUDA,
    ``__shfl_down`` on HIP).
    """
    return tirx.call_intrin(value.dtype, tirx.op.Op.get("tl.shfl_down_sync"), _as_uint32_mask(mask), value, delta, width)


def shfl_up(
    value: int | PrimExpr | tirx.Call,
    delta: int | PrimExpr | tirx.Call,
    width: int | PrimExpr = _DEFAULT_SHFL_WIDTH,
    mask: int | PrimExpr = _FULL_WARP_MASK,
):
    """Shift ``value`` up by ``delta`` lanes (``__shfl_up_sync`` on CUDA,
    ``__shfl_up`` on HIP).
    """
    return tirx.call_intrin(value.dtype, tirx.op.Op.get("tl.shfl_up_sync"), _as_uint32_mask(mask), value, delta, width)


def sync_threads(barrier_id: int = None, arrive_count: int = None):
    """Synchronize all threads in a block."""
    args = []
    if barrier_id is not None:
        args.append(barrier_id)
    if arrive_count is not None:
        args.append(arrive_count)
    return tirx.call_intrin("int32", "tirx.tvm_storage_sync", "shared", *args)


def named_barrier_arrive(barrier_id, thread_count):
    """CTA named barrier one-sided arrive (bar.arrive).

    Signals that the calling threads have arrived at the named barrier without
    waiting for other participants.  Unlike ``T.sync_threads(barrier_id, n)``
    which maps to ``bar.sync`` (arrive + wait), this call only *arrives* and
    returns immediately, allowing the calling warp group to continue working
    while the other side waits.

    This is useful in warp-specialized producer/consumer pipelines:

    .. code-block:: python

        # Producer warp group: signal readiness, keep going
        T.named_barrier_arrive(ready_barrier, total_threads)

        # Consumer warp group: block until producer has arrived
        T.sync_threads(ready_barrier, total_threads)

    Args:
        barrier_id:   Named barrier index (0-15). May be a variable (PrimExpr).
        thread_count: Total number of CTA threads participating in the barrier.
                      May be a variable (PrimExpr).

    Lowers to: ``asm volatile("bar.arrive %0, %1;" : : "r"(id), "r"(cnt));``
    """
    return tirx.call_intrin("handle", tirx.op.Op.get("tl.named_barrier_arrive"), barrier_id, thread_count)


def sync_warp(mask: int = None):
    """Synchronize all threads in a warp."""
    if mask is not None:
        return tirx.call_intrin("void", tirx.op.Op.get("tl.sync_warp"), mask)
    return tirx.call_intrin("void", tirx.op.Op.get("tl.sync_warp"))


def shfl_sync(
    value: int | PrimExpr,
    srcLane: int | PrimExpr,
    width: int | PrimExpr = _DEFAULT_SHFL_WIDTH,
    mask: int | PrimExpr = _FULL_WARP_MASK,
):
    """Broadcast ``value`` from ``srcLane`` to all lanes in the subgroup of
    ``width`` lanes (``__shfl_sync`` on CUDA, ``__shfl`` on HIP — mask ignored
    on HIP).
    """
    return tirx.call_intrin(value.dtype, tirx.op.Op.get("tl.shfl_sync"), _as_uint32_mask(mask), value, srcLane, width)


# ---------------------------------------------------------------------------
# Warp-vote / warp-ballot intrinsics
#
# The CUDA/HIP split and the uint32->uint64 zero-extension for ballot_sync and
# activemask are handled in codegen (src/target/codegen_{cuda,hip}.cc). These
# Python wrappers simply emit the backend-agnostic tl.* ops.
# ---------------------------------------------------------------------------


def any_sync(
    predicate: int | PrimExpr,
    mask: int | PrimExpr = _FULL_WARP_MASK,
) -> PrimExpr:
    """Non-zero if ANY active lane in ``mask`` has a non-zero ``predicate``.

    Lowers to ``__any_sync(mask, predicate)`` on CUDA and ``__any(predicate)``
    on HIP (the mask is ignored on HIP because the full wavefront is always
    convergent).

    Args:
        predicate: Integer condition to test.
        mask: Warp lane mask (defaults to ``0xFFFFFFFF``, i.e. all 32 lanes).

    Returns:
        int32: Non-zero if any thread in the mask has a non-zero predicate.
    """
    return tirx.call_intrin("int32", tirx.op.Op.get("tl.any_sync"), _as_uint32_mask(mask), predicate)


def all_sync(
    predicate: int | PrimExpr,
    mask: int | PrimExpr = _FULL_WARP_MASK,
) -> PrimExpr:
    """Non-zero only if ALL active lanes in ``mask`` have a non-zero predicate.

    Lowers to ``__all_sync(mask, predicate)`` on CUDA and ``__all(predicate)``
    on HIP.

    Args:
        predicate: Integer condition to test.
        mask: Warp lane mask (defaults to ``0xFFFFFFFF``, i.e. all 32 lanes).

    Returns:
        int32: Non-zero if all threads in the mask have a non-zero predicate.
    """
    return tirx.call_intrin("int32", tirx.op.Op.get("tl.all_sync"), _as_uint32_mask(mask), predicate)


def ballot_sync(
    predicate: int | PrimExpr,
    mask: int | PrimExpr = _FULL_WARP_MASK,
) -> PrimExpr:
    """Return a ``uint64`` bitmask of lanes in ``mask`` whose predicate is set.

    CUDA: ``__ballot_sync(mask, predicate)`` returns ``unsigned int``; codegen
    zero-extends it to ``uint64`` (upper 32 bits always zero for 32-wide warps).
    HIP: ``__ballot(predicate)`` returns ``uint64`` natively, covering all
    64 wavefront lanes. The mask argument is ignored on HIP.

    Returns:
        uint64: Bitmask with bit N set if lane N's predicate is non-zero.
    """
    return tirx.call_intrin("uint64", tirx.op.Op.get("tl.ballot_sync"), _as_uint32_mask(mask), predicate)


def ballot(predicate: int | PrimExpr) -> PrimExpr:
    """Full-warp / full-wavefront ballot. Equivalent to
    ``ballot_sync(predicate)`` (i.e. with the default full warp mask).

    Returns:
        uint64: Bitmask with bit N set if lane N's predicate is non-zero.
    """
    return tirx.call_intrin("uint64", tirx.op.Op.get("tl.ballot"), predicate)


def activemask() -> PrimExpr:
    """Return a ``uint64`` bitmask of currently active (non-exited) lanes.

    Lowers to ``__activemask()`` (zero-extended to ``uint64``) on CUDA and
    ``__ballot(1)`` on HIP.
    """
    return tirx.call_intrin("uint64", tirx.op.Op.get("tl.activemask"))


# ---------------------------------------------------------------------------
# Thread-block synchronization with predicate (shared across CUDA & HIP)
# ---------------------------------------------------------------------------


def syncthreads_count(predicate: int | PrimExpr) -> PrimExpr:
    """Block barrier that returns the number of threads whose ``predicate``
    evaluates to non-zero (``__syncthreads_count`` on CUDA and HIP).
    """
    return tirx.call_intrin("int32", tirx.op.Op.get("tl.syncthreads_count"), predicate)


def syncthreads_and(predicate: int | PrimExpr) -> PrimExpr:
    """Block barrier that returns non-zero only if ALL threads have a non-zero
    ``predicate`` (``__syncthreads_and`` on CUDA and HIP).
    """
    return tirx.call_intrin("int32", tirx.op.Op.get("tl.syncthreads_and"), predicate)


def syncthreads_or(predicate: int | PrimExpr) -> PrimExpr:
    """Block barrier that returns non-zero if ANY thread has a non-zero
    ``predicate`` (``__syncthreads_or`` on CUDA and HIP).
    """
    return tirx.call_intrin("int32", tirx.op.Op.get("tl.syncthreads_or"), predicate)


# ---------------------------------------------------------------------------
# Warp match intrinsics (CUDA sm_70+; unsupported on HIP)
# ---------------------------------------------------------------------------


def match_any_sync(
    value: int | PrimExpr,
    mask: int | PrimExpr = _FULL_WARP_MASK,
) -> PrimExpr:
    """Return a ``uint32`` bitmask of lanes in ``mask`` whose ``value`` equals
    the calling lane's value. Lowers to ``__match_any_sync`` on CUDA
    (compute capability >= 7.0). Not supported on HIP.
    """
    return tirx.call_intrin("uint32", tirx.op.Op.get("tl.match_any_sync"), _as_uint32_mask(mask), value)


def match_all_sync(
    value: int | PrimExpr,
    mask: int | PrimExpr = _FULL_WARP_MASK,
) -> PrimExpr:
    """Return ``mask`` if all lanes in ``mask`` agree on ``value``, else 0.

    Lowers to ``__match_all_sync`` on CUDA (compute capability >= 7.0); the
    trailing ``int*`` predicate output is hidden in codegen and discarded.
    Callers can reconstruct the predicate as ``result != 0``. Not supported
    on HIP.
    """
    return tirx.call_intrin("uint32", tirx.op.Op.get("tl.match_all_sync"), _as_uint32_mask(mask), value)


def sync_global():
    """Synchronize all threads in the entire grid."""
    tx, ty, tz = get_thread_bindings()
    ex, ey, ez = get_block_extents()
    print(tx, ty, tz, ex, ey, ez)
    args = ["global", tx == 0 and ty == 0 and tz == 0, ex * ey * ez]
    return evaluate(tirx.Call("handle", "tirx.tvm_storage_sync", args))


def sync_grid():
    """Synchronize all threads in a grid."""
    return tirx.call_intrin("handle", tirx.op.Op.get("tl.sync_grid"))


def initialize_wgmma_descriptor(
    descriptor: tirx.Buffer,
    start_address: PrimExpr,
    layout_type_: int = 0,
    leading_byte_offset: int = 0,
    stride_byte_offset: int = 0,
) -> PrimExpr:
    """Initialize a WGMMA/UTCMMA shared-memory descriptor."""

    if not isinstance(descriptor, (BufferLoad, tirx.Buffer)):
        raise TypeError("Descriptor must be a tvm.tirx.Buffer or tvm.tirx.BufferLoad.")

    if isinstance(descriptor, tirx.Buffer) and (len(descriptor.shape) != 1 or descriptor.shape[0] != 1):
        raise ValueError("Descriptor must be a 1D buffer of size 1.")

    descriptor = descriptor if isinstance(descriptor, BufferLoad) else tirx.BufferLoad(descriptor, [0])

    return evaluate(
        tirx.call_intrin(
            "handle",
            tirx.op.Op.get("tl.initialize_wgmma_descriptor"),
            descriptor,
            start_address,
            layout_type_,
            int(leading_byte_offset),
            int(stride_byte_offset),
        )
    )


def initialize_tcgen05_descriptor(
    descriptor: tirx.Buffer,
    start_address: PrimExpr,
    leading_byte_offset: int,
    stride_byte_offset: int,
    base_offset: int = 0,
    leading_is_absolute: bool = False,
    swizzle_mode: int = 0,
) -> PrimExpr:
    """Initialize a TCGEN05 shared-memory descriptor."""

    if not isinstance(descriptor, (BufferLoad, tirx.Buffer)):
        raise TypeError("Descriptor must be a tvm.tirx.Buffer or tvm.tirx.BufferLoad.")

    if isinstance(descriptor, tirx.Buffer) and (len(descriptor.shape) != 1 or descriptor.shape[0] != 1):
        raise ValueError("Descriptor must be a 1D buffer of size 1.")

    descriptor = descriptor if isinstance(descriptor, BufferLoad) else tirx.BufferLoad(descriptor, [0])

    return evaluate(
        tirx.call_intrin(
            "handle",
            tirx.op.Op.get("tl.initialize_tcgen05_descriptor"),
            descriptor,
            start_address,
            int(leading_byte_offset),
            int(stride_byte_offset),
            int(base_offset),
            tirx.IntImm("int32", 1 if leading_is_absolute else 0),
            int(swizzle_mode),
        )
    )


def increase_descriptor_offset(descriptor: PrimExpr, offset: PrimExpr) -> PrimExpr:
    """
    Increase the offset of a memory descriptor.

    Parameters:
        descriptor (PrimExpr): The memory descriptor to modify.
        offset (PrimExpr): The offset value to increase.

    Returns:
        PrimExpr: A handle representing the modified descriptor.
    """
    if not isinstance(descriptor, (BufferLoad, tirx.Buffer)):
        raise TypeError("Descriptor must be a tvm.tirx.Buffer or tvm.tirx.BufferLoad.")

    if isinstance(descriptor, tirx.Buffer) and len(descriptor.shape) != 1 or descriptor.shape[0] != 1:
        raise ValueError("Descriptor must be a 1D buffer of size 1.")

    descriptor = descriptor if isinstance(descriptor, BufferLoad) else tirx.BufferLoad(descriptor, [0])

    return evaluate(tirx.call_intrin("handle", tirx.op.Op.get("tl.increase_descriptor_offset"), descriptor, offset))


def loop_break():
    """Break out of the innermost loop."""
    return tirx.call_intrin("handle", tirx.op.Op.get("tl.loop_break"))


def cp_async_barrier_noinc(barrier: BarrierType):
    """Perform a ptx async copy barrier using cp.async.mbarrier.arrive.noinc."""
    barrier = _mbar_to_buffer_load(barrier)
    return tirx.call_intrin("handle", tirx.op.Op.get("tl.ptx_cp_async_barrier_noinc"), barrier)


def tcgen05_mma_arrive(mbar: tirx.Buffer | BufferLoad | PrimExpr, arrive_2cta: bool = False):
    """Signal UMMA (TCGEN05) barrier arrival for a shared-memory mbarrier pointer.

    Parameters
    ----------
    mbar: tirx.Buffer | BufferLoad | PrimExpr
        The mbarrier object in shared memory (e.g., Barrier*) or its address.
    arrive_2cta: bool
        Whether to also arrive at the peer CTA's barrier.
        If set, will be lowered to umma_arrive_multicast_2x1SM.
    """
    if isinstance(mbar, (tirx.Buffer, BufferLoad)):
        mbar = retrieve_ptr(mbar, access_type="rw")
    ann = {"use_2cta": 1} if arrive_2cta else {}
    return tirx.call_intrin("void", tirx.op.Op.get("tl.tcgen05_mma_arrive"), mbar, annotations=ann)


def tcgen05_before_thread_sync():
    return tirx.call_intrin("void", tirx.op.Op.get("tl.tcgen05_before_thread_sync"))


def tcgen05_after_thread_sync():
    return tirx.call_intrin("void", tirx.op.Op.get("tl.tcgen05_after_thread_sync"))


def _tcgen05_num_smem_chunks(smem_src, chunk_elems: int):
    if isinstance(smem_src, tirx.Buffer):
        shape = list(smem_src.shape)
    elif isinstance(smem_src, tirx.BufferRegion):
        shape = [r.extent for r in smem_src.region]
    elif isinstance(smem_src, tirx.BufferLoad):
        region = get_buffer_region_from_load(smem_src)
        if region is None:
            raise TypeError("T.tcgen05_cp_warpx4 requires Buffer/BufferRegion-like scale-factor sources.")
        shape = [r.extent for r in region.region]
    else:
        raise TypeError(f"Unsupported scale-factor buffer type: {type(smem_src)}")

    total_elems = 1
    for extent in shape:
        if not isinstance(extent, tirx.IntImm):
            raise ValueError("Packed scale-factor helpers require a static extent.")
        total_elems *= extent.value
    if total_elems % chunk_elems != 0:
        raise ValueError(f"Packed scale-factor helpers require total extent to be a multiple of {chunk_elems}, got {total_elems}.")
    return total_elems // chunk_elems


def tcgen05_cp_warpx4(smem_src, tmem_dst, tmem_col_offset=0, *, use_2cta: bool = False):
    """Copy one or more packed scale-factor chunks from shared memory to tensor memory.

    The helper lowers to one or more ``tcgen05.cp.cta_group::{1,2}.32x128b.warpx4``
    instructions. For 1D packed ``uint32`` scale buffers, each 128-word chunk maps to
    4 TMEM columns and the column offset is advanced automatically.
    """
    num_chunks = _tcgen05_num_smem_chunks(smem_src, 128)
    if isinstance(tmem_dst, tirx.Buffer):
        tmem_ptr = tmem_dst.data
    elif isinstance(tmem_dst, (BufferLoad, BufferRegion)):
        tmem_ptr = tmem_dst.buffer.data
    else:
        tmem_ptr = tmem_dst
    ann = {"use_2cta": 1} if use_2cta else None
    buffer, base_offset = retrieve_buffer_and_offset(smem_src)

    @macro
    def _tcgen05_cp_warpx4_chunked(buffer, tmem_ptr, tmem_col_offset, base_offset):
        for i in T.unroll(num_chunks):
            chunk_ptr = buffer.access_ptr("r", offset=base_offset + i * 128)
            tirx.call_intrin(
                "void",
                tirx.op.Op.get("tl.ptx_tcgen05_cp_warpx4"),
                chunk_ptr,
                tmem_ptr,
                tmem_col_offset + i * 4,
                annotations=ann,
            )

    return _tcgen05_cp_warpx4_chunked(buffer, tmem_ptr, tmem_col_offset, base_offset)


def tcgen05_sf_warp_transpose(smem_src):
    """Warp-level transpose for one or more packed scale-factor chunks in shared memory.

    For 1D packed ``uint32`` scale buffers, the helper automatically applies the
    transpose to each 128-word chunk in order.
    """
    num_chunks = _tcgen05_num_smem_chunks(smem_src, 128)

    buffer, base_offset = retrieve_buffer_and_offset(smem_src)

    @macro
    def _tcgen05_sf_warp_transpose_chunked(buffer, base_offset):
        for i in T.unroll(num_chunks):
            chunk_ptr = buffer.access_ptr("rw", offset=base_offset + i * 128)
            tirx.call_intrin("void", tirx.op.Op.get("tl.ptx_tcgen05_sf_warp_transpose"), chunk_ptr)

    return _tcgen05_sf_warp_transpose_chunked(buffer, base_offset)


def ptx_mma_sm70(
    shape,
    A_layout,
    B_layout,
    A_dtype,
    B_dtype,
    C_dtype,
    multiplicand_a,
    a_index,
    multiplicand_b,
    b_index,
    accumulator,
    c_index,
):
    """TVM intrinsic for ptx tensor core mma instructions on SM70 (Volta).

    This intrinsic provides SM70-specific MMA operations that support m16n16k4 shape
    with FP16 inputs and FP16/FP32 accumulation.

    Parameters
    ----------

    shape : str
        The shape of mma fragment (e.g., "m16n16k4").

    A_layout : str
        The layout of multiplicand fragment A ("row" or "col").

    B_layout : str
        The layout of multiplicand fragment B ("row" or "col").

    A_dtype : str
        The data type of multiplicand fragment A (typically "fp16").

    B_dtype : str
        The data type of multiplicand fragment B (typically "fp16").

    C_dtype : str
        The data type of accumulator fragment C ("fp16" or "fp32").

    multiplicand_a : Var
        The multiplicand fragment A variable.

    a_index : Expr
        The index of multiplicand fragment A.

    multiplicand_b : Var
        The multiplicand fragment B variable.

    b_index : Expr
        The index of multiplicand fragment B.

    accumulator : Var
        The accumulator fragment C variable.

    c_index : Expr
        The index of accumulator fragment C.

    Returns
    -------
    call : PrimExpr
        The call expression.

    Examples
    --------
    >>> T.ptx_mma_sm70(
    ...     "float16",
    ...     "m16n16k4",
    ...     "row",
    ...     "col",
    ...     "fp16",
    ...     "fp16",
    ...     "fp16",
    ...     A_local.data,
    ...     0,
    ...     B_local.data,
    ...     0,
    ...     C_local.data,
    ...     0,
    ... )
    """
    return tirx.call_intrin(
        "handle",
        tirx.op.Op.get("tl.ptx_mma_sm70"),
        shape,
        A_layout,
        B_layout,
        A_dtype,
        B_dtype,
        C_dtype,
        multiplicand_a,
        a_index,
        multiplicand_b,
        b_index,
        accumulator,
        c_index,
    )


def ds_read_tr16_b64(src: BufferLikeType) -> PrimExpr:
    """LDS transpose read, 64-bit, 16-element transpose (gfx950 only).

    Reads 8 bytes from LDS (__shared__ memory) with a 16-element transpose.
    Used for FP16/BF16 MFMA matrix B-loads on MI350/MI355X (gfx950).

    Args:
        src: A `Buffer`, `BufferRegion`, or `BufferLoad` in shared memory.

    Returns:
        PrimExpr: The loaded 64-bit value as uint32x2.

    Example:
        >>> val = T.ds_read_tr16_b64(smem[i])
    """
    if not isinstance(src, BufferLikeTypeTuple):
        raise TypeError(f"T.ds_read_tr16_b64 expects Buffer, BufferRegion, or BufferLoad. Got {type(src)}: {src}")
    ptr = retrieve_ptr(src, access_type="r")
    return tirx.call_intrin("uint32x2", tirx.op.Op.get("tl.ds_read_tr16_b64"), ptr)


def ds_read_tr8_b64(src: BufferLikeType) -> PrimExpr:
    """LDS transpose read, 64-bit, 8-element transpose (gfx950 only).

    Reads 8 bytes from LDS (__shared__ memory) with an 8-element transpose.
    Used for FP32 MFMA matrix B-loads on MI350/MI355X (gfx950).

    Args:
        src: A `Buffer`, `BufferRegion`, or `BufferLoad` in shared memory.

    Returns:
        PrimExpr: The loaded 64-bit value as uint32x2.

    Example:
        >>> val = T.ds_read_tr8_b64(smem[i])
    """
    if not isinstance(src, BufferLikeTypeTuple):
        raise TypeError(f"T.ds_read_tr8_b64 expects Buffer, BufferRegion, or BufferLoad. Got {type(src)}: {src}")
    ptr = retrieve_ptr(src, access_type="r")
    return tirx.call_intrin("uint32x2", tirx.op.Op.get("tl.ds_read_tr8_b64"), ptr)


def ldg32(src: BufferLikeType, pred: PrimExpr = None) -> PrimExpr:
    """Load 32 bits (4 bytes) from global memory using explicit PTX instructions.

    Usage: `T.ldg32(x[i])` or `T.ldg32(x[i:i+2])` emits `tl::ldg32(ptr)`.

    Args:
        src: A `Buffer`, `BufferRegion`, or `BufferLoad`.
        pred: Optional predicate condition. If False, the load is skipped.

    Returns:
        PrimExpr: The loaded 32-bit value.

    Example:
        >>> val = T.ldg32(x[i])
        >>> val = T.ldg32(x[i:i+2])  # load 2 x fp16
        >>> val = T.ldg32(x[i], pred=i < N)  # predicated load
    """
    if not isinstance(src, BufferLikeTypeTuple):
        raise TypeError(f"T.ldg32 expects Buffer, BufferRegion, or BufferLoad. Got {type(src)}: {src}")
    ptr = retrieve_ptr(src, access_type="r")
    if pred is None:
        return tirx.call_intrin("uint32", tirx.op.Op.get("tl.ldg32"), ptr)
    else:
        return tirx.call_intrin("uint32", tirx.op.Op.get("tl.ldg32"), ptr, pred)


def ldg64(src: BufferLikeType, pred: PrimExpr = None) -> PrimExpr:
    """Load 64 bits (8 bytes) from global memory using explicit PTX instructions.

    Usage: `T.ldg64(x[i])` or `T.ldg64(x[i:i+4])` emits `tl::ldg64(ptr)`.

    Args:
        src: A `Buffer`, `BufferRegion`, or `BufferLoad`.
        pred: Optional predicate condition. If False, the load is skipped.

    Returns:
        PrimExpr: The loaded 64-bit value.

    Example:
        >>> val = T.ldg64(x[i])
        >>> val = T.ldg64(x[i:i+4])  # load 4 x fp16
        >>> val = T.ldg64(x[i], pred=i < N)  # predicated load
    """
    if not isinstance(src, BufferLikeTypeTuple):
        raise TypeError(f"T.ldg64 expects Buffer, BufferRegion, or BufferLoad. Got {type(src)}: {src}")
    ptr = retrieve_ptr(src, access_type="r")
    if pred is None:
        return tirx.call_intrin("uint32x2", tirx.op.Op.get("tl.ldg64"), ptr)
    else:
        return tirx.call_intrin("uint32x2", tirx.op.Op.get("tl.ldg64"), ptr, pred)


def ldg128(src: BufferLikeType, pred: PrimExpr = None) -> PrimExpr:
    """Load 128 bits (16 bytes) from global memory using explicit PTX instructions.

    Usage: `T.ldg128(x[i])` or `T.ldg128(x[i:i+8])` emits `tl::ldg128(ptr)`.

    Args:
        src: A `Buffer`, `BufferRegion`, or `BufferLoad`.
        pred: Optional predicate condition. If False, the load is skipped.

    Returns:
        PrimExpr: The loaded 128-bit value.

    Example:
        >>> val = T.ldg128(x[i])
        >>> val = T.ldg128(x[i:i+8])  # load 8 x fp16
        >>> val = T.ldg128(x[i], pred=i < N)  # predicated load
    """
    if not isinstance(src, BufferLikeTypeTuple):
        raise TypeError(f"T.ldg128 expects Buffer, BufferRegion, or BufferLoad. Got {type(src)}: {src}")
    ptr = retrieve_ptr(src, access_type="r")
    if pred is None:
        return tirx.call_intrin("uint32x4", tirx.op.Op.get("tl.ldg128"), ptr)
    else:
        return tirx.call_intrin("uint32x4", tirx.op.Op.get("tl.ldg128"), ptr, pred)


def lds32(src: BufferLikeType) -> PrimExpr:
    """Load 32 bits (4 bytes) from shared memory and return a uint32."""
    if not isinstance(src, BufferLikeTypeTuple):
        raise TypeError(f"T.lds32 expects Buffer, BufferRegion, or BufferLoad. Got {type(src)}: {src}")
    ptr = retrieve_ptr(src, access_type="r")
    return tirx.call_intrin("uint32", tirx.op.Op.get("tl.lds32"), ptr)


def lds64(src: BufferLikeType) -> PrimExpr:
    """Load 64 bits (8 bytes) from shared memory and return a uint32x2."""
    if not isinstance(src, BufferLikeTypeTuple):
        raise TypeError(f"T.lds64 expects Buffer, BufferRegion, or BufferLoad. Got {type(src)}: {src}")
    ptr = retrieve_ptr(src, access_type="r")
    return tirx.call_intrin("uint32x2", tirx.op.Op.get("tl.lds64"), ptr)


def lds128(src: BufferLikeType) -> PrimExpr:
    """Load 128 bits (16 bytes) from shared memory and return a uint32x4."""
    if not isinstance(src, BufferLikeTypeTuple):
        raise TypeError(f"T.lds128 expects Buffer, BufferRegion, or BufferLoad. Got {type(src)}: {src}")
    ptr = retrieve_ptr(src, access_type="r")
    return tirx.call_intrin("uint32x4", tirx.op.Op.get("tl.lds128"), ptr)


def ldg256(src: BufferLikeType, pred: PrimExpr = None) -> PrimExpr:
    """Load 256 bits (32 bytes) from global memory using explicit PTX instructions.

    Usage: `T.ldg256(x[i])` or `T.ldg256(x[i:i+16])` emits `tl::ldg256(ptr)`.

    Args:
        src: A `Buffer`, `BufferRegion`, or `BufferLoad`.
        pred: Optional predicate condition. If False, the load is skipped.

    Returns:
        PrimExpr: The loaded 256-bit value.

    Example:
        >>> val = T.ldg256(x[i])
        >>> val = T.ldg256(x[i:i+16])  # load 16 x fp16
        >>> val = T.ldg256(x[i], pred=i < N)  # predicated load
    """
    if not isinstance(src, BufferLikeTypeTuple):
        raise TypeError(f"T.ldg256 expects Buffer, BufferRegion, or BufferLoad. Got {type(src)}: {src}")
    ptr = retrieve_ptr(src, access_type="r")
    if pred is None:
        return tirx.call_intrin("uint32x8", tirx.op.Op.get("tl.ldg256"), ptr)
    else:
        return tirx.call_intrin("uint32x8", tirx.op.Op.get("tl.ldg256"), ptr, pred)


def stg32(dst: BufferLikeType, value: PrimExpr, pred: PrimExpr = None) -> None:
    """Store 32 bits (4 bytes) to global memory using explicit PTX instructions.

    Usage: `T.stg32(y[i], value)` emits `tl::stg32(ptr, value)`.

    Args:
        dst: A `Buffer`, `BufferRegion`, or `BufferLoad` indicating the destination.
        value: The 32-bit value to store.
        pred: Optional predicate condition. If False, the store is skipped.

    Example:
        >>> T.stg32(y[i], val)
        >>> T.stg32(y[i], val, pred=i < N)  # predicated store
    """
    if not isinstance(dst, BufferLikeTypeTuple):
        raise TypeError(f"T.stg32 expects Buffer, BufferRegion, or BufferLoad. Got {type(dst)}: {dst}")
    ptr = retrieve_ptr(dst, access_type="w")
    if pred is None:
        return tirx.call_intrin("handle", tirx.op.Op.get("tl.stg32"), ptr, value)
    else:
        return tirx.call_intrin("handle", tirx.op.Op.get("tl.stg32"), ptr, value, pred)


def stg64(dst: BufferLikeType, value: PrimExpr, pred: PrimExpr = None) -> None:
    """Store 64 bits (8 bytes) to global memory using explicit PTX instructions.

    Usage: `T.stg64(y[i:i+2], value)` emits `tl::stg64(ptr, value)`.

    Args:
        dst: A `Buffer`, `BufferRegion`, or `BufferLoad` indicating the destination.
        value: The 64-bit value to store (e.g., uint2).
        pred: Optional predicate condition. If False, the store is skipped.

    Example:
        >>> T.stg64(y[i:i+2], val)
        >>> T.stg64(y[i:i+2], val, pred=i < N)  # predicated store
    """
    if not isinstance(dst, BufferLikeTypeTuple):
        raise TypeError(f"T.stg64 expects Buffer, BufferRegion, or BufferLoad. Got {type(dst)}: {dst}")
    ptr = retrieve_ptr(dst, access_type="w")
    if pred is None:
        return tirx.call_intrin("handle", tirx.op.Op.get("tl.stg64"), ptr, value)
    else:
        return tirx.call_intrin("handle", tirx.op.Op.get("tl.stg64"), ptr, value, pred)


def stg128(dst: BufferLikeType, value: PrimExpr, pred: PrimExpr = None) -> None:
    """Store 128 bits (16 bytes) to global memory using explicit PTX instructions.

    Usage: `T.stg128(y[i:i+4], value)` emits `tl::stg128(ptr, value)`.

    Args:
        dst: A `Buffer`, `BufferRegion`, or `BufferLoad` indicating the destination.
        value: The 128-bit value to store (e.g., uint4).
        pred: Optional predicate condition. If False, the store is skipped.

    Example:
        >>> T.stg128(y[i:i+4], val)
        >>> T.stg128(y[i:i+4], val, pred=i < N)  # predicated store
    """
    if not isinstance(dst, BufferLikeTypeTuple):
        raise TypeError(f"T.stg128 expects Buffer, BufferRegion, or BufferLoad. Got {type(dst)}: {dst}")
    ptr = retrieve_ptr(dst, access_type="w")
    if pred is None:
        return tirx.call_intrin("handle", tirx.op.Op.get("tl.stg128"), ptr, value)
    else:
        return tirx.call_intrin("handle", tirx.op.Op.get("tl.stg128"), ptr, value, pred)


def sts32(dst: BufferLikeType, value: PrimExpr) -> None:
    """Store a uint32 value to shared memory as 32 bits (4 bytes)."""
    if not isinstance(dst, BufferLikeTypeTuple):
        raise TypeError(f"T.sts32 expects Buffer, BufferRegion, or BufferLoad. Got {type(dst)}: {dst}")
    ptr = retrieve_ptr(dst, access_type="w")
    return tirx.call_intrin("handle", tirx.op.Op.get("tl.sts32"), ptr, value)


def sts64(dst: BufferLikeType, value: PrimExpr) -> None:
    """Store a uint32x2 value to shared memory as 64 bits (8 bytes)."""
    if not isinstance(dst, BufferLikeTypeTuple):
        raise TypeError(f"T.sts64 expects Buffer, BufferRegion, or BufferLoad. Got {type(dst)}: {dst}")
    ptr = retrieve_ptr(dst, access_type="w")
    return tirx.call_intrin("handle", tirx.op.Op.get("tl.sts64"), ptr, value)


def sts128(dst: BufferLikeType, value: PrimExpr) -> None:
    """Store a uint32x4 value to shared memory as 128 bits (16 bytes)."""
    if not isinstance(dst, BufferLikeTypeTuple):
        raise TypeError(f"T.sts128 expects Buffer, BufferRegion, or BufferLoad. Got {type(dst)}: {dst}")
    ptr = retrieve_ptr(dst, access_type="w")
    return tirx.call_intrin("handle", tirx.op.Op.get("tl.sts128"), ptr, value)


def stg256(dst: BufferLikeType, value: PrimExpr, pred: PrimExpr = None) -> None:
    """Store 256 bits (32 bytes) to global memory using explicit PTX instructions.

    Usage: `T.stg256(y[i:i+8], value)` emits `tl::stg256(ptr, value)`.

    Args:
        dst: A `Buffer`, `BufferRegion`, or `BufferLoad` indicating the destination.
        value: The 256-bit value to store (e.g., ulonglong4).
        pred: Optional predicate condition. If False, the store is skipped.

    Example:
        >>> T.stg256(y[i:i+8], val)
        >>> T.stg256(y[i:i+8], val, pred=i < N)  # predicated store
    """
    if not isinstance(dst, BufferLikeTypeTuple):
        raise TypeError(f"T.stg256 expects Buffer, BufferRegion, or BufferLoad. Got {type(dst)}: {dst}")
    ptr = retrieve_ptr(dst, access_type="w")
    if pred is None:
        return tirx.call_intrin("handle", tirx.op.Op.get("tl.stg256"), ptr, value)
    else:
        return tirx.call_intrin("handle", tirx.op.Op.get("tl.stg256"), ptr, value, pred)
