# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
# pylint: disable=invalid-name, too-many-arguments
"""CUDA, PTX, and NVSHMEM TIR intrinsic builders."""

from __future__ import annotations

from tvm import tirx
from tvm.ir import Call, Op, is_prim_expr
from tvm.runtime import const
from tvm.tirx.op import bitwise_and, call_intrin, tvm_access_ptr
from tvm.tirx.operator.intrinsics._common import CLUSTER_BARRIER_SEM as _CLUSTER_BARRIER_SEM
from tvm.tirx.operator.intrinsics._common import (
    CP_ASYNC_BULK_CACHE_HINT as _CP_ASYNC_BULK_CACHE_HINT,
)
from tvm.tirx.operator.intrinsics._common import CP_ASYNC_BULK_RED_OP as _CP_ASYNC_BULK_RED_OP
from tvm.tirx.operator.intrinsics._common import CP_ASYNC_CACHE_HINT as _CP_ASYNC_CACHE_HINT
from tvm.tirx.operator.intrinsics._common import CP_ASYNC_FILL_MODE as _CP_ASYNC_FILL_MODE
from tvm.tirx.operator.intrinsics._common import CP_ASYNC_PREFETCH_SIZE as _CP_ASYNC_PREFETCH_SIZE
from tvm.tirx.operator.intrinsics._common import F32X2_ROUND as _F32X2_ROUND
from tvm.tirx.operator.intrinsics._common import FENCE_PROXY_ASYNC_SPACE as _FENCE_PROXY_ASYNC_SPACE
from tvm.tirx.operator.intrinsics._common import FENCE_SCOPE as _FENCE_SCOPE
from tvm.tirx.operator.intrinsics._common import FENCE_SEM as _FENCE_SEM
from tvm.tirx.operator.intrinsics._common import LDMATRIX_DTYPE as _LDMATRIX_DTYPE
from tvm.tirx.operator.intrinsics._common import LDMATRIX_NUM as _LDMATRIX_NUM
from tvm.tirx.operator.intrinsics._common import NVSHMEM_CMP as _NVSHMEM_CMP
from tvm.tirx.operator.intrinsics._common import NVSHMEM_SIG_OP as _NVSHMEM_SIG_OP
from tvm.tirx.operator.intrinsics._common import TCGEN05_CP_DECOMPRESS as _TCGEN05_CP_DECOMPRESS
from tvm.tirx.operator.intrinsics._common import TCGEN05_CP_MULTICAST as _TCGEN05_CP_MULTICAST
from tvm.tirx.operator.intrinsics._common import TCGEN05_CP_SHAPES as _TCGEN05_CP_SHAPES
from tvm.tirx.operator.intrinsics._common import TCGEN05_CTA_GROUP as _TCGEN05_CTA_GROUP
from tvm.tirx.operator.intrinsics._common import TCGEN05_LDST_SHAPES as _TCGEN05_LDST_SHAPES

tir = tirx

########################################################
# CUDA native builtins
########################################################


def cuda_func_call(func_name, *args, source_code, return_type="void"):
    """TVM intrinsic to call a CUDA function. Source code is provided as a string.

    Parameters
    ----------
    func_name: str
        The name of the CUDA function.

    args: Expr
        The arguments to the CUDA function.

    source_code: str
        The source code of the CUDA function.

    return_type: str
        The return type of the CUDA function.
    """
    return call_intrin(return_type, "tirx.cuda.func_call", func_name, *args, source_code)


def cuda_warp_reduce(value, op, width=32):
    """Warp-level butterfly shuffle-XOR reduction.

    Reduces ``value`` across ``width`` adjacent lanes using the specified
    operation.  Codegen emits ``log2(width)`` steps of
    ``__shfl_xor_sync(0xFFFFFFFF, val, mask)`` with descending XOR masks.

    Parameters
    ----------
    value : Expr
        The per-thread scalar value to reduce.

    op : str
        Reduction operation: ``"sum"``, ``"max"``, or ``"min"``.

    width : int
        Number of lanes participating in each reduction group.
        Must be a power of two in [2, 32].  Defaults to 32 (full warp).

    Returns
    -------
    call : Expr
        The reduced value (same dtype as *value*).
    """
    return call_intrin(value.ty, "tirx.cuda.warp_reduce", value, op, width)


def cuda_warp_sum(value, width=32):
    """Convenience wrapper: ``cuda_warp_reduce(value, "sum", width)``."""
    return cuda_warp_reduce(value, "sum", width)


def cuda_warp_max(value, width=32):
    """Convenience wrapper: ``cuda_warp_reduce(value, "max", width)``."""
    return cuda_warp_reduce(value, "max", width)


def cuda_warp_min(value, width=32):
    """Convenience wrapper: ``cuda_warp_reduce(value, "min", width)``."""
    return cuda_warp_reduce(value, "min", width)


def cuda_cta_reduce(value, op, num_warps, scratch):
    """CTA-wide reduction via warp shuffle + shared memory.

    Two-step reduction: (1) intra-warp shuffle reduction, (2) warp-0
    collects per-warp partials from ``scratch``, reduces, broadcasts via
    ``__syncthreads()``.  All CTA threads must participate.

    Parameters
    ----------
    value : Expr
        Per-thread scalar value to reduce.

    op : str
        Reduction operation: ``"sum"``, ``"max"``, or ``"min"``.

    num_warps : int
        Number of warps in the CTA.  Must be a power of two in [1, 32].

    scratch : Var
        Data pointer to shared-memory scratch space (>= num_warps elements).

    Returns
    -------
    call : Expr
        The reduced value broadcast to all threads (same dtype as *value*).
    """
    return call_intrin(value.ty, "tirx.cuda.cta_reduce", value, op, num_warps, scratch)


def cuda_cta_sum(value, num_warps, scratch):
    """Convenience wrapper: ``cuda_cta_reduce(value, "sum", num_warps, scratch)``."""
    return cuda_cta_reduce(value, "sum", num_warps, scratch)


def cuda_cta_max(value, num_warps, scratch):
    """Convenience wrapper: ``cuda_cta_reduce(value, "max", num_warps, scratch)``."""
    return cuda_cta_reduce(value, "max", num_warps, scratch)


def cuda_cta_min(value, num_warps, scratch):
    """Convenience wrapper: ``cuda_cta_reduce(value, "min", num_warps, scratch)``."""
    return cuda_cta_reduce(value, "min", num_warps, scratch)


def cuda_warp_sync():
    """TVM intrinsic to synchronize threads within the current warp.

    This lowers to a CUDA `__syncwarp()` call.

    Returns
    -------
    call : Expr
        The call expression.
    """
    return call_intrin("", "tirx.cuda.warp_sync")


def cuda_cta_sync():
    """TVM intrinsic to call CUDA syncthreads (block-wide barrier)

    Returns
    -------
    call : Expr
        The call expression.
    """
    return call_intrin("", "tirx.cuda.cta_sync")


def cuda_grid_sync():
    """TVM intrinsic to call CUDA grid-wide sync (cooperative groups)

    Returns
    -------
    call : Expr
        The call expression.
    """
    return call_intrin("", "tirx.cuda.grid_sync")


def cuda_cluster_sync():
    """TVM intrinsic to call CUDA cluster-wide barrier sync

    Returns
    -------
    call : Expr
        The call expression.
    """
    return call_intrin("", "tirx.cuda.cluster_sync")


def cuda_thread_rank():
    """TVM intrinsic that returns ``cooperative_groups::thread_rank()``
    for the enclosing CTA -- the linear thread index within the block.

    Useful for building "single thread of CTA" predicates without
    referencing user-declared scope_id vars. For example, the idiomatic
    mbarrier.init leader predicate is::

        T.cuda.thread_rank() == 0

    Returns
    -------
    call : Expr
        The call expression (``int32``).
    """
    return call_intrin("int32", "tirx.cuda.thread_rank")


def cuda_half2float(src):
    """TVM intrinsic to convert half to float

    Parameters
    ----------
    src : Expr
        Source pointer.

    Returns
    -------
    call : Expr
        The call expression.
    """
    return call_intrin("float32", "tirx.cuda.half2float", src)


def cuda_bfloat162float(src):
    """TVM intrinsic to convert bfloat16 to float

    Parameters
    ----------
    src : Expr
        Source pointer.

    Returns
    -------
    call : Expr
        The call expression.
    """
    return call_intrin("float32", "tirx.cuda.bfloat162float", src)


def cuda_float22half2(dst, src):
    """TVM intrinsic to convert float2 to half2 with rounding

    Parameters
    ----------
    dst : Expr
        Destination pointer.

    src : Expr
        Source pointer.

    Returns
    -------
    call : Expr
        The call expression.
    """
    return call_intrin("", "tirx.cuda.float22half2", dst, src)


def cuda_trap_when_assert_failed(cond):
    """TVM intrinsic to trap when assertion failed (cond == false)

    Parameters
    ----------
    cond : Expr
        Condition to check.

    Returns
    -------
    call : Expr
        The call expression.
    """
    return call_intrin("", "tirx.cuda.trap_when_assert_failed", cond)


def cuda_runtime_instr_desc(desc, sf_id):
    """TVM intrinsic to update runtime instruction descriptor

    Parameters
    ----------
    desc : Expr
        Pointer to the descriptor (uint32*).

    sf_id : Expr
        The subfragment id.

    Returns
    -------
    call : Expr
        The call expression.
    """
    return call_intrin("", "tirx.cuda.runtime_instr_desc", desc, sf_id)


def cuda_half8tofloat8(src_addr, dst_addr):
    """TVM intrinsic to convert 8 half2s to 8 float2s

    Parameters
    ----------
    src_addr : Expr
        Source pointer.

    dst_addr : Expr
        Destination pointer.

    Returns
    -------
    call : Expr
        The call expression.
    """
    return call_intrin("", "tirx.cuda.half8tofloat8", src_addr, dst_addr)


def cuda_float8tohalf8(src_addr, dst_addr):
    """TVM intrinsic to convert 8 float2s to 8 half2s

    Parameters
    ----------
    src_addr : Expr
        Source pointer.

    dst_addr : Expr
        Destination pointer.

    Returns
    -------
    call : Expr
        The call expression.
    """
    return call_intrin("", "tirx.cuda.float8tohalf8", src_addr, dst_addr)


def ptx_mma_sp(
    dtype,
    shape,
    A_layout,
    B_layout,
    A_dtype,
    B_dtype,
    C_dtype,
    multiplicand_a,
    a_index,
    multiplicand_b,
    b_index,
    accumulator,
    c_index,
    metadata,
    meta_index,
    sparse_selector,
    saturate,
):
    """TVM intrinsic for sparse tensor core ptx instructions
    https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-instructions-for-sparse-mma

    Parameters
    ----------
    dtype : str
        The data type of the result.

    shape : str
        The shape of mma fragment.

    A_layout : Literal["row", "col"]
        The layout of multiplicand fragment A.

    B_layout : Literal["row", "col"]
        The layout of multiplicand fragment B.

    A_dtype : str
        The data type of multiplicand fragment A.

    B_dtype : str
        The data type of multiplicand fragment B.

    C_dtype : str
        The data type of multiplicand fragment C.

    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.

    metadata : Expr
        The metadata of operand.

    meta_index : Expr
        The metadata index of operand.

    sparse_selector : Expr
        The sparse selector indicating the thread that stores the metadata.

    saturate : bool
        The optional saturation at the output.

    Returns
    -------
    call : Expr
        The call expression.
    """
    return call_intrin(
        dtype,
        "tirx.ptx.mma_sp",
        shape,
        A_layout,
        B_layout,
        A_dtype,
        B_dtype,
        C_dtype,
        multiplicand_a,
        a_index,
        multiplicand_b,
        b_index,
        accumulator,
        c_index,
        metadata,
        meta_index,
        sparse_selector,
        saturate,
    )


def ptx_cp_async_bulk(
    dtype, shared_ptr, shared_offset, global_ptr, global_offset, bytes, barrier_id
):
    """TVM intrinsic for ptx async copy from global to shared memory using cp.async.bulk
    https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk

    Parameters
    ----------
    dtype : str
       The data type of the result.

    shared_ptr : Var
        The shared memory pointer variable.

    shared_offset : Expr
        The offset of shared memory pointer.

    global_ptr : Var
        The global memory pointer variable.

    global_offset : Expr
        The offset of global memory pointer.

    bytes : int
        The data size to copy.

    barrier_id : int
        The ID of the barrier shared memory pointer.

    Returns
    -------
    call : Expr
        The call expression.
    """
    return call_intrin(
        dtype,
        "tirx.ptx.cp_async_bulk",
        shared_ptr,
        shared_offset,
        global_ptr,
        global_offset,
        bytes,
        barrier_id,
    )


def ptx_cp_async_bulk_shared_to_cluster(dst_ptr, src_ptr, size, mbar):
    """PTX cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes

    Asynchronous bulk copy from executing CTA's shared memory to a remote
    CTA's shared memory within the same cluster.

    Parameters
    ----------
    dst_ptr : Expr
        Destination pointer in shared::cluster address space (remote CTA).

    src_ptr : Expr
        Source pointer in shared::cta address space (local CTA).

    size : Expr
        Number of bytes to copy (must be multiple of 16).

    mbar : Expr
        Mbarrier address in shared::cluster space for completion signaling,
        usually produced by ``T.ptx.map_shared_rank``.

    Returns
    -------
    call : Expr
        The call expression.
    """
    return call_intrin("", "tirx.ptx.cp_async_bulk_shared_to_cluster", dst_ptr, src_ptr, size, mbar)


def ptx_cp_async_mbarrier_arrive(barrier_id):
    """TVM intrinsic for ptx async copy barrier using cp.async.mbarrier.arrive
    https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-cp-async-mbarrier-arrive

    Parameters
    ----------
    barrier_id : int
        The ID of the barrier shared memory pointer.

    Returns
    -------
    call : Expr
        The call expression.
    """
    return call_intrin("", "tirx.ptx.cp_async_mbarrier_arrive", barrier_id)


def ptx_fence(sem: str, scope: str):
    """TVM intrinsic for PTX fence instruction.

    Generates: fence.{sem}.{scope};

    Parameters
    ----------
    sem : str
        The semantics of the fence. One of "sc", "acq_rel".
    scope : str
        The scope of the fence. One of "cta", "cluster", "gpu", "sys".

    Returns
    -------
    call : Expr
        The call expression.
    """
    _choice("sem", sem, _FENCE_SEM)
    _choice("scope", scope, _FENCE_SCOPE)
    return call_intrin("", "tirx.ptx.fence", sem, scope)


def ptx_fence_proxy_async(space: str = ""):
    """TVM intrinsic for PTX fence.proxy.async instruction.

    Generates: fence.proxy.async[.{space}];

    Parameters
    ----------
    space : str
        The address space qualifier. One of "", "global", "shared::cta", "shared::cluster".
        Empty string means no qualifier.

    Returns
    -------
    call : Expr
        The call expression.
    """
    _choice("space", space, _FENCE_PROXY_ASYNC_SPACE)
    return call_intrin("", "tirx.ptx.fence_proxy_async", space)


def ptx_mbarrier_init(bar, thread_count):
    """TVM intrinsic to call mbarrier.init.shared::cta.b64

    Parameters
    ----------
    bar : Var
        The pointer to barrier variable.

    thread_count : int
        The number of threads expected to arrive at the barrier.

    Returns
    -------
    call : Expr
        The call expression.
    """
    return call_intrin("", "tirx.ptx.mbarrier_init", bar, thread_count)


def ptx_mbarrier_arrive(bar, cta_id=None, pred=None, count=None):
    """TVM intrinsic to call
        mbarrier.arrive.shared::cta.b64
    or
        @p mapa.shared::cluster.u32
        @p mbarrier.arrive.shared::cluster.b64 [, count]

    Parameters
    ----------
    bar : Var
        The pointer to barrier variable.

    cta_id : Optional[Expr]
        The cta id.

    pred : Optional[Expr]
        The predicate to guard the operation.

    count : Optional[Expr]
        Explicit arrival count operand for the cross-CTA (cluster) form. When
        ``None`` the implicit count-of-1 form is emitted; when given, emits
        ``mbarrier.arrive.shared::cluster.b64 _, [addr], count``.
    """
    if cta_id is None and pred is None:
        return call_intrin("", "tirx.ptx.mbarrier_arrive", bar)
    assert cta_id is not None and pred is not None
    if count is None:
        return call_intrin("", "tirx.ptx.mbarrier_arrive", bar, cta_id, pred)
    return call_intrin("", "tirx.ptx.mbarrier_arrive", bar, cta_id, pred, count)


def ptx_mbarrier_arrive_cluster_count(bar, cta_id, count):
    """Cross-CTA ``mbarrier.arrive`` on CTA ``cta_id`` with an explicit count.

    Convenience for an already-elected thread: emits
    ``@p mapa.shared::cluster.u32`` + ``@p mbarrier.arrive.shared::cluster.b64 _,
    [addr], count`` with the guard defaulted to 1.
    """
    return call_intrin("", "tirx.ptx.mbarrier_arrive", bar, cta_id, True, count)


def ptx_mbarrier_arrive_expect_tx(bar, byte_count, cta_id=None, pred=None):
    """TVM intrinsic to call
        mbarrier.arrive_expect_tx.shared::cta.b64
    or
        @p mapa.shared::cluster.u32
        @p mbarrier.arrive_expect_tx.shared::cluster.b64

    Parameters
    ----------
    bar : Var
        The pointer to barrier variable.

    byte_count : int
        Increases the tx count of the mbarrier object to track completion of
        addtional async transactions.

    cta_id : Optional[Expr]
        The cta id.

    pred : Optional[Expr]
        The predicate to guard the operation.

    Returns
    -------
    call : Expr
        The call expression.
    """
    if cta_id is None and pred is None:
        return call_intrin("", "tirx.ptx.mbarrier_arrive_expect_tx", bar, byte_count)
    assert cta_id is not None
    # Cross-CTA expect_tx from an already-elected thread: default the guard to 1
    # (the caller has elected a single lane), so callers can pass cta_id alone.
    if pred is None:
        pred = True
    return call_intrin("", "tirx.ptx.mbarrier_arrive_expect_tx", bar, byte_count, cta_id, pred)


def ptx_mbarrier_try_wait(bar, phase):
    """TVM intrinsic to call mbarrier.try_wait.parity repeatedly until it returns true

    Parameters
    ----------
    bar : Var
        The pointer to barrier variable.

    phase : int
        The phase of the barrier.

    Returns
    -------
    call : Expr
        The call expression.
    """
    return call_intrin("", "tirx.ptx.mbarrier_try_wait", bar, phase)


def ptx_mbarrier_try_wait_acquire_cluster(bar, phase):
    """``mbarrier.try_wait.parity.acquire.cluster`` retry loop.

    Cluster-scope acquire wait — used to wait on a barrier that a remote CTA in
    the cluster arrives on (a group cluster wait).

    Parameters
    ----------
    bar : Var
        The pointer to barrier variable.

    phase : int
        The phase of the barrier.
    """
    return call_intrin("", "tirx.ptx.mbarrier_try_wait_acquire_cluster", bar, phase)


def ptx_mbarrier_try_wait_once(bar, phase, ticks):
    """TVM intrinsic for one-shot non-blocking ``mbarrier.try_wait.parity``.

    Returns ``1`` if the requested parity has been reached and ``0`` otherwise.
    This is intended for bounded debug waits; production waits should use
    :func:`ptx_mbarrier_try_wait`.
    """
    return call_intrin("uint32", "tirx.ptx.mbarrier_try_wait_once", bar, phase, ticks)


def ptx_bar_arrive(name_bar_id, thread_count):
    """TVM intrinsic to call bar.arrive a, b

    Parameters
    ----------
    name_bar_id : int
        The ID of the named barrier.

    thread_count : int
        The number of threads expected to arrive at the barrier.

    Returns
    -------
    call : Expr
        The call expression.
    """
    return call_intrin("", "tirx.ptx.bar_arrive", name_bar_id, thread_count)


def ptx_bar_sync(name_bar_id, thread_count):
    """TVM intrinsic to call bar.sync a, {b}

    Parameters
    ----------
    name_bar_id : int
        The ID of the named barrier.

    thread_count : int
        The number of threads expected to arrive at the barrier.

    Returns
    -------
    call : Expr
        The call expression.
    """
    return call_intrin("", "tirx.ptx.bar_sync", name_bar_id, thread_count)


def ptx_cp_async(
    dst_ptr,
    src_ptr,
    cp_size,
    *,
    cache_hint="",
    cache_policy=None,
    prefetch_size=-1,
    predicate=-1,
    fill_mode="",
):
    """TVM intrinsic for ptx async copy from global to shared memory using cp.async
    https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async

    Dispatches to one of three PTX-form-aligned ops:

    * ``ptx_cp_async_src_size`` for ``fill_mode == "zero"`` (zero-fill via
      ``src_size = pred ? cp_size : 0``).
    * ``ptx_cp_async_ignore_src`` for a non-empty ``predicate`` with no
      fill_mode (``setp+@p`` guards the asm).
    * ``ptx_cp_async_plain`` for the no-predicate / no-fill_mode case.

    Parameters
    ----------
    shared_ptr : Expr
        The pointer to the shared memory.

    global_ptr : Expr
        The pointer to the global memory.

    cp_size : int
        The data size to copy.

    cache_hint : str["evict_last", "evict_first", "evict_normal", ""]
        The cache hint.

    prefetch_size : int[-1, 64, 128, 256]
        The prefetch size.

    predicate : Expr
        The predicate to guard the operation.

    fill_mode : str["zero", ""]
        The fill mode.

    Returns
    -------
    call : Expr
        The call expression.
    """
    cache_policy, has_cache_policy = _resolve_cache_policy(cache_hint, cache_policy)
    _choice("prefetch_size", prefetch_size, _CP_ASYNC_PREFETCH_SIZE)
    _choice("fill_mode", fill_mode, _CP_ASYNC_FILL_MODE)
    return call_intrin(
        "",
        "tirx.ptx.cp_async",
        dst_ptr,
        src_ptr,
        cp_size,
        cache_policy,
        int(has_cache_policy),
        prefetch_size,
        predicate,
        fill_mode,
    )


def ptx_cp_async_legacy(*all_args):
    """Legacy ``ptx_cp_async`` API taking explicit src/dst offsets.

    Signature: ``(dst_ptr, dst_offset, src_ptr, src_offset, cp_size)``.
    Offsets are folded into the pointers via ``tvm_access_ptr`` then
    dispatched to fork-native :func:`ptx_cp_async`.

    ``T.ptx.cp_async_legacy`` runs through ``_dtype_forward`` which
    prepends a ``dtype=`` kwarg as a leading positional. The dtype names
    the *element* type of the buffer (offsets are in elements of that
    dtype, not bytes), so this function accepts either 5 or 6 positional
    args.
    """
    args = list(all_args)
    elem_dtype = "int8"
    if len(args) == 6:
        # Leading positional is the buffer element dtype, used to scale
        # offsets correctly when folding via ``tvm_access_ptr``.
        elem_dtype = args.pop(0)
    if len(args) != 5:
        raise ValueError(
            f"ptx_cp_async_legacy expects 5 args (or 6 with dtype= kwarg "
            f"prepended); got {len(all_args)}"
        )
    dst_ptr, dst_offset, src_ptr, src_offset, cp_size = args
    dst_ptr = _wrap_or_fold_access_ptr(dst_ptr, dst_offset, elem_dtype)
    src_ptr = _wrap_or_fold_access_ptr(src_ptr, src_offset, elem_dtype)
    return ptx_cp_async(dst_ptr, src_ptr, cp_size)


def ptx_cp_async_commit_group():
    """TVM intrinsic for ptx async copy commit
    https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-commit-group

    Returns
    -------
    call : Expr
        The call expression.
    """
    return call_intrin("", "tirx.ptx.cp_async_commit_group")


def ptx_cp_async_wait_group(num=0):
    """TVM intrinsic for ptx async copy wait
    https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-wait-group

    Parameters
    ----------
    num : int, optional
        The number of the most recent uncommitted pending cp.async groups to wait.

    Returns
    -------
    call : Expr
        The call expression.
    """
    return call_intrin("", "tirx.ptx.cp_async_wait_group", num)


def ptx_cp_async_bulk_tensor_global_to_cluster(
    dim, dst_ptr, bar, tensormap_addr, cta_mask, cta_group, cache_hint, *coords, cache_policy=None
):
    """TVM intrinsic to call cp.async.bulk.tensor.dim.shared::cluster.global.tile.mbarrier::complete_tx::bytes

    Parameters
    ----------
    dim : int
        The dimension of the source tensor.

    dst_ptr : Expr
        The destination pointer to the shared memory.

    bar : Expr
        The pointer to mbarrier variable.

    tensormap_addr : Expr
        The generic address of the tensor map object.

    cta_mask : int
        The mask of the cta for multicast.

    cta_group : int
        Must be either 1 or 2.
        If set to 1, mbarrier must be in the shared memory of the same CTA
        as the shared memory destination.  If set to 2, mbarrier can be in
        shared memory of either the same CTA as the shared memory destination
        or the shared memory of the peer CTA.

    cache_hint : str
        The cache hint.

    coords : List[Expr]
        specifies the starting coordinates in the tensor data in the global memory

    Returns
    -------
    call : Expr
        The call expression.
    """  # noqa: E501
    _choice("cta_group", cta_group, _TCGEN05_CTA_GROUP)
    if is_prim_expr(cache_hint):
        has_cache_policy, *coords = coords
        return call_intrin(
            "",
            "tirx.ptx.cp_async_bulk_tensor_global_to_cluster",
            dim,
            dst_ptr,
            bar,
            tensormap_addr,
            cta_mask,
            cta_group,
            cache_hint,
            has_cache_policy,
            *coords,
        )
    cache_policy, has_cache_policy = _resolve_cache_policy(cache_hint, cache_policy)
    return call_intrin(
        "",
        "tirx.ptx.cp_async_bulk_tensor_global_to_cluster",
        dim,
        dst_ptr,
        bar,
        tensormap_addr,
        cta_mask,
        cta_group,
        cache_policy,
        int(has_cache_policy),
        *coords,
    )


def ptx_cp_async_bulk_tensor_tile_gather4_global_to_cluster(
    dim, dst_ptr, bar, tensormap_addr, cta_mask, cta_group, cache_hint, *coords, cache_policy=None
):
    """TVM intrinsic to call
    cp.async.bulk.tensor.dim.shared::cluster.global.tile::gather4.mbarrier::complete_tx::bytes

    Parameters
    ----------
    dim : int
        The dimension of the source tensor.

    dst_ptr : Expr
        The destination pointer to the shared memory.

    bar : Expr
        The pointer to mbarrier variable.

    tensormap_addr : Expr
        The generic address of the tensor map object.

    cta_mask : int
        The mask of the cta for multicast.

    cta_group : int
        Must be either 1 or 2.

    cache_hint : str
        The cache hint.

    coords : List[Expr]
        The TMA coordinates followed by the 4 gather row indices.

    Returns
    -------
    call : Expr
        The call expression.
    """
    _choice("cta_group", cta_group, _TCGEN05_CTA_GROUP)
    if is_prim_expr(cache_hint):
        has_cache_policy, *coords = coords
        return call_intrin(
            "",
            "tirx.ptx.cp_async_bulk_tensor_tile_gather4_global_to_cluster",
            dim,
            dst_ptr,
            bar,
            tensormap_addr,
            cta_mask,
            cta_group,
            cache_hint,
            has_cache_policy,
            *coords,
        )
    cache_policy, has_cache_policy = _resolve_cache_policy(cache_hint, cache_policy)
    return call_intrin(
        "",
        "tirx.ptx.cp_async_bulk_tensor_tile_gather4_global_to_cluster",
        dim,
        dst_ptr,
        bar,
        tensormap_addr,
        cta_mask,
        cta_group,
        cache_policy,
        int(has_cache_policy),
        *coords,
    )


def ptx_cp_async_bulk_tensor_shared_to_global(
    dim, src_ptr, tensormap_addr, cache_hint, *coords, cache_policy=None
):
    """TVM intrinsic to call cp.async.bulk.tensor.dim.global.shared::cta.tile.bulk_group

    Parameters
    ----------
    dim : int
        The dimension of the copy tensor.

    src_ptr : Expr
        The source pointer to the shared memory.

    tensormap_addr : Expr
        The generic address of the tensor map object.

    cache_hint : str
        The cache hint.

    coords : List[Expr]
        specifies the starting coordinates in the tensor data in the global memory

    Returns
    -------
    call : Expr
        The call expression.
    """
    if is_prim_expr(cache_hint):
        has_cache_policy, *coords = coords
        return call_intrin(
            "",
            "tirx.ptx.cp_async_bulk_tensor_shared_to_global",
            dim,
            src_ptr,
            tensormap_addr,
            cache_hint,
            has_cache_policy,
            *coords,
        )
    cache_policy, has_cache_policy = _resolve_cache_policy(cache_hint, cache_policy)
    return call_intrin(
        "",
        "tirx.ptx.cp_async_bulk_tensor_shared_to_global",
        dim,
        src_ptr,
        tensormap_addr,
        cache_policy,
        int(has_cache_policy),
        *coords,
    )


def ptx_cp_async_bulk_tensor_global_to_cluster_prefetch(
    dim, tensormap_addr, cache_hint, *coords, cache_policy=None
):
    """TVM intrinsic to call cp.async.bulk.prefetch.tensor.dim.L2.global.tile

    Parameters
    ----------
    dim : int
        The dimension of the source tensor.

    tensormap_addr : Expr
        The generic address of the tensor map object.

    cache_hint : str
        The cache hint.

    coords : List[Expr]
        specifies the starting coordinates in the tensor data in the global memory

    Returns
    -------
    call : Expr
        The call expression.
    """
    if is_prim_expr(cache_hint):
        has_cache_policy, *coords = coords
        return call_intrin(
            "",
            "tirx.ptx.cp_async_bulk_tensor_global_to_cluster_prefetch",
            dim,
            tensormap_addr,
            cache_hint,
            has_cache_policy,
            *coords,
        )
    cache_policy, has_cache_policy = _resolve_cache_policy(cache_hint, cache_policy)
    return call_intrin(
        "",
        "tirx.ptx.cp_async_bulk_tensor_global_to_cluster_prefetch",
        dim,
        tensormap_addr,
        cache_policy,
        int(has_cache_policy),
        *coords,
    )


def ptx_cp_async_bulk_tensor_shared_to_global_reduce(
    dim, src_ptr, tensormap_addr, cache_hint, red_op, *coords, cache_policy=None
):
    """TVM intrinsic to call cp.reduce.async.bulk.tensor.dim.dst.src.redOp

    Parameters
    ----------
    dim : int
        The dimension of the copy tensor.

    src_ptr : Expr
        The source pointer to the shared memory.

    tensormap_addr : Expr
        The generic address of the tensor map object.

    cache_hint: str
        The cache hint.

    red_op: str
        The reduction operator.

    coords: List[Expr]
        The coordinates of the tensor.

    Returns
    -------
    call : Expr
        The call expression.
    """
    if is_prim_expr(cache_hint):
        has_cache_policy = red_op
        red_op, *coords = coords
        _choice("red_op", red_op, _CP_ASYNC_BULK_RED_OP)
        return call_intrin(
            "",
            "tirx.ptx.cp_async_bulk_tensor_shared_to_global_reduce",
            dim,
            src_ptr,
            tensormap_addr,
            cache_hint,
            has_cache_policy,
            red_op,
            *coords,
        )
    cache_policy, has_cache_policy = _resolve_cache_policy(cache_hint, cache_policy)
    _choice("red_op", red_op, _CP_ASYNC_BULK_RED_OP)
    return call_intrin(
        "",
        "tirx.ptx.cp_async_bulk_tensor_shared_to_global_reduce",
        dim,
        src_ptr,
        tensormap_addr,
        cache_policy,
        int(has_cache_policy),
        red_op,
        *coords,
    )


def ptx_cp_async_bulk_commit_group():
    """TVM intrinsic to call cp.async.bulk.tensor.commit_group

    Returns
    -------
    call : Expr
        The call expression.
    """
    return call_intrin("", "tirx.ptx.cp_async_bulk_commit_group")


def ptx_cp_async_bulk_wait_group(n=0, read=True):
    """TVM intrinsic to call cp.async.bulk.tensor.wait_group

    Parameters
    ----------
    n : int
        The number of the most recent uncommitted pending cp.async groups to wait.

    read : bool
        Whether the wait is for read.

    Returns
    -------
    call : Expr
        The call expression.
    """
    return call_intrin("", "tirx.ptx.cp_async_bulk_wait_group", n, read)


def ptx_barrier_cluster_arrive(sem="", aligned=True):
    """TVM intrinsic to call barrier.cluster.arrive{.sem}{.aligned}

    Parameters
    ----------
    sem : str
        Either release or relaxed or empty string.

    aligned : bool
        Whether all threads in the warp must execute the same instruction.
    """
    _choice("sem", sem, _CLUSTER_BARRIER_SEM)
    return call_intrin("", "tirx.ptx.barrier_cluster_arrive", sem, aligned)


def ptx_barrier_cluster_wait(acquire=False, aligned=True):
    """TVM intrinsic to call barrier.cluster.wait{.acquire}{.aligned}

    Parameters
    ----------
    acquire : bool
        The memory synchronization

    aligned : bool
        Whether all threads in the warp must execute the same instruction.
    """
    return call_intrin("", "tirx.ptx.barrier_cluster_wait", acquire, aligned)


def ptx_clc_try_cancel(handle, mbar):
    """TVM intrinsic to call clusterlaunchcontrol.try_cancel.

    Async-requests cancelling the next cluster's launch (work-stealing): writes the
    16B response handle to smem and signals ``mbar`` (complete_tx, multicast to both
    cluster CTAs).

    Parameters
    ----------
    handle : Expr
        Pointer to the 16B (uint4) smem response handle.

    mbar : Expr
        Pointer to the mbarrier signalled when the handle lands.
    """
    return call_intrin("", "tirx.ptx.clc_try_cancel", handle, mbar)


def ptx_clc_query_cancel(handle):
    """TVM intrinsic to call clusterlaunchcontrol.query_cancel.

    Decodes the response handle written by :func:`ptx_clc_try_cancel`. Returns the
    cancelled cluster's first ``ctaid.x``, or ``0xFFFFFFFF`` when no work was stolen.

    Parameters
    ----------
    handle : Expr
        Pointer to the 16B (uint4) smem response handle.
    """
    return call_intrin("uint32", "tirx.ptx.clc_query_cancel", handle)


def ptx_elect_sync():
    """TVM intrinsic to call elect.sync"""
    return call_intrin("uint32", "tirx.ptx.elect_sync")


def ptx_fence_mbarrier_init():
    """TVM intrinsic for PTX fence.mbarrier_init.release.cluster instruction.

    Generates: fence.mbarrier_init.release.cluster;

    Returns
    -------
    call : Expr
        The call expression.
    """
    return call_intrin("", "tirx.ptx.fence_mbarrier_init")


def ptx_fetch_register(bits, reg_name):
    """TVM intrinsic to tvm instrinsics to fetch PTX pre-defined registers

    Parameters
    ----------
    bits : int
        The number of bits of the register.

    reg_name : str
        The name of the register.

    Returns
    -------
    call : Expr
        The call expression.
    """
    return call_intrin("int" + str(bits), "tirx.ptx.fetch_register", bits, reg_name)


def ptx_mma(
    shape,
    a_layout,
    b_layout,
    d_type,
    a_type,
    b_type,
    c_type,
    d_ptrs,
    a_ptrs,
    b_ptrs,
    c_ptrs=None,
    saturate=False,
    bit_op=None,
):
    """TVM intrinsic for ptx tensor core mma instructions.
    https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-instructions-for-mma

    Each per-thread register of every operand is addressed by its OWN pointer
    (one ``void*`` per b32/f32 register), so the register fragments need not be
    contiguous in the register file. ``d_ptrs`` / ``a_ptrs`` / ``b_ptrs`` /
    ``c_ptrs`` are lists of one pointer per 32-bit register (b32 for
    fp16/bf16/tf32/int8 multiplicands, f32/f64 for the accumulator), enumerated
    in the fixed PTX register order (see the gemm dispatch /
    ``tests/python/tirx-base/test_tir_ptx_mma.py``).

    Within one b32 register the packed elements (e.g. 2 fp16 along k_pack)
    must stay contiguous (stride 1); only the b32 registers themselves may be
    scattered.

    Parameters
    ----------
    shape : str
        The shape of mma fragment.

    a_layout : Literal["row", "col"]
        The layout of multiplicand fragment A.

    b_layout : Literal["row", "col"]
        The layout of multiplicand fragment B.

    d_type : str
        The data type of result fragment D.

    a_type : str
        The data type of multiplicand fragment A.

    b_type : str
        The data type of multiplicand fragment B.

    c_type : str
        The data type of accumulator fragment C.

    d_ptrs : List[Expr]
        One pointer per result-fragment D register, in PTX order.

    a_ptrs : List[Expr]
        One pointer per multiplicand-A register, in PTX order.

    b_ptrs : List[Expr]
        One pointer per multiplicand-B register, in PTX order.

    c_ptrs : Optional[List[Expr]]
        One pointer per accumulator-C register, in PTX order. ``None`` (the
        default) means the accumulator is not used (beta == 0): codegen feeds
        a literal 0 for each C slot.

    saturate : bool
        The optional saturation at the output.

    bit_op : Optional[Literal["xor", "and"]]
        The 1-bit operator (for the b1 subbyte form). ``None`` means unused.

    Returns
    -------
    call : Expr
        The call expression.
    """
    d_ptrs = list(d_ptrs)
    a_ptrs = list(a_ptrs)
    b_ptrs = list(b_ptrs)
    has_c = c_ptrs is not None
    c_ptrs = list(c_ptrs) if has_c else []

    # Encode group register counts as leading attrs so codegen can slice the
    # flat pointer tail. ``no_c_ptr`` mirrors the legacy IntImm(0) sentinel.
    no_c_ptr = not has_c
    # Flattened pointer list: D regs, A regs, B regs, then C regs (if any).
    ptrs = [*d_ptrs, *a_ptrs, *b_ptrs, *c_ptrs]

    base = [
        "",
        "tirx.ptx.mma",
        shape,
        a_layout,
        b_layout,
        d_type,
        a_type,
        b_type,
        c_type,
        len(d_ptrs),
        len(a_ptrs),
        len(b_ptrs),
        len(c_ptrs),
        no_c_ptr,
        *ptrs,
        saturate,
    ]
    if bit_op is None:
        return call_intrin(*base)
    return call_intrin(*base, bit_op)


def ptx_mma_legacy(*all_args, operator=None):
    """Legacy ``ptx_mma`` API.

    Signature: ``(shape, A_layout, B_layout, A_dtype, B_dtype, C_dtype,
    multiplicand_a, a_index, multiplicand_b, b_index, accumulator,
    c_index, saturate, operator=None)``. The accumulator is reused as
    both input and output (no separate ``d``/``c`` slot), unlike
    fork-native :func:`ptx_mma` which distinguishes them. Translation:

    * ``a_dtype, b_dtype, c_dtype`` → fork ``a_type, b_type, c_type``
      (and reuse ``c_dtype`` as fork ``d_type`` since the accumulator
      dtype is the output dtype here).
    * ``(a_ptr, a_offset)`` and ``(b_ptr, b_offset)`` → folded via
      :func:`tvm_access_ptr`.
    * ``(accumulator, c_index)`` → folded; passed for both ``d_ptr`` and
      ``c_ptr`` since the accumulator is reused as the output.

    ``T.ptx.mma.legacy`` runs through ``_dtype_forward`` which prepends a
    ``dtype=`` kwarg as a leading positional, so this function accepts
    either 13 or 14 positional args.
    """
    args = list(all_args)
    # ``T.ptx.mma.legacy(..., dtype="...")`` has the dtype prepended by
    # ``_dtype_forward``; strip it here.
    if len(args) in (14, 15):
        _ = args.pop(0)
    if len(args) == 14:
        # operator passed positionally as the trailing arg.
        operator = args.pop()
    if len(args) != 13:
        raise ValueError(
            f"ptx_mma_legacy expects 13-15 positional args (with optional "
            f"leading ``call_dtype`` from dtype= kwarg and optional trailing "
            f"``operator``); got {len(all_args)}"
        )
    (
        shape,
        a_layout,
        b_layout,
        a_dtype,
        b_dtype,
        c_dtype,
        a_ptr,
        a_offset,
        b_ptr,
        b_offset,
        acc_ptr,
        c_offset,
        saturate,
    ) = args
    # Emit tirx.ptx_mma_legacy directly with separate (ptr_var, offset)
    # pairs. codegen_cuda.cc uses C pointer arithmetic ``ptr + offset``
    # so element offsets stay element-accurate, and lower_warp_memory
    # rewrites the offset's group component to a thread-local index.
    call_args = [
        shape,
        a_layout,
        b_layout,
        a_dtype,
        b_dtype,
        c_dtype,
        a_ptr,
        a_offset,
        b_ptr,
        b_offset,
        acc_ptr,
        c_offset,
        saturate,
    ]
    if operator is not None:
        call_args.append(operator)
    return call_intrin("", "tirx.ptx.mma_legacy", *call_args)


def ptx_mma_sp_legacy(*all_args):
    """Legacy ``ptx_mma_sp`` API.

    Signature: ``(shape, A_layout, B_layout, A_dtype, B_dtype, C_dtype,
    multiplicand_a, a_index, multiplicand_b, b_index, accumulator,
    c_index, metadata, meta_index, sparse_selector, saturate)``.

    ``T.ptx.mma_sp.legacy`` runs through ``_dtype_forward`` which prepends
    a ``dtype=`` kwarg as a leading positional, so this function accepts
    either 16 or 17 positional args.
    """
    args = list(all_args)
    if len(args) == 17:
        _ = args.pop(0)
    if len(args) != 16:
        raise ValueError(
            f"ptx_mma_sp_legacy expects 16 args (or 17 with dtype= kwarg "
            f"prepended); got {len(all_args)}"
        )
    (
        shape,
        a_layout,
        b_layout,
        a_dtype,
        b_dtype,
        c_dtype,
        a_ptr,
        a_offset,
        b_ptr,
        b_offset,
        acc_ptr,
        c_offset,
        meta_ptr,
        meta_offset,
        sparse_selector,
        saturate,
    ) = args
    return ptx_mma_sp(
        c_dtype,
        shape,
        a_layout,
        b_layout,
        a_dtype,
        b_dtype,
        c_dtype,
        a_ptr,
        a_offset,
        b_ptr,
        b_offset,
        acc_ptr,
        c_offset,
        meta_ptr,
        meta_offset,
        sparse_selector,
        saturate,
    )


def mma_store(dtype, m, n, dst_ptr, src_ptr, src_offset, dst_stride):
    """Store the result of PTX MMA into a destination pointer."""

    return call_intrin(dtype, "tirx.mma_store", m, n, dst_ptr, src_ptr, src_offset, dst_stride)


def mma_store_legacy(dtype, m, n, dst_ptr, src_ptr, src_offset, dst_stride):
    """mma_store with apache-style pointer/offset semantics."""

    return call_intrin(
        dtype,
        "tirx.mma_store_legacy",
        m,
        n,
        dst_ptr,
        src_ptr,
        src_offset,
        dst_stride,
    )


def mma_fill(dtype, local_size, local_ptr, offset):
    """Zero-initialize an MMA accumulation register."""

    return call_intrin(dtype, "tirx.mma_fill", local_size, local_ptr, offset)


def mma_fill_legacy(dtype, local_size, local_ptr, offset):
    """mma_fill with apache-style pointer/offset semantics."""

    return call_intrin(dtype, "tirx.mma_fill_legacy", local_size, local_ptr, offset)


def ptx_ldmatrix(trans, num, dtype, smem_ptr, *dst_handles):
    """TVM intrinsic for ldmatrix.sync.aligned.m8n8.x{num}{.trans}.shared.{dtype}.

    Mirrors the PTX ISA destination form: each output register is a separate
    operand. Pass ``T.address_of(buf[idx])`` (or ``buf.ptr_to([idx])``) for
    each destination — the slots may be non-contiguous.

    Parameters
    ----------
    trans : bool
        Apply the ``.trans`` modifier.
    num : int
        One of 1, 2, 4 — number of m8n8 fragments.
    dtype : str
        ``"b16"`` (4 bytes per fragment register) or ``"b8"`` (2 bytes per).
    smem_ptr : Expr
        Generic pointer to source shared memory.
    *dst_handles : Expr
        N pointer-to-uint32 destinations, where
        ``N = num if dtype == "b16" else num // 2``.

    https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-instructions-ldmatrix
    """
    _choice("num", num, _LDMATRIX_NUM)
    _choice("dtype", dtype, _LDMATRIX_DTYPE)
    # _LDMATRIX_DTYPE entries carry leading dot (".b16" / ".b8").
    dtype_bare = dtype.lstrip(".") if isinstance(dtype, str) else dtype
    n_regs = int(num) if dtype_bare == "b16" else int(num) // 2
    if len(dst_handles) != n_regs:
        raise ValueError(
            f"ldmatrix .x{int(num)}.{dtype_bare} expects {n_regs} destination "
            f"handles, got {len(dst_handles)}"
        )
    return call_intrin("", "tirx.ptx.ldmatrix", trans, num, dtype, smem_ptr, *dst_handles)


_PTX_TO_NUMPY_DTYPE = {
    "fp16": "float16",
    "fp32": "float32",
    "fp64": "float64",
    "bf16": "bfloat16",
    "tf32": "float32",
    "s8": "int8",
    "u8": "uint8",
    "s32": "int32",
    "s4": "int4",
    "u4": "uint4",
    "b1": "int1",
    "b16": "uint16",
    "e4m3": "float8_e4m3fn",
    "e5m2": "float8_e5m2",
}


def _ptx_to_numpy_dtype(dtype_str):
    """Map a PTX-abbreviation or numpy dtype string to a numpy dtype string
    suitable for ``tvm_access_ptr`` (which scales the offset by the element
    bit width). Unknown strings pass through unchanged so a caller may also
    pass an already-numpy dtype."""
    s = dtype_str if isinstance(dtype_str, str) else str(dtype_str)
    return _PTX_TO_NUMPY_DTYPE.get(s, s)


def _wrap_or_fold_access_ptr(ptr, offset, elem_dtype):
    """Wrap ``ptr`` with ``tvm_access_ptr`` unless it already is one.

    Several s_tir tensor intrinsics already pass ``buffer.access_ptr(...)``
    (an ``tvm_access_ptr`` Call) for the pointer argument. Naively wrapping
    that again yields a nested ``tvm_access_ptr(... access_ptr(...) ...)``
    whose ``args[1]`` is a Call rather than a Var, which crashes the
    lowering rule (Downcast<Var> at intrin_rule.cc) and several s_tir
    passes that assume a raw buffer var. Detect that case and fold the
    outer offset into the inner one.
    """

    is_access_ptr_call = (
        isinstance(ptr, Call) and isinstance(ptr.op, Op) and ptr.op.name == "tirx.tvm_access_ptr"
    )
    if is_access_ptr_call:
        # Inner Call already wraps the buffer var. Reuse its inner var and
        # inner element dtype (the marker type_annotation), and add the
        # outer offset (which is in `elem_dtype` units, same convention as
        # the inner since both come from the same buffer).
        inner_args = ptr.args
        inner_marker = inner_args[0]
        inner_var = inner_args[1]
        inner_offset = inner_args[2]
        rw_mask = inner_args[4]
        return call_intrin(
            ptr.ty,
            "tirx.tvm_access_ptr",
            inner_marker,
            inner_var,
            inner_offset + offset,
            1,
            rw_mask,
        )
    return tvm_access_ptr(elem_dtype, ptr, offset, 1, 1)


def ptx_ldmatrix_legacy(*all_args):
    """Legacy ``ptx_ldmatrix`` API taking explicit offsets.

    Signature: ``(trans, num, dtype, local_ptr, local_offset, smem_ptr,
    smem_offset)``. Offsets are folded into the pointers via
    ``tvm_access_ptr`` and dispatched to the fork-native
    :func:`ptx_ldmatrix`.

    ``T.ptx.ldmatrix_legacy`` runs through ``_dtype_forward`` which
    prepends a ``dtype=`` kwarg as a leading positional naming the buffer
    element type — offsets are in elements of that dtype, not bytes, so
    we forward it to ``tvm_access_ptr`` for correct scaling.
    """
    if len(all_args) == 8:
        elem_dtype, trans, num, dtype, local_ptr, local_offset, smem_ptr, smem_offset = all_args
    elif len(all_args) == 7:
        trans, num, dtype, local_ptr, local_offset, smem_ptr, smem_offset = all_args
        elem_dtype = "int8"
    else:
        raise ValueError(
            f"ptx_ldmatrix_legacy expects 7 args (or 8 with dtype= kwarg "
            f"prepended); got {len(all_args)}"
        )
    # Call.dtype carries the buffer element type so codegen can pick the
    # int8+trans manual-loop fallback (ldmatrix can't transpose int8).
    return call_intrin(
        elem_dtype,
        "tirx.ptx.ldmatrix_legacy",
        trans,
        num,
        dtype,
        local_ptr,
        local_offset,
        smem_ptr,
        smem_offset,
    )


def ptx_stmatrix(trans, num, dtype, smem_ptr, *src_handles, shape="m8n8", space="shared"):
    """TVM intrinsic for ``stmatrix.sync.aligned.shape.x{num}{.trans}.space.{dtype}``.

    Mirrors :func:`ptx_ldmatrix`: each source register is a separate operand.
    Pass ``T.address_of(buf[idx])`` (or ``buf.ptr_to([idx])``) for each
    source — the slots may be non-contiguous.

    Parameters
    ----------
    trans : bool
        Apply the ``.trans`` modifier (required for ``shape == "m16n8"``).
    num : int
        One of 1, 2, 4 — number of m8n8 fragments per warp.
    dtype : str
        ``".b16"`` (4 bytes per fragment register) or ``".b8"`` (2 bytes per).
    smem_ptr : Expr
        Destination pointer in shared memory.
    *src_handles : Expr
        ``num`` pointer-to-uint32 sources.
    shape : str, keyword-only, default "m8n8"
        ``"m8n8"`` or ``"m16n8"``.
    space : str, keyword-only, default "shared"
        ``"shared"`` or ``"shared::cta"``.

    https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-instructions-stmatrix
    """
    _choice("num", num, _LDMATRIX_NUM)
    _choice("dtype", dtype, _LDMATRIX_DTYPE)
    if shape not in ("m8n8", "m16n8"):
        raise ValueError(f"Unsupported stmatrix shape {shape!r}")
    if space not in ("shared", "shared::cta"):
        raise ValueError(f"Unsupported stmatrix state space {space!r}")
    if shape == "m16n8" and not trans:
        raise ValueError("stmatrix .m16n8 requires .trans")
    n_regs = int(num)
    if len(src_handles) != n_regs:
        dtype_bare = dtype.lstrip(".") if isinstance(dtype, str) else dtype
        raise ValueError(
            f"stmatrix .x{int(num)}.{dtype_bare} expects {n_regs} source "
            f"handles, got {len(src_handles)}"
        )
    return call_intrin(
        "", "tirx.ptx.stmatrix", trans, num, dtype, shape, space, smem_ptr, *src_handles
    )


def ptx_wgmma_encode_matrix_descriptor(desc, addr, ldo, sdo, swizzle):
    """TVM intrinsic to create memory descriptor for wgmma instructions

    Parameters
    ----------
    desc : Expr
        The pointer to the shared memory descriptor.

    addr : Expr
        The address of the matrix.

    ldo : Expr
        The leading dimension offset.

    sdo : Expr
        The stride dimension offset.

    swizzle : int
        The swizzle value (CUtensorMapSwizzle_enum).
    """
    return call_intrin("", "tirx.ptx.wgmma_encode_matrix_descriptor", desc, addr, ldo, sdo, swizzle)


def ptx_wgmma_noop_barrier(reg):
    """TVM intrinsic to call "" : "+{format}"(reg)::"memory"

    Parameters
    ----------
    reg : Expr
        The register to fence.

    Returns
    -------
    call : Expr
        The call expression.
    """
    return call_intrin("", "tirx.ptx.wgmma_noop_barrier", reg)


def ptx_wgmma_mma_async_ss(
    descA, descB, *accums, M, N, K, in_dtype, out_dtype, transA, transB, scaleA, scaleB, scaleD
):
    """TVM intrinsic to call wgmma.mma_async.sync.aligned.shape.dtype.atype.btype over 2 smem operators

    Parameters
    ----------
    M : int
        The number of rows in matrix A and D.

    N : int
        The number of columns in matrix B and D.

    K : int
        The number of columns in matrix A and rows in matrix B.

    in_dtype : str
        The data type of the input matrices.

    out_type : str
        The data type of the output matrices.

    transA : bool
        True for M/N major, False for K major.

    transB : bool
        True for M/N major, False for K major.

    scaleA : float
        The scaling factor for matrix A.

    scaleB : float
        The scaling factor for matrix B.

    scaleD : Expr
        True: D = A * B + D, False: D = A * B.

    descA : Expr
        The SMEM descriptor of matrix A

    descB : Expr
        The SMEM descriptor of matrix B

    accums : list
        The accumulators registers.
    """  # noqa: E501
    return call_intrin(
        "",
        "tirx.ptx.wgmma_mma_async_ss",
        M,
        N,
        K,
        in_dtype,
        out_dtype,
        transA,
        transB,
        scaleA,
        scaleB,
        scaleD,
        descA,
        descB,
        *accums,
    )


def ptx_wgmma_mma_async_rs(
    descB, *reg_list, M, N, K, in_dtype, out_dtype, transA, transB, scaleA, scaleB, scaleD
):
    """TVM intrinsic to call wgmma.mma_async.sync.aligned.shape.dtype.atype.btype
        When A is in register and B is in shared memory

    Parameters
    ----------
    M : int
        The number of rows in matrix A and D.

    N : int
        The number of columns in matrix B and D.

    K : int
        The number of columns in matrix A and rows in matrix B.

    in_dtype : str
        The data type of the input matrices.

    out_type : str
        The data type of the output matrices.

    transA : bool
        True for M/N major, False for K major.

    transB : bool
        True for M/N major, False for K major.

    scaleA : float
        The scaling factor for matrix A.

    scaleB : float
        The scaling factor for matrix B.

    scaleD : Expr
        True: D = A * B + D, False: D = A * B.

    descB : Expr
        The SMEM descriptor of matrix B

    reg_list : list
        The A registers and accumulators registers.
    """
    return call_intrin(
        "",
        "tirx.ptx.wgmma_mma_async_rs",
        M,
        N,
        K,
        in_dtype,
        out_dtype,
        transA,
        transB,
        scaleA,
        scaleB,
        scaleD,
        descB,
        *reg_list,
    )


def ptx_wgmma_fence():
    """TVM intrinsic to call wgmma.fence.sync.aligned

    Returns
    -------
    call : Expr
        The call expression.
    """
    return call_intrin("", "tirx.ptx.wgmma_fence")


def ptx_wgmma_commit_group():
    """TVM intrinsic to call wgmma.commit_group.sync.aligned

    Returns
    -------
    call : Expr
        The call expression.
    """
    return call_intrin("", "tirx.ptx.wgmma_commit_group")


def ptx_wgmma_wait_group(n):
    """TVM intrinsic to call wgmma.wait_group.sync.aligned

    Parameters
    ----------
    n : int
        The number of the most recent uncommitted pending wgmma groups to wait.

    Returns
    -------
    call : Expr
        The call expression.
    """
    return call_intrin("", "tirx.ptx.wgmma_wait_group", n)


def ptx_setmaxnreg(inc: bool, reg_count):
    """TVM intrinsic to call setmaxnreg.action.sync.aligned.u32 imm-reg-count

    Parameters
    ----------
    inc : bool
        True to increase the register count, False to decrease.

    reg_count : int
        The register count.
    """
    return call_intrin("", "tirx.ptx.setmaxnreg", inc, reg_count)


def ptx_tcgen05_alloc(dst_ptr, n_cols, cta_group=1):
    """TVM intrinsic to call tcgen05.alloc.cta_group.sync.aligned
        Dynamically allocates the number of cols in tensor memory, and write
        the address of allocated memory to shared memory.

    Parameters
    ----------
    dst_ptr : Var
        The pointer to the destination shared memory.

    n_cols : int
        The number of columns to allocate in tensor memory.
        Must be a multiple of 32 and a power of 2, and within the range [32, 512].

    cta_group : int
        The number of CTA groups involved in the allocation.
        If cta_group=1, one warp from CTA performs the allocation. Else, if cta_group=2,
        one warp from each of the peer CTAs perform the allocation.
    """
    _choice("cta_group", cta_group, _TCGEN05_CTA_GROUP)
    return call_intrin("", "tirx.ptx.tcgen05_alloc", dst_ptr, n_cols, cta_group)


def ptx_tcgen05_dealloc(taddr, n_cols, cta_group=1):
    """TVM intrinsic to call tcgen05.dealloc.cta_group.sync.aligned
        Deallocates the tensor memory specified by the tensor memory address taddr.

    Parameters
    ----------
    taddr : Expr
        The address of previously allocated tensor memory, should be uint32_t.

    n_cols : int
        The number of columns to deallocate in tensor memory.
        Must be a multiple of 32 and a power of 2, and within the range [32, 512].

    cta_group : int
        The number of CTA groups involved in the deallocation.
        If cta_group=1, one warp from CTA performs the deallocation. Else, if cta_group=2,
        one warp from each of the peer CTAs perform the deallocation.
    """
    _choice("cta_group", cta_group, _TCGEN05_CTA_GROUP)
    return call_intrin("", "tirx.ptx.tcgen05_dealloc", taddr, n_cols, cta_group)


def ptx_tcgen05_relinquish_alloc_permit(cta_group=1):
    """TVM intrinsic to call tcgen05.relinquish_alloc_permit.cta_group.sync.aligned
        The CTA of the executing thread is relinquishing the right to allocate
        Tensor Memory after calling this op.

    Parameters
    ----------
    cta_group : int
        The number of CTA groups involved in relinquishing.
        If cta_group=1, one warp from CTA performs the relinquishing. Else, if cta_group=2,
        one warp from each of the peer CTAs perform the relinquishing.
    """
    _choice("cta_group", cta_group, _TCGEN05_CTA_GROUP)
    return call_intrin("", "tirx.ptx.tcgen05_relinquish_alloc_permit", cta_group)


def ptx_tcgen05_encode_matrix_descriptor(desc, addr, ldo, sdo, swizzle):
    """TVM intrinsic to create memory descriptor for tcgen05 instructions

    Parameters
    ----------
    desc : Expr
        The pointer to the shared memory descriptor.

    addr : Expr
        The address of the matrix.

    ldo : Expr
        The leading dimension offset.

    sdo : Expr
        The stride dimension offset.

    swizzle : int
        The swizzle value (CUtensorMapSwizzle_enum).
    """
    return call_intrin(
        "", "tirx.ptx.tcgen05_encode_matrix_descriptor", desc, addr, ldo, sdo, swizzle
    )


def ptx_tcgen05_encode_instr_descriptor(
    desc,
    *,
    d_dtype,
    a_dtype,
    b_dtype,
    M,
    N,
    K,
    trans_a,
    trans_b,
    n_cta_groups=1,
    neg_a=False,
    neg_b=False,
    sat_d=False,
    is_sparse=False,
):
    """TVM intrinsic to create instruction descriptor for tcgen05 MMA without block scaling

    Parameters
    ----------
    desc : Expr
        The pointer to the instruction descriptor.

    d_dtype : str
        The datatype of resultant matrix D.

    a_dtype : str
        The datatype of multiplicand matrix A.

    b_dtype : str
        The datatype of multiplicand matrix B.

    M : int
        The size of non-reduction dimension of Matrix A.

    N : int
        The size of non-reduction dimension of Matrix B.

    K : int
        The size of reduction dimension of Matrix A/B.

    trans_a : bool
        Whether the multiplicand matrix A is transposed.
        True for M/N major, False for K major.

    trans_b : bool
        Whether the multiplicand matrix B is transposed.
        True for M/N major, False for K major.

    n_cta_groups : int
        The number of CTA groups involved in the MMA operation.

    neg_a : bool
        Whether to negate the multiplicand matrix A.

    neg_b : bool
        Whether to negate the multiplicand matrix B.

    sat_d : bool
        Whether to saturate the resultant matrix D.

    is_sparse : bool
        Whether the MMA operation is sparse.
    """
    _choice("n_cta_groups", n_cta_groups, _TCGEN05_CTA_GROUP)
    return call_intrin(
        "",
        "tirx.ptx.tcgen05_encode_instr_descriptor",
        desc,
        d_dtype,
        a_dtype,
        b_dtype,
        M,
        N,
        K,
        trans_a,
        trans_b,
        n_cta_groups,
        neg_a,
        neg_b,
        sat_d,
        is_sparse,
    )


def ptx_tcgen05_encode_instr_descriptor_block_scaled(
    desc,
    *,
    d_dtype,
    a_dtype,
    b_dtype,
    sfa_dtype,
    sfb_dtype,
    sfa_tmem_addr,
    sfb_tmem_addr,
    M,
    N,
    K,
    trans_a,
    trans_b,
    n_cta_groups=1,
    neg_a=False,
    neg_b=False,
    is_sparse=False,
):
    """TVM intrinsic to create instruction descriptor for tcgen05 MMA with block scaling

    Parameters
    ----------
    desc : Expr
        The pointer to the instruction descriptor.

    d_dtype : str
        The datatype of resultant matrix D.

    a_dtype : str
        The datatype of multiplicand matrix A.

    b_dtype : str
        The datatype of multiplicand matrix B.

    sfa_dtype : str
        The datatype of scale factor matrix A.

    sfb_dtype : str
        The datatype of scale factor matrix B.

    sfa_tmem_addr : Expr
        The address of the scale factor matrix A in tensor memory, should be uint32_t.

    sfb_tmem_addr : Expr
        The address of the scale factor matrix B in tensor memory, should be uint32_t.

    M : int
        The size of non-reduction dimension of Matrix A.

    N : int
        The size of non-reduction dimension of Matrix B.

    K : int
        The size of reduction dimension of Matrix A/B.

    trans_a : bool
        Whether the multiplicand matrix A is transposed.
        True for M/N major, False for K major.

    trans_b : bool
        Whether the multiplicand matrix B is transposed.
        True for M/N major, False for K major.

    n_cta_groups : int
        The number of CTA groups involved in the MMA operation.

    neg_a : bool
        Whether to negate the multiplicand matrix A.

    neg_b : bool
        Whether to negate the multiplicand matrix B.

    is_sparse : bool
        Whether the MMA operation is sparse.
    """
    _choice("n_cta_groups", n_cta_groups, _TCGEN05_CTA_GROUP)
    return call_intrin(
        "",
        "tirx.ptx.tcgen05_encode_instr_descriptor_block_scaled",
        desc,
        d_dtype,
        a_dtype,
        b_dtype,
        sfa_dtype,
        sfb_dtype,
        sfa_tmem_addr,
        sfb_tmem_addr,
        M,
        N,
        K,
        trans_a,
        trans_b,
        n_cta_groups,
        neg_a,
        neg_b,
        is_sparse,
    )


def ptx_tcgen05_mma(
    d_tmem_addr,
    a_operand,
    b_desc,
    i_desc,
    *disable_output_lane,
    d_dtype,
    a_dtype,
    b_dtype,
    use_a_tmem,
    cta_group,
    enable_input_d=1,
    scale_input_d=0,
    pred=None,
):
    """TVM intrinsic to call tcgen05.mma.cta_group.kind without block scaling.

    Parameters
    ----------
    d_dtype : str
        The datatype of resultant matrix D.

    a_dtype : str
        The datatype of multiplicand matrix A.

    b_dtype : str
        The datatype of multiplicand matrix B.

    d_tmem_addr : Expr
        The address of the resultant matrix D in tensor memory, should be uint32_t.

    a_operand : Expr
        Either the matrix descriptor of multiplicand matrix A in shared memory,
        or the address of the multiplicand matrix A in tensor memory (uint32_t).

    b_desc : Expr
        The matrix descriptor of multiplicand matrix B in shared memory.

    i_desc : Expr
        The instruction descriptor of the MMA operation.

    use_a_tmem : bool
        Whether the multiplicand matrix A is in tensor memory.

    cta_group : int
        The number of CTA groups involved in the MMA operation.

    enable_input_d : Expr
        Scale operand for the input accumulator C/D. The inline asm tests
        `enable_input_d != 0`: zero means D = A*B, non-zero means D = A*B + D.

    scale_input_d : int
        The optional scaling factor to scale input matrix D.
        D = A*B+D * (2 ^ - scale-input-d)

    disable_output_lane : list
        The lanes that should not be updated in the resultant matrix D.

    pred : Optional[Expr]
        Runtime ``uint32`` instruction-level predicate. When given, emit
        ``@p_issue tcgen05.mma...`` with ``p_issue = (pred != 0)``. Preserves
        PTX-level predicate semantics (single predicated SASS instruction).
    """

    _choice("cta_group", cta_group, _TCGEN05_CTA_GROUP)

    # default value for disable_output_lane
    if len(disable_output_lane) == 0:
        disable_output_lane = [0] * (4 if cta_group == 1 else 8)

    args = [
        d_dtype,
        a_dtype,
        b_dtype,
        d_tmem_addr,
        a_operand,
        b_desc,
        i_desc,
        use_a_tmem,
        cta_group,
        enable_input_d,
        scale_input_d,
        *disable_output_lane,
    ]
    if pred is not None:
        args.append(pred)
    return call_intrin("", "tirx.ptx.tcgen05_mma", *args)


def ptx_tcgen05_mma_block_scale(
    d_tmem_addr,
    a_operand,
    b_desc,
    sfa_tmem_addr,
    sfb_tmem_addr,
    i_desc,
    *,
    d_dtype,
    a_dtype,
    b_dtype,
    sfa_dtype,
    sfb_dtype,
    use_a_tmem,
    cta_group,
    enable_input_d=1,
):
    """TVM intrinsic to call tcgen05.mma.cta_group.kind.block_scale
        Performs matrix multiplication with block scaling:
        (A * scale_A)  * (B * scale_B) + D

    Parameters
    ----------
    d_dtype : str
        The datatype of resultant matrix D.

    a_dtype : str
        The datatype of multiplicand matrix A.

    b_dtype : str
        The datatype of multiplicand matrix B.

    sfa_dtype : str
        The datatype of scale factor matrix A.

    sfb_dtype : str
        The datatype of scale factor matrix B.

    d_tmem_addr : Expr
        The address of the resultant matrix D in tensor memory, should be uint32_t.

    a_operand : Expr
        Either the matrix descriptor of multiplicand matrix A in shared memory,
        or the address of the multiplicand matrix A in tensor memory (uint32_t).

    b_desc : Expr
        The matrix descriptor of multiplicand matrix B in shared memory.

    sfa_tmem_addr : Expr
        The address of the scale factor matrix A in tensor memory, should be uint32_t.

    sfb_tmem_addr : Expr
        The address of the scale factor matrix B in tensor memory, should be uint32_t.

    i_desc : Expr
        The instruction descriptor of the MMA operation.

    use_a_tmem : bool
        Whether the multiplicand matrix A is in tensor memory.

    cta_group : int
        The number of CTA groups involved in the MMA operation.

    enable_input_d : Expr
        Scale operand for the input accumulator C/D. Zero means D = A*B,
        non-zero means D = A*B + D.
    """

    _choice("cta_group", cta_group, _TCGEN05_CTA_GROUP)
    return call_intrin(
        "",
        "tirx.ptx.tcgen05_mma_block_scale",
        d_dtype,
        a_dtype,
        b_dtype,
        sfa_dtype,
        sfb_dtype,
        d_tmem_addr,
        a_operand,
        b_desc,
        sfa_tmem_addr,
        sfb_tmem_addr,
        i_desc,
        use_a_tmem,
        cta_group,
        enable_input_d,
    )


def ptx_tcgen05_mma_sp(
    d_tmem_addr,
    a_operand,
    b_desc,
    sp_tmem_addr,
    i_desc,
    *disable_output_lane,
    d_dtype,
    a_dtype,
    b_dtype,
    use_a_tmem,
    cta_group,
    enable_input_d=1,
    scale_input_d=0,
):
    """TVM intrinsic to call tcgen05.mma.sp.cta_group.kind without block scaling.

    Parameters
    ----------
    d_dtype : str
        The datatype of resultant matrix D.

    a_dtype : str
        The datatype of multiplicand matrix A.

    b_dtype : str
        The datatype of multiplicand matrix B.

    d_tmem_addr : Expr
        The address of the resultant matrix D in tensor memory, should be uint32_t.

    a_operand : Expr
        Either the matrix descriptor of multiplicand matrix A in shared memory,
        or the address of the multiplicand matrix A in tensor memory (uint32_t).

    b_desc : Expr
        The matrix descriptor of multiplicand matrix B in shared memory.

    sp_tmem_addr : Expr
        The address of the metadata of sparse matrix in tensor memory, should be uint32_t.

    i_desc : Expr
        The instruction descriptor of the MMA operation.

    use_a_tmem : bool
        Whether the multiplicand matrix A is in tensor memory.

    cta_group : int
        The number of CTA groups involved in the MMA operation.

    enable_input_d : Expr
        Scale operand for the input accumulator C/D. The inline asm tests
        `enable_input_d != 0`: zero means D = A*B, non-zero means D = A*B + D.

    scale_input_d : int
        The optional scaling factor to scale input matrix D.
        D = A*B+D * (2 ^ - scale-input-d)

    disable_output_lane : list
        The lanes that should not be updated in the resultant matrix D.
    """

    _choice("cta_group", cta_group, _TCGEN05_CTA_GROUP)

    # default value for disable_output_lane
    if len(disable_output_lane) == 0:
        disable_output_lane = [0] * (4 if cta_group == 1 else 8)

    return call_intrin(
        "",
        "tirx.ptx.tcgen05_mma_sp",
        d_dtype,
        a_dtype,
        b_dtype,
        d_tmem_addr,
        a_operand,
        b_desc,
        sp_tmem_addr,
        i_desc,
        use_a_tmem,
        cta_group,
        enable_input_d,
        scale_input_d,
        *disable_output_lane,
    )


def ptx_tcgen05_mma_sp_block_scale(
    d_tmem_addr,
    a_operand,
    b_desc,
    sfa_tmem_addr,
    sfb_tmem_addr,
    sp_tmem_addr,
    i_desc,
    *,
    d_dtype,
    a_dtype,
    b_dtype,
    sfa_dtype,
    sfb_dtype,
    use_a_tmem,
    cta_group,
    enable_input_d=1,
):
    """TVM intrinsic to call tcgen05.mma.sp.cta_group.kind.block_scale
        Performs sparse matrix multiplication with block scaling:
        (A * scale_A)  * (B * scale_B) + D

    Parameters
    ----------
    d_dtype : str
        The datatype of resultant matrix D.

    a_dtype : str
        The datatype of multiplicand matrix A.

    b_dtype : str
        The datatype of multiplicand matrix B.

    sfa_dtype : str
        The datatype of scale factor matrix A.

    sfb_dtype : str
        The datatype of scale factor matrix B.

    d_tmem_addr : Expr
        The address of the resultant matrix D in tensor memory, should be uint32_t.

    a_operand : Expr
        Either the matrix descriptor of multiplicand matrix A in shared memory,
        or the address of the multiplicand matrix A in tensor memory (uint32_t).

    b_desc : Expr
        The matrix descriptor of multiplicand matrix B in shared memory.

    sfa_tmem_addr : Expr
        The address of the scale factor matrix A in tensor memory, should be uint32_t.

    sfb_tmem_addr : Expr
        The address of the scale factor matrix B in tensor memory, should be uint32_t.

    sp_tmem_addr : Expr
        The address of the metadata of sparse matrix in tensor memory, should be uint32_t.

    i_desc : Expr
        The instruction descriptor of the MMA operation.

    use_a_tmem : bool
        Whether the multiplicand matrix A is in tensor memory.

    cta_group : int
        The number of CTA groups involved in the MMA operation.

    enable_input_d : Expr
        Scale operand for the input accumulator C/D. Zero means D = A*B,
        non-zero means D = A*B + D.
    """
    _choice("cta_group", cta_group, _TCGEN05_CTA_GROUP)
    return call_intrin(
        "",
        "tirx.ptx.tcgen05_mma_sp_block_scale",
        d_dtype,
        a_dtype,
        b_dtype,
        sfa_dtype,
        sfb_dtype,
        d_tmem_addr,
        a_operand,
        b_desc,
        sfa_tmem_addr,
        sfb_tmem_addr,
        sp_tmem_addr,
        i_desc,
        use_a_tmem,
        cta_group,
        enable_input_d,
    )


def ptx_tcgen05_fence_before_thread_sync():
    """TVM intrinsic to call tcgen05.fence::before_thread_sync
    Orders all prior asynchronous tcgen05 operations relative to subsequent operations.
    """
    return call_intrin("", "tirx.ptx.tcgen05_fence_before_thread_sync")


def ptx_tcgen05_fence_after_thread_sync():
    """TVM intrinsic to call tcgen05.fence::after_thread_sync
    Orders all subsequent asynchronous tcgen05 operations relative to previous operations.
    """
    return call_intrin("", "tirx.ptx.tcgen05_fence_after_thread_sync")


def _choice(name: str, value, options):
    """Validate `value` is one of `options`. Raise a clear ValueError otherwise.

    Symbolic values (Var, non-constant Expr) are accepted without
    validation; specialization later replaces them with concrete values
    that the C-side intrinsic body re-checks.
    """
    # Concrete int / IntImm value: validate.
    try:
        concrete = int(value)
    except (TypeError, ValueError):
        return  # symbolic; defer check
    if concrete not in options:
        raise ValueError(f"invalid {name}={concrete!r}; expected one of {tuple(options)}")


# See top-of-file imports for `_FENCE_SEM` etc. (re-exported from _common).
# Note: TCGEN05_LDST_SHAPES values must stay in sync with the shape branches
# of codegen_ptx_tcgen05_ld/_st in intrinsics/cuda/tcgen05.py.


def ptx_tcgen05_cp(
    taddr, src_desc, *, shape, cta_group=1, multicast="", decompress="", row=0, col=0
):
    """TVM intrinsic for the Blackwell `tcgen05.cp` PTX instruction.

    The emitted PTX is::

        tcgen05.cp.cta_group::{cta_group}.{shape}[.{multicast}][.{decompress}] [taddr], src_desc;

    Each keyword argument maps 1:1 to a PTX token: read the call and you
    know what instruction is emitted.

    Parameters
    ----------
    taddr : Expr
        Destination tensor-memory address (uint32). Callers typically pass
        ``tmem_base + column_offset_in_uint32s`` directly. Use the optional
        ``row`` / ``col`` keyword arguments only when the address needs
        runtime row/col composition via ``get_tmem_addr`` (high 16 bits row,
        low 16 bits col).

    src_desc : Expr
        The 64-bit shared-memory matrix descriptor.

    shape : str
        One of ``"32x128b"``, ``"4x256b"``, ``"128x128b"``, ``"128x256b"``,
        ``"64x128b"``.

    cta_group : int
        1 or 2.

    multicast : str
        One of ``""``, ``"warpx4"``, ``"warpx2::02_13"``, ``"warpx2::01_23"``.
        ``"32x128b"`` requires ``"warpx4"``; ``"64x128b"`` requires one of the
        ``warpx2::*`` values; other shapes require ``""``.

    decompress : str
        Trailing PTX suffix for fp4/fp6 → fp8 on-the-fly decompression.
        One of ``""``, ``"b8x16.b4x16_p64"``, ``"b8x16.b6x16_p32"``.

    row, col : Expr
        Optional row/col offsets added to ``taddr`` at runtime. Default 0.
    """
    _choice("shape", shape, _TCGEN05_CP_SHAPES)
    _choice("cta_group", cta_group, _TCGEN05_CTA_GROUP)
    _choice("multicast", multicast, _TCGEN05_CP_MULTICAST)
    _choice("decompress", decompress, _TCGEN05_CP_DECOMPRESS)
    if shape == "32x128b" and multicast != "warpx4":
        raise ValueError(f"shape=32x128b requires multicast='warpx4', got {multicast!r}")
    if shape == "64x128b" and multicast not in ("warpx2::02_13", "warpx2::01_23"):
        raise ValueError(f"shape=64x128b requires multicast in warpx2::*, got {multicast!r}")
    if shape in ("128x128b", "128x256b", "4x256b") and multicast != "":
        raise ValueError(f"shape={shape} requires multicast='', got {multicast!r}")

    return call_intrin(
        "",
        "tirx.ptx.tcgen05_cp",
        taddr,
        src_desc,
        shape,
        cta_group,
        multicast,
        decompress,
        row,
        col,
    )


def ptx_tcgen05_shift(taddr, cta_group=1):
    """TVM intrinsic to call tcgen05.shift.cta_group.down
        Asynchronously shift down the rows of the matrix in Tensor Memory for a warp.

    Parameters
    ----------
    taddr : Expr
        The address of matrix in tensor memory, should be uint32_t.

    cta_group : int
        The number of CTA groups involved in the shift.
        If cta_group=1, shift operation is performed in the Tensor Memory of current CTA.
        Else, shift operation is performed in the Tensor Memory of both the current CTA and
        the peer CTA.
    """
    _choice("cta_group", cta_group, _TCGEN05_CTA_GROUP)
    return call_intrin("", "tirx.ptx.tcgen05_shift", taddr, cta_group)


def ptx_tcgen05_ld(src_addr, *regs, shape, num, row=0, col=0, pack=False):
    """TVM intrinsic for tcgen05.ld.sync.aligned — async collective load from TMEM.

    Emits ``tcgen05.ld.sync.aligned.{shape}.x{num}[.pack::16b].b32 {regs}, [addr];``

    Parameters
    ----------
    src_addr : Expr
        Tensor-memory source address (uint32).

    regs : list[Expr]
        Destination registers. Count depends on shape x num.

    shape : str
        One of ``"16x32bx2"``, ``"16x64b"``, ``"16x128b"``, ``"16x256b"``, ``"32x32b"``.

    num : int
        Repeat factor along the columns. Power-of-two in [1, 128].

    row, col : Expr
        Optional TMEM row/col offsets added to ``src_addr`` at runtime (row must be
        a multiple of 32). Default 0.

    pack : bool
        Pack two 16-bit chunks into a single 32-bit register.
    """
    _choice("shape", shape, _TCGEN05_LDST_SHAPES)
    return call_intrin("", "tirx.ptx.tcgen05_ld", src_addr, row, col, shape, num, pack, *regs)


def ptx_tcgen05_st(dst_addr, *regs, shape, num, row=0, col=0, unpack=False):
    """TVM intrinsic for tcgen05.st.sync.aligned — async collective store to TMEM.

    Emits ``tcgen05.st.sync.aligned.{shape}.x{num}[.unpack::16b].b32 [addr], {regs};``

    Parameters
    ----------
    dst_addr : Expr
        Tensor-memory destination address (uint32).

    regs : list[Expr]
        Source registers. Count depends on shape x num.

    shape : str
        One of ``"16x32bx2"``, ``"16x64b"``, ``"16x128b"``, ``"16x256b"``, ``"32x32b"``.

    num : int
        Repeat factor along the columns. Power-of-two in [1, 128].

    row, col : Expr
        Optional TMEM row/col offsets added to ``dst_addr`` at runtime (row must be
        a multiple of 32). Default 0.

    unpack : bool
        Unpack a 32-bit register into two 16-bit chunks.
    """
    _choice("shape", shape, _TCGEN05_LDST_SHAPES)
    return call_intrin("", "tirx.ptx.tcgen05_st", dst_addr, row, col, shape, num, unpack, *regs)


def ptx_tcgen05_wait_ld():
    """TVM intrinsic to call tcgen05.wait::ld.sync.aligned
    Wait for the completion of all prior async tcgen05.ld operations.
    """
    return call_intrin("", "tirx.ptx.tcgen05_wait_ld")


def ptx_tcgen05_wait_st():
    """TVM intrinsic to call tcgen05.wait::st.sync.aligned
    Wait for the completion of all prior async tcgen05.st operations.
    """
    return call_intrin("", "tirx.ptx.tcgen05_wait_st")


def ptx_tcgen05_commit(bar, cta_group=1, cta_mask=0, *, pred=None):
    """TVM intrinsic to call tcgen05.commit.cta_group

    Parameters
    ----------
    bar : Expr
        The pointer to mbarrier variable.

    cta_group: int
        The number of CTA groups involved in previous tcgen05 operations.

    cta_mask : int
        The mask of the CTAs in the cluster, used for multicast.

    pred : Optional[Expr]
        Runtime ``uint32`` predicate. When given, emit
        ``@p tcgen05.commit...`` with ``p = (pred != 0)``. This preserves
        PTX-level instruction predicate semantics (single predicated
        instruction in SASS), distinct from a C-level ``if`` branch.

    Returns
    -------
    call : Expr
        The call expression.
    """
    _choice("cta_group", cta_group, _TCGEN05_CTA_GROUP)
    args = [bar, cta_group, cta_mask]
    if pred is not None:
        args.append(pred)
    return call_intrin("", "tirx.ptx.tcgen05_commit", *args)


def timer_init_cuda(profiler_buffer, profiler_tag, profiler_write_offset, num_groups, group_id):
    """TVM intrinsic for initializing the CUDA profiler, and store profiling result in a buffer.

    Parameters
    ----------
    profiler_buffer: Var
        The buffer to store the profiling result.

    profiler_tag: Var
        Buffer of length 1 storing the base tag of the current thread.

    profiler_write_offset: Var
        Buffer of length 1 storing the offset in buffer to write the next
        profiling result for the current thread.

    num_groups: int
        The number of groups in the profiler.

    group_id: Expr
        The group id of the current thread.

    Returns
    -------
    call : Expr
        The call expression.
    """

    return call_intrin(
        "void",
        "tirx.timer_init_cuda",
        profiler_buffer,
        profiler_tag,
        profiler_write_offset,
        num_groups,
        group_id,
    )


def timer_start_cuda(
    event_type,
    profiler_buffer,
    profiler_tag,
    profiler_write_offset,
    profiler_write_stride,
    leader_cond,
):
    """TVM intrinsic for starting the timer for profiling a specific event, and storing profiling result in a buffer.

    Parameters
    ----------
    event_type: Enum
        The event to profile.

    profiler_buffer: Var
        The buffer to store the profiling result.

    profiler_tag: Var
        Buffer of length 1 storing the base tag of the current thread.

    profiler_write_offset: Var
        Buffer of length 1 storing the offset in buffer to write the next
        profiling result for the current thread.

    profiler_write_stride: int
        The stride to advance in buffer in the next write.

    leader_cond: Expr
        The condition to check if the current thread is the leader.

    Returns
    -------
    call : Expr
        The call expression.
    """  # noqa: E501

    return call_intrin(
        "void",
        "tirx.timer_start_cuda",
        event_type.value,
        profiler_buffer,
        profiler_tag,
        profiler_write_offset,
        profiler_write_stride,
        leader_cond,
    )


def timer_end_cuda(
    event_type,
    profiler_buffer,
    profiler_tag,
    profiler_write_offset,
    profiler_write_stride,
    leader_cond,
):
    """TVM intrinsic for ending the timer for profiling a specific event, and storing profiling result in a buffer.

    Parameters
    ----------
    event_type: Enum
        The event to profile.

    profiler_buffer: Var
        The buffer to store the profiling result.

    profiler_tag: Var
        Buffer of length 1 storing the base tag of the current thread.

    profiler_write_offset: Var
        Buffer of length 1 storing the offset in buffer to write the next
        profiling result for the current thread.

    profiler_write_stride: int
        The stride to advance in buffer in the next write.

    leader_cond: Expr
        The condition to check if the current thread is the leader.

    Returns
    -------
    call : Expr
        The call expression.
    """  # noqa: E501

    return call_intrin(
        "void",
        "tirx.timer_end_cuda",
        event_type.value,
        profiler_buffer,
        profiler_tag,
        profiler_write_offset,
        profiler_write_stride,
        leader_cond,
    )


def timer_finalize_cuda(
    profiler_buffer, profiler_tag, profiler_write_offset, profiler_write_stride, leader_cond
):
    """TVM intrinsic for finalizing the CUDA profiler, and store profiling result in a buffer.

    Parameters
    ----------
    profiler_buffer: Var
        The buffer to store the profiling result.

    profiler_tag: Var
        Buffer of length 1 storing the base tag of the current thread.

    profiler_write_offset: Var
        Buffer of length 1 storing the offset in buffer to write the next
        profiling result for the current thread.

    profiler_write_stride: int
        The stride to advance in buffer in the next write.

    leader_cond: Expr
        The condition to check if the current thread is the leader.

    Returns
    -------
    call : Expr
        The call expression.
    """

    return call_intrin(
        "void",
        "tirx.timer_finalize_cuda",
        profiler_buffer,
        profiler_tag,
        profiler_write_offset,
        profiler_write_stride,
        leader_cond,
    )


def cuda_atomic_add(res_addr, value):
    """TVM intrinsic to call cuda atomic add instruction

    Parameters
    ----------
    res_addr : Expr
        The result address.

    value: Expr
        The value to add.

    Returns
    -------
    call : Expr
        The call expression.
    """
    value = tir.convert(value)
    return call_intrin(value.ty, "tirx.cuda.atomic_add", res_addr, value)


def cuda_thread_fence():
    """TVM intrinsic to call cuda thread fence instruction

    Returns
    -------
    call : Expr
        The call expression.
    """
    return call_intrin("", "tirx.cuda.thread_fence")


def cuda_warpgroup_sync(bar_no):
    """TVM intrinsic to synchronize a CUDA warpgroup via a named barrier.

    Parameters
    ----------
    bar_no : Expr
        The named barrier id to use for the warpgroup.

    Notes
    -----
    Synchronizes 128 threads in a warpgroup using `bar.sync bar_no, 128`.

    Returns
    -------
    call : Expr
        The call expression.
    """
    return call_intrin("", "tirx.cuda.warpgroup_sync", bar_no)


def cuda_syncthreads_and(cond):
    """TVM intrinsic to call cuda syncthreads_and instruction

    Parameters
    ----------
    cond: Expr
        The condition.

    Returns
    -------
    call : Expr
        The call expression.
    """
    return call_intrin("int64", "tirx.cuda.syncthreads_and", cond)


def cuda_syncthreads_or(cond):
    """TVM intrinsic to call cuda syncthreads_or instruction

    Parameters
    ----------
    cond: Expr
        The condition.

    Returns
    -------
    call : Expr
        The call expression.
    """
    return call_intrin("int64", "tirx.cuda.syncthreads_or", cond)


def cuda_nano_sleep(time):
    """TVM intrinsic to call cuda nano sleep instruction

    Parameters
    ----------
    time: Expr
        The time to sleep.

    Returns
    -------
    call : Expr
        The call expression.
    """
    return call_intrin("", "tirx.cuda.nano_sleep", time)


def cuda_printf(fmt, *args):
    """TVM intrinsic to call cuda printf instruction

    Parameters
    ----------
    fmt: str
        The format string.

    *args: list
        The arguments to the format string.

    Returns
    -------
    call : Expr
        The call expression.
    """
    return call_intrin("", "tirx.cuda.printf", fmt, *args)


def cuda_ldg(addr, dtype):
    """TVM intrinsic to call CUDA C++ __ldg() function

    Parameters
    ----------
    addr : Expr
        The memory address to load.

    dtype : str
        The data type of the loaded value.

    Returns
    """
    return call_intrin(dtype, "tirx.cuda.ldg", addr, dtype)


def cuda_get_tmem_addr(addr, row_offset, col_offset):
    """TVM intrinsic to call cuda tmem address calculation

    Parameters
    ----------
    addr: Expr
        The memory address to calculate.

    row_offset: Expr
        The row offset to calculate.

    col_offset: Expr
        The column offset to calculate.

    Returns
    -------
    call : Expr
        The call expression.
    """
    return call_intrin("uint32", "tirx.cuda.get_tmem_addr", addr, row_offset, col_offset)


def cuda_cvta_generic_to_shared(ptr):
    """Convert a generic pointer to a shared-memory address (uint32).

    Wraps ``__cvta_generic_to_shared(ptr)``. Used by op-wrappers that
    precompute the shared-memory address at the wrapper layer instead of
    inside the asm helper body.
    """
    return call_intrin("uint32", "tirx.cuda.cvta_generic_to_shared", ptr)


def cuda_smem_addr_from_uint64(cluster_addr):
    """Narrow a 64-bit cluster-mapped SMEM address to a 32-bit SMEM address.

    Wraps ``static_cast<unsigned int>(cluster_addr)``. Used by
    cp.async.bulk.shared::cluster.* op-wrappers.
    """
    return call_intrin("uint32", "tirx.cuda.smem_addr_from_uint64", cluster_addr)


def cuda_sm100_tma_2sm_mbarrier_addr(bar):
    """Compute the SM100 2SM TMA mbarrier shared-address operand."""
    return bitwise_and(cuda_cvta_generic_to_shared(bar), const(0xFEFFFFFF, dtype="uint32"))


def ptx_exp2(x):
    """TVM intrinsic for PTX fast exp2 approximation (ex2.approx.ftz.f32)

    Parameters
    ----------
    x : Expr
        The float32 input value.

    Returns
    -------
    call : Expr
        The call expression returning 2^x (approximate).
    """
    return call_intrin("float32", "tirx.ptx.exp2", x)


def ptx_rcp(x):
    """TVM intrinsic for PTX fast reciprocal approximation (rcp.approx.ftz.f32)

    Parameters
    ----------
    x : Expr
        The float32 input value.

    Returns
    -------
    call : Expr
        The call expression returning 1/x (approximate).
    """
    return call_intrin("float32", "tirx.ptx.rcp", x)


def ptx_any_sync(mask, pred):
    """TVM intrinsic for PTX warp-wide any predicate (__any_sync)

    Parameters
    ----------
    mask : Expr
        The thread mask (uint32).
    pred : Expr
        The predicate value (int32).

    Returns
    -------
    call : Expr
        The call expression returning 1 if any thread in mask has pred != 0.
    """
    return call_intrin("int32", "tirx.ptx.any_sync", mask, pred)


def ptx_reduce3_max_f32(a, b, c):
    """TVM intrinsic to call 3-input max.f32 PTX instruction (sm_100a+)

    Parameters
    ----------
    a, b, c : Expr
        The three float32 values to compare.

    Returns
    -------
    call : Expr
        The call expression returning max(a, b, c).
    """
    return call_intrin("float32", "tirx.ptx.reduce3_max_f32", a, b, c)


def ptx_reduce3_min_f32(a, b, c):
    """TVM intrinsic to call 3-input min.f32 PTX instruction (sm_100a+)

    Parameters
    ----------
    a, b, c : Expr
        The three float32 values to compare.

    Returns
    -------
    call : Expr
        The call expression returning min(a, b, c).
    """
    return call_intrin("float32", "tirx.ptx.reduce3_min_f32", a, b, c)


def _ptx_binary_arith(op_name, dtype, d, a, b, *, rounding="rn", ftz=False, sat=False):
    """Shared helper for add/sub/mul over (f32 | f32x2 | f64), DPS form."""
    _choice("rounding", rounding, _F32X2_ROUND)
    if dtype == "f64" and (ftz or sat):
        raise ValueError(f"PTX {op_name}.f64 does not accept .ftz or .sat")
    if dtype == "f32x2" and sat:
        raise ValueError(f"PTX {op_name}.f32x2 does not accept .sat")
    return call_intrin(
        "",
        f"tirx.ptx.{op_name}_{dtype}",
        d,
        a,
        b,
        rounding,
        int(ftz),
        int(sat),
    )


def _ptx_fma(dtype, d, a, b, c, *, rounding="rn", ftz=False, sat=False):
    """Shared helper for fma over (f32 | f32x2 | f64), DPS form."""
    _choice("rounding", rounding, _F32X2_ROUND)
    if dtype == "f64" and (ftz or sat):
        raise ValueError("PTX fma.f64 does not accept .ftz or .sat")
    if dtype == "f32x2" and sat:
        raise ValueError("PTX fma.f32x2 does not accept .sat")
    return call_intrin(
        "",
        f"tirx.ptx.fma_{dtype}",
        d,
        a,
        b,
        c,
        rounding,
        int(ftz),
        int(sat),
    )


def ptx_add_f32(d_addr, a, b, *, rounding="rn", ftz=False, sat=False):
    """PTX ``add{.rnd}{.ftz}{.sat}.f32 [d_addr], a, b`` — DPS form."""
    return _ptx_binary_arith("add", "f32", d_addr, a, b, rounding=rounding, ftz=ftz, sat=sat)


def ptx_add_f32x2(d_addr, a, b, *, rounding="rn", ftz=False):
    """PTX ``add{.rnd}{.ftz}.f32x2 [d_addr], a, b`` — DPS form.

    a, b are packed-as-uint64 register operands (2 fp32 each).
    """
    return _ptx_binary_arith("add", "f32x2", d_addr, a, b, rounding=rounding, ftz=ftz)


def ptx_add_f64(d_addr, a, b, *, rounding="rn"):
    """PTX ``add{.rnd}.f64 [d_addr], a, b`` — DPS form (no .ftz / .sat)."""
    return _ptx_binary_arith("add", "f64", d_addr, a, b, rounding=rounding)


def ptx_sub_f32(d_addr, a, b, *, rounding="rn", ftz=False, sat=False):
    """PTX ``sub{.rnd}{.ftz}{.sat}.f32 [d_addr], a, b`` — DPS form."""
    return _ptx_binary_arith("sub", "f32", d_addr, a, b, rounding=rounding, ftz=ftz, sat=sat)


def ptx_sub_f32x2(d_addr, a, b, *, rounding="rn", ftz=False):
    """PTX ``sub{.rnd}{.ftz}.f32x2 [d_addr], a, b`` — DPS form."""
    return _ptx_binary_arith("sub", "f32x2", d_addr, a, b, rounding=rounding, ftz=ftz)


def ptx_sub_f64(d_addr, a, b, *, rounding="rn"):
    """PTX ``sub{.rnd}.f64 [d_addr], a, b`` — DPS form."""
    return _ptx_binary_arith("sub", "f64", d_addr, a, b, rounding=rounding)


def ptx_mul_f32(d_addr, a, b, *, rounding="rn", ftz=False, sat=False):
    """PTX ``mul{.rnd}{.ftz}{.sat}.f32 [d_addr], a, b`` — DPS form."""
    return _ptx_binary_arith("mul", "f32", d_addr, a, b, rounding=rounding, ftz=ftz, sat=sat)


def ptx_mul_f32x2(d_addr, a, b, *, rounding="rn", ftz=False):
    """PTX ``mul{.rnd}{.ftz}.f32x2 [d_addr], a, b`` — DPS form."""
    return _ptx_binary_arith("mul", "f32x2", d_addr, a, b, rounding=rounding, ftz=ftz)


def ptx_mul_f64(d_addr, a, b, *, rounding="rn"):
    """PTX ``mul{.rnd}.f64 [d_addr], a, b`` — DPS form."""
    return _ptx_binary_arith("mul", "f64", d_addr, a, b, rounding=rounding)


def ptx_fma_f32(d_addr, a, b, c, *, rounding="rn", ftz=False, sat=False):
    """PTX ``fma{.rnd}{.ftz}{.sat}.f32 [d_addr], a, b, c`` — DPS form."""
    return _ptx_fma("f32", d_addr, a, b, c, rounding=rounding, ftz=ftz, sat=sat)


def ptx_fma_f32x2(d_addr, a, b, c, *, rounding="rn", ftz=False):
    """PTX ``fma{.rnd}{.ftz}.f32x2 [d_addr], a, b, c`` — DPS form.

    a, b, c are packed-as-uint64 register operands.
    """
    return _ptx_fma("f32x2", d_addr, a, b, c, rounding=rounding, ftz=ftz)


def ptx_fma_f64(d_addr, a, b, c, *, rounding="rn"):
    """PTX ``fma{.rnd}.f64 [d_addr], a, b, c`` — DPS form."""
    return _ptx_fma("f64", d_addr, a, b, c, rounding=rounding)


def ptx_max_f32(a, b, *, ftz=False, nan=False):
    """TVM intrinsic for PTX ``max{.ftz}{.NaN}.f32 d, a, b``.

    2-operand form (distinct from :func:`ptx_reduce3_max_f32` which is the
    3-operand SM_100+ form). ``.NaN`` qualifier propagates NaN inputs to
    the output; without it, NaN inputs are silently ignored.

    Parameters
    ----------
    a, b : Expr
        Float32 inputs.
    ftz : bool
        If True, flush subnormals to zero (``.ftz``).
    nan : bool
        If True, propagate NaN inputs (``.NaN``).
    """
    return call_intrin("float32", "tirx.ptx.max_f32", a, b, int(ftz), int(nan))


def ptx_griddepcontrol_wait():
    """TVM intrinsic for PTX ``griddepcontrol.wait`` (sm_90+).

    Blocks the current grid until prerequisite grids signalled via
    :func:`ptx_griddepcontrol_launch_dependents` have finished. Acts as a
    full memory barrier.
    """
    return call_intrin("", "tirx.ptx.griddepcontrol_wait")


def ptx_griddepcontrol_launch_dependents():
    """TVM intrinsic for PTX ``griddepcontrol.launch_dependents`` (sm_90+).

    Signals that the current grid has reached a point where dependent
    grids may begin execution.
    """
    return call_intrin("", "tirx.ptx.griddepcontrol_launch_dependents")


_PTX_LD_SCOPE = {"cta", "cluster", "gpu", "sys"}
_PTX_LD_SPACE = {"global", "shared", "shared::cta", "shared::cluster", "local"}
_PTX_LD_VOLATILE_SPACE = _PTX_LD_SPACE | {"const"}
_PTX_SCALAR_TYPE = {
    "b8",
    "u8",
    "s8",
    "b16",
    "u16",
    "s16",
    "b32",
    "b64",
    "u32",
    "u64",
    "s32",
    "s64",
    "f32",
    "f64",
}
_PTX_LD_TYPE = set(_PTX_SCALAR_TYPE)
_PTX_LD_COP = {"", "ca", "cg", "cs", "lu", "cv"}
_PTX_LD_VEC = {"", "v2", "v4", "v8"}
_PTX_L1_EVICT = {
    "",
    "L1::evict_normal",
    "L1::evict_unchanged",
    "L1::evict_first",
    "L1::evict_last",
    "L1::no_allocate",
}
_PTX_L2_EVICT = {"", "L2::evict_normal", "L2::evict_first", "L2::evict_last"}
_PTX_PREFETCH = {"", "L2::64B", "L2::128B", "L2::256B"}
_PTX_MEM_SCOPE = {"", "cta", "cluster", "gpu", "sys"}
_PTX_MEM_SPACE = {"global", "shared", "shared::cta", "shared::cluster"}
_PTX_RED_OP = {"and", "or", "xor", "add", "inc", "dec", "min", "max"}
_PTX_ATOM_OP = {"and", "or", "xor", "exch", "add", "inc", "dec", "min", "max"}
_PTX_ST_VEC = {"", "v2", "v4", "v8"}
_PTX_ST_COP = {"", "wb", "cg", "cs", "wt"}
_PTX_PREFETCH_TENSORMAP_SPACE = {"", "const", "param"}
_PTX_SCALAR_RETURN_TYPE = {
    "b32": "uint32",
    "u32": "uint32",
    "s32": "int32",
    "b64": "uint64",
    "u64": "uint64",
    "s64": "int64",
    "f32": "float32",
    "f64": "float64",
}
_PTX_CACHE_POLICY = {
    "evict_normal": 0x1000000000000000,
    "evict_first": 0x12F0000000000000,
    "evict_last": 0x14F0000000000000,
}


def _resolve_cache_policy(cache_hint, cache_policy, choices=_CP_ASYNC_BULK_CACHE_HINT):
    _choice("cache_hint", cache_hint, choices)
    if cache_policy is not None:
        return cache_policy, True
    if cache_hint:
        if cache_hint not in _PTX_CACHE_POLICY:
            raise ValueError(
                f"Unsupported built-in cache policy {cache_hint!r}; pass cache_policy explicitly"
            )
        return const(_PTX_CACHE_POLICY[cache_hint], dtype="uint64"), True
    return const(0, dtype="uint64"), False


def ptx_ld_acquire(
    addr,
    return_type,
    ptx_type,
    *,
    scope="gpu",
    space="global",
    vec="",
    dst=None,
    cache_hint="",
    cache_policy=None,
    l1_evict="",
    l2_evict="",
    prefetch_size="",
):
    """TVM intrinsic for PTX ``ld.acquire.scope{.ss}...`` loads.

    ``scope``, state ``space``, PTX ``type`` and TVM ``return_type`` are
    explicit so callers can request either raw-bit or typed loads.  The
    optional ``vec``/``dst`` and cache arguments cover vector and cache-policy
    forms of the same PTX instruction.

    Parameters
    ----------
    addr : Expr
        The memory address to load.

    return_type : str
        TVM dtype returned by the load.

    ptx_type : str
        PTX type suffix such as ``"b32"``, ``"u64"``, or ``"s32"``.

    scope : str
        PTX memory scope: ``"cta"``, ``"cluster"``, ``"gpu"``, or ``"sys"``.

    space : str
        PTX state space suffix.

    Returns
    -------
    call : Expr
        The loaded value.
    """
    _choice("scope", scope, _PTX_LD_SCOPE)
    _choice("space", space, _PTX_LD_SPACE)
    _choice("ptx_type", ptx_type, _PTX_LD_TYPE)
    _choice("vec", vec, _PTX_LD_VEC)
    cache_policy, has_cache_policy = _resolve_cache_policy(cache_hint, cache_policy)
    to_dst = int(dst is not None)
    if vec and dst is None:
        raise ValueError("vec ld.acquire requires dst")
    if to_dst:
        return call_intrin(
            "",
            "tirx.ptx.ld_acquire",
            dst,
            addr,
            cache_policy,
            return_type,
            scope,
            space,
            vec,
            ptx_type,
            int(has_cache_policy),
            to_dst,
            l1_evict,
            l2_evict,
            prefetch_size,
        )
    return call_intrin(
        return_type,
        "tirx.ptx.ld_acquire",
        addr,
        return_type,
        scope,
        space,
        vec,
        ptx_type,
        int(has_cache_policy),
        0,
        l1_evict,
        l2_evict,
        prefetch_size,
    )


def ptx_ld(
    addr,
    return_type,
    ptx_type,
    *,
    dst=None,
    weak=False,
    space="global",
    cop="",
    vec="",
    cache_hint="",
    cache_policy=None,
    l1_evict="",
    l2_evict="",
    prefetch_size="",
):
    """TVM intrinsic for PTX ``ld{.weak}{.ss}{.cop}...`` loads."""
    _choice("space", space, _PTX_LD_SPACE | {"const", "param::entry", "param::func"})
    _choice("cop", cop, _PTX_LD_COP)
    _choice("ptx_type", ptx_type, _PTX_LD_TYPE)
    _choice("vec", vec, _PTX_LD_VEC)
    cache_policy, has_cache_policy = _resolve_cache_policy(cache_hint, cache_policy)
    to_dst = int(dst is not None)
    if vec and dst is None:
        raise ValueError("vec ld requires dst")
    if to_dst:
        return call_intrin(
            "",
            "tirx.ptx.ld",
            dst,
            addr,
            cache_policy,
            return_type,
            int(bool(weak)),
            space,
            cop,
            vec,
            ptx_type,
            int(has_cache_policy),
            to_dst,
            l1_evict,
            l2_evict,
            prefetch_size,
        )
    return call_intrin(
        return_type,
        "tirx.ptx.ld",
        addr,
        cache_policy,
        return_type,
        int(bool(weak)),
        space,
        cop,
        "",
        ptx_type,
        int(has_cache_policy),
        0,
        l1_evict,
        l2_evict,
        prefetch_size,
    )


def ptx_ld_relaxed(
    addr,
    return_type,
    ptx_type,
    *,
    scope="gpu",
    space="global",
    vec="",
    dst=None,
    cache_hint="",
    cache_policy=None,
    l1_evict="",
    l2_evict="",
    prefetch_size="",
):
    _choice("scope", scope, _PTX_LD_SCOPE)
    _choice("space", space, _PTX_LD_SPACE)
    _choice("ptx_type", ptx_type, _PTX_LD_TYPE)
    _choice("vec", vec, _PTX_LD_VEC)
    cache_policy, has_cache_policy = _resolve_cache_policy(cache_hint, cache_policy)
    to_dst = int(dst is not None)
    if vec and dst is None:
        raise ValueError("vec ld.relaxed requires dst")
    if to_dst:
        return call_intrin(
            "",
            "tirx.ptx.ld_relaxed",
            dst,
            addr,
            cache_policy,
            return_type,
            scope,
            space,
            vec,
            ptx_type,
            int(has_cache_policy),
            to_dst,
            l1_evict,
            l2_evict,
            prefetch_size,
        )
    return call_intrin(
        return_type,
        "tirx.ptx.ld_relaxed",
        addr,
        cache_policy,
        return_type,
        scope,
        space,
        vec,
        ptx_type,
        int(has_cache_policy),
        0,
        l1_evict,
        l2_evict,
        prefetch_size,
    )


def ptx_ld_volatile(
    addr,
    return_type,
    ptx_type,
    *,
    space="global",
    vec="",
    dst=None,
    prefetch_size="",
):
    """TVM intrinsic for PTX ``ld.volatile{.ss}...`` loads."""
    _choice("space", space, _PTX_LD_VOLATILE_SPACE)
    _choice("ptx_type", ptx_type, _PTX_LD_TYPE)
    _choice("vec", vec, _PTX_LD_VEC)
    to_dst = int(dst is not None)
    if vec and dst is None:
        raise ValueError("vec ld.volatile requires dst")
    if to_dst:
        return call_intrin(
            "",
            "tirx.ptx.ld_volatile",
            dst,
            addr,
            return_type,
            space,
            vec,
            ptx_type,
            to_dst,
            prefetch_size,
        )
    return call_intrin(
        return_type,
        "tirx.ptx.ld_volatile",
        addr,
        return_type,
        space,
        vec,
        ptx_type,
        0,
        prefetch_size,
    )


def ptx_ld_mmio(addr, return_type, ptx_type, *, sem="acquire", dst=None):
    if sem not in ("acquire", "relaxed"):
        raise ValueError(f"Unsupported PTX ld.mmio sem {sem!r}")
    _choice("ptx_type", ptx_type, _PTX_LD_TYPE)
    to_dst = int(dst is not None)
    if to_dst:
        return call_intrin(
            "",
            "tirx.ptx.ld_mmio",
            dst,
            addr,
            return_type,
            sem,
            "sys",
            "global",
            ptx_type,
            1,
        )
    return call_intrin(
        return_type,
        "tirx.ptx.ld_mmio",
        addr,
        return_type,
        sem,
        "sys",
        "global",
        ptx_type,
        0,
    )


def ptx_ld_global_acquire(res, addr):
    """TVM intrinsic to call the legacy ptx ld.global.acquire helper.

    Parameters
    ----------
    res : Expr
        The result of the load.

    addr : Expr
        The memory address to load.

    Returns
    -------
    call : Expr
        The call expression.
    """
    return call_intrin("", "tirx.ptx.ld_global_acquire", res, addr)


def ptx_red_scalar(
    address,
    value,
    *,
    sem="",
    scope="",
    space="global",
    op,
    ptx_type,
    cache_hint="",
    cache_policy=None,
):
    _choice("scope", scope, _PTX_MEM_SCOPE)
    _choice("space", space, _PTX_MEM_SPACE)
    _choice("op", op, _PTX_RED_OP)
    _choice("ptx_type", ptx_type, _PTX_SCALAR_TYPE)
    cache_policy, has_cache_policy = _resolve_cache_policy(
        cache_hint, cache_policy, _CP_ASYNC_CACHE_HINT
    )
    if sem not in ("", "relaxed", "release"):
        raise ValueError(f"Unsupported PTX red sem {sem!r}")
    return call_intrin(
        "",
        "tirx.ptx.red_scalar",
        address,
        value,
        cache_policy,
        sem,
        scope,
        space,
        op,
        ptx_type,
        int(has_cache_policy),
    )


def ptx_atom_scalar(
    address,
    value,
    *,
    sem="",
    scope="",
    space="global",
    op,
    ptx_type,
    cache_hint="",
    cache_policy=None,
):
    _choice("scope", scope, _PTX_MEM_SCOPE)
    _choice("space", space, _PTX_MEM_SPACE)
    _choice("op", op, _PTX_ATOM_OP)
    _choice("ptx_type", ptx_type, _PTX_SCALAR_TYPE)
    cache_policy, has_cache_policy = _resolve_cache_policy(cache_hint, cache_policy)
    if sem not in ("", "relaxed", "acquire", "release", "acq_rel"):
        raise ValueError(f"Unsupported PTX atom sem {sem!r}")
    return call_intrin(
        _PTX_SCALAR_RETURN_TYPE[ptx_type],
        "tirx.ptx.atom_scalar",
        address,
        value,
        cache_policy,
        sem,
        scope,
        space,
        op,
        ptx_type,
        int(has_cache_policy),
    )


def ptx_st(
    address,
    *values,
    src=None,
    weak=False,
    space="shared",
    cop="",
    vec="",
    ptx_type,
    cache_hint="",
    cache_policy=None,
    l1_evict="",
    l2_evict="",
):
    _choice("space", space, _PTX_MEM_SPACE | {"local", "param::func"})
    _choice("cop", cop, _PTX_ST_COP)
    _choice("vec", vec, _PTX_ST_VEC)
    _choice("ptx_type", ptx_type, _PTX_SCALAR_TYPE)
    cache_policy, has_cache_policy = _resolve_cache_policy(cache_hint, cache_policy)
    attrs = (
        cache_policy,
        int(bool(weak)),
        space,
        cop,
        vec,
        ptx_type,
        int(has_cache_policy),
        l1_evict,
        l2_evict,
        int(src is not None),
    )
    if src is not None:
        if values:
            raise ValueError("ptx_st expects values or src, not both")
        return call_intrin("", "tirx.ptx.st", address, src, *attrs)
    return call_intrin("", "tirx.ptx.st", address, *values, *attrs)


def ptx_st_relaxed(
    address,
    *values,
    src=None,
    scope="gpu",
    space="global",
    vec="",
    ptx_type,
    cache_hint="",
    cache_policy=None,
    l1_evict="",
    l2_evict="",
):
    _choice("scope", scope, _PTX_LD_SCOPE)
    _choice("space", space, _PTX_MEM_SPACE | {"local", "param::func"})
    _choice("vec", vec, _PTX_ST_VEC)
    _choice("ptx_type", ptx_type, _PTX_SCALAR_TYPE)
    cache_policy, has_cache_policy = _resolve_cache_policy(cache_hint, cache_policy)
    attrs = (
        cache_policy,
        scope,
        space,
        vec,
        ptx_type,
        int(has_cache_policy),
        l1_evict,
        l2_evict,
        int(src is not None),
    )
    if src is not None:
        if values:
            raise ValueError("ptx_st_relaxed expects values or src, not both")
        return call_intrin("", "tirx.ptx.st_relaxed", address, src, *attrs)
    return call_intrin("", "tirx.ptx.st_relaxed", address, *values, *attrs)


def ptx_st_release(
    address,
    *values,
    src=None,
    scope="gpu",
    space="global",
    vec="",
    ptx_type,
    cache_hint="",
    cache_policy=None,
    l1_evict="",
    l2_evict="",
):
    _choice("scope", scope, _PTX_LD_SCOPE)
    _choice("space", space, _PTX_MEM_SPACE | {"local", "param::func"})
    _choice("vec", vec, _PTX_ST_VEC)
    _choice("ptx_type", ptx_type, _PTX_SCALAR_TYPE)
    cache_policy, has_cache_policy = _resolve_cache_policy(cache_hint, cache_policy)
    attrs = (
        cache_policy,
        scope,
        space,
        vec,
        ptx_type,
        int(has_cache_policy),
        l1_evict,
        l2_evict,
        int(src is not None),
    )
    if src is not None:
        if values:
            raise ValueError("ptx_st_release expects values or src, not both")
        return call_intrin("", "tirx.ptx.st_release", address, src, *attrs)
    return call_intrin("", "tirx.ptx.st_release", address, *values, *attrs)


def ptx_st_volatile(address, *values, src=None, space="global", vec="", ptx_type):
    _choice("space", space, _PTX_MEM_SPACE | {"local", "param::func"})
    _choice("vec", vec, _PTX_ST_VEC)
    _choice("ptx_type", ptx_type, _PTX_SCALAR_TYPE)
    attrs = (space, vec, ptx_type, int(src is not None))
    if src is not None:
        if values:
            raise ValueError("ptx_st_volatile expects values or src, not both")
        return call_intrin("", "tirx.ptx.st_volatile", address, src, *attrs)
    return call_intrin("", "tirx.ptx.st_volatile", address, *values, *attrs)


def ptx_st_mmio(address, value, *, sem="release", ptx_type):
    if sem not in ("release", "relaxed"):
        raise ValueError(f"Unsupported PTX st.mmio sem {sem!r}")
    _choice("ptx_type", ptx_type, _PTX_SCALAR_TYPE)
    return call_intrin("", "tirx.ptx.st_mmio", address, value, sem, "sys", "global", ptx_type)


def ptx_st_bulk(ptr, num_bytes, *, weak=False, space="shared::cta"):
    if space not in ("", "shared::cta"):
        raise ValueError(f"Unsupported PTX st.bulk space {space!r}")
    return call_intrin("", "tirx.ptx.st_bulk", ptr, num_bytes, int(bool(weak)), space)


def ptx_prefetch_tensormap(tensormap_addr, space=""):
    _choice("space", space, _PTX_PREFETCH_TENSORMAP_SPACE)
    return call_intrin("", "tirx.ptx.prefetch_tensormap", tensormap_addr, space)


def ptx_mbarrier_test_wait_parity(barrier, phase, *, sem="", scope="", space="shared::cta"):
    if sem not in ("", "acquire", "relaxed"):
        raise ValueError(f"Unsupported mbarrier.test_wait.parity sem {sem!r}")
    if scope not in ("", "cta", "cluster"):
        raise ValueError(f"Unsupported mbarrier.test_wait.parity scope {scope!r}")
    if bool(sem) != bool(scope):
        raise ValueError("mbarrier.test_wait.parity sem and scope must be set together")
    if space not in ("shared", "shared::cta"):
        raise ValueError(f"Unsupported mbarrier.test_wait.parity space {space!r}")
    return call_intrin(
        "uint32", "tirx.ptx.mbarrier_test_wait_parity", barrier, phase, sem, scope, space
    )


def ptx_cp_async_bulk_g2s_cta(
    dst_ptr,
    src_ptr,
    num_bytes,
    mbarrier_ptr,
    *,
    cache_hint="",
    cache_policy=None,
    ignore_oob=False,
    ignore_bytes_left=0,
    ignore_bytes_right=0,
):
    cache_policy, has_cache_policy = _resolve_cache_policy(cache_hint, cache_policy)
    return call_intrin(
        "",
        "tirx.ptx.cp_async_bulk_g2s_cta",
        dst_ptr,
        src_ptr,
        num_bytes,
        ignore_bytes_left,
        ignore_bytes_right,
        mbarrier_ptr,
        cache_policy,
        int(has_cache_policy),
        int(bool(ignore_oob)),
    )


def ptx_cp_async_bulk_g2s_cluster(
    dst_ptr,
    src_ptr,
    num_bytes,
    mbarrier_ptr,
    *,
    cache_hint="",
    cache_policy=None,
    multicast=False,
    cta_mask=0,
):
    cache_policy, has_cache_policy = _resolve_cache_policy(cache_hint, cache_policy)
    return call_intrin(
        "",
        "tirx.ptx.cp_async_bulk_g2s_cluster",
        dst_ptr,
        src_ptr,
        num_bytes,
        mbarrier_ptr,
        cta_mask,
        cache_policy,
        int(has_cache_policy),
        int(bool(multicast)),
    )


def ptx_cp_async_bulk_s2s_cluster(dst_ptr, src_ptr, num_bytes, mbarrier):
    return call_intrin(
        "", "tirx.ptx.cp_async_bulk_s2s_cluster", dst_ptr, src_ptr, num_bytes, mbarrier
    )


def ptx_cp_async_bulk_s2g(
    dst_ptr, src_ptr, num_bytes, *, cache_hint="", cache_policy=None, cp_mask=False, byte_mask=0
):
    cache_policy, has_cache_policy = _resolve_cache_policy(cache_hint, cache_policy)
    return call_intrin(
        "",
        "tirx.ptx.cp_async_bulk_s2g",
        dst_ptr,
        src_ptr,
        num_bytes,
        byte_mask,
        cache_policy,
        int(has_cache_policy),
        int(bool(cp_mask)),
    )


def ptx_fns_b32(mask, base, offset):
    return call_intrin("uint32", "tirx.ptx.fns_b32", mask, base, offset)


def ptx_add_rn_f32_bf16(acc, x):
    return call_intrin("float32", "tirx.ptx.add_rn_f32_bf16", acc, x)


def cuda_uint_as_float(bits):
    return call_intrin("float32", "tirx.cuda.uint_as_float", bits)


def cuda_float_as_uint(x):
    return call_intrin("uint32", "tirx.cuda.float_as_uint", x)


def cuda_ballot_sync(mask, pred):
    return call_intrin("uint32", "tirx.cuda.ballot_sync", mask, pred)


def cuda_ffs_u32(value):
    return call_intrin("int32", "tirx.cuda.ffs_u32", value)


def cuda_reduce_add_sync_u32(mask, value):
    return call_intrin("uint32", "tirx.cuda.reduce_add_sync_u32", mask, value)


def cuda_reduce_min_sync_u32(mask, value):
    return call_intrin("uint32", "tirx.cuda.reduce_min_sync_u32", mask, value)


def cuda_clock64():
    return call_intrin("uint64", "tirx.cuda.clock64")


def cuda_make_float2(x, y):
    return call_intrin("uint64", "tirx.cuda.make_float2", x, y)


def cuda_float2_x(packed):
    return call_intrin("float32", "tirx.cuda.float2_x", packed)


def cuda_float2_y(packed):
    return call_intrin("float32", "tirx.cuda.float2_y", packed)


def cuda_fmul2_rn(a, b):
    return call_intrin("uint64", "tirx.cuda.fmul2_rn", a, b)


def cuda_fadd2_rn(a, b):
    return call_intrin("uint64", "tirx.cuda.fadd2_rn", a, b)


def cuda_float22bfloat162_rn(v0, v1):
    return call_intrin("uint32", "tirx.cuda.float22bfloat162_rn", v0, v1)


def cuda_float22bfloat162_rn_from_float2(packed):
    return call_intrin("uint32", "tirx.cuda.float22bfloat162_rn_from_float2", packed)


def cuda_bfloat1622float2(packed):
    return call_intrin("uint64", "tirx.cuda.bfloat1622float2", packed)


def cuda_hmin2(a, b):
    return call_intrin("uint32", "tirx.cuda.hmin2", a, b)


def cuda_hmax2(a, b):
    return call_intrin("uint32", "tirx.cuda.hmax2", a, b)


def cuda_fp8x4_e4m3_from_float4(x, y, z, w):
    return call_intrin("uint32", "tirx.cuda.fp8x4_e4m3_from_float4", x, y, z, w)


def ptx_map_shared_rank(ptr, rank):
    """TVM intrinsic to call ptx map_shared_rank instruction

    Parameters
    ----------
    ptr: Expr
        The generic pointer to the local shared memory, handle type

    rank: int
        The rank of the distributed shared memory.

    Returns
    -------
    call : Expr
        The call expression.
    """

    return ptx_mapa(ptr, rank, space="", ptx_type="u64", return_type="uint64")


def ptx_mapa(ptr, rank, *, space="", ptx_type="u64", return_type="uint64"):
    """TVM intrinsic for PTX ``mapa{.space}.type d, a, b``."""
    if space not in ("", "shared::cluster"):
        raise ValueError(f"Unsupported mapa space {space!r}")
    if ptx_type not in ("u32", "u64"):
        raise ValueError(f"Unsupported mapa type {ptx_type!r}")
    return call_intrin(return_type, "tirx.ptx.mapa", ptr, rank, space, ptx_type, return_type)


def cuda_atomic_cas(ptr, old_val, new_val):
    """TVM intrinsic to call cuda atomic cas instruction

    Parameters
    ----------
    ptr: Expr
        The pointer to the memory location.

    old_val: Expr
        The old value.

    new_val: Expr
        The new value.

    Returns
    -------
    call : Expr
        The call expression.
    """
    old_val = tir.convert(old_val)
    return call_intrin(old_val.ty, "tirx.cuda.atomic_cas", ptr, old_val, new_val)


########################################################
# NVSHMEM builtins
########################################################


def nvshmem_my_pe():
    """TVM intrinsic to call nvshmem_my_pe()

    Returns
    -------
    call : Expr
        The call expression.
    """

    return call_intrin("int32", "tirx.nvshmem.my_pe")


def nvshmem_n_pes():
    """TVM intrinsic to call nvshmem_n_pes()

    Returns
    -------
    call : Expr
        The call expression.
    """

    return call_intrin("int32", "tirx.nvshmem.n_pes")


def nvshmem_getmem_nbi(dst, src, nelems, pe):
    """TVM intrinsic to call nvshmem_getmem_nbi()

    Parameters
    ----------
    dst: Expr
        The pointer to the symmetric address or host/device address of the data object to be updated.

    src: Expr
        The pointer to the symmetric address of the source data object.

    nelems: int
        The number of bytes to get per thread.

    pe: int
        The PE number of the remote PE.

    Returns
    -------
    call : Expr
        The call expression.
    """  # noqa: E501

    return call_intrin("", "tirx.nvshmem.getmem_nbi", dst, src, nelems, pe)


def nvshmem_putmem_nbi(dst, src, nelems, pe):
    """TVM intrinsic to call nvshmem_putmem_nbi()

    Parameters
    ----------
    dst: Expr
        The pointer to the symmetric address of the destination data object.

    src: Expr
        The pointer to the symmetric address or host/device address of the data object to be copied.

    nelems: int
        The number of bytes to put per thread.

    pe: int
        The PE number of the remote PE.

    Returns
    -------
    call : Expr
        The call expression.
    """

    return call_intrin("", "tirx.nvshmem.putmem_nbi", dst, src, nelems, pe)


def nvshmem_getmem_nbi_warp(dst, src, nelems, pe):
    """TVM intrinsic to call nvshmem_getmem_nbi_warp()

    Parameters
    ----------
    dst: Expr
        The pointer to the symmetric address or host/device address of the data object to be updated.

    src: Expr
        The pointer to the symmetric address of the source data object.

    nelems: int
        The number of bytes to get per warp.

    pe: int
        The PE number of the remote PE.

    Returns
    -------
    call : Expr
        The call expression.
    """  # noqa: E501

    return call_intrin("", "tirx.nvshmem.getmem_nbi_warp", dst, src, nelems, pe)


def nvshmem_putmem_nbi_warp(dst, src, nelems, pe):
    """TVM intrinsic to call nvshmem_putmem_nbi_warp()

    Parameters
    ----------
    dst: Expr
        The pointer to the symmetric address of the destination data object.

    src: Expr
        The pointer to the symmetric address or host/device address of the data object to be copied.

    nelems: int
        The number of bytes to put per warp.

    pe: int
        The PE number of the remote PE.

    Returns
    -------
    call : Expr
        The call expression.
    """

    return call_intrin("", "tirx.nvshmem.putmem_nbi_warp", dst, src, nelems, pe)


def nvshmem_getmem_nbi_block(dst, src, nelems, pe):
    """TVM intrinsic to call nvshmem_getmem_nbi_block()

    Parameters
    ----------
    dst: Expr
        The pointer to the symmetric address or host/device address of the data object to be updated.

    src: Expr
        The pointer to the symmetric address of the source data object.

    nelems: int
        The number of bytes to get per block.

    pe: int
        The PE number of the remote PE.

    Returns
    -------
    call : Expr
        The call expression.
    """  # noqa: E501

    return call_intrin("", "tirx.nvshmem.getmem_nbi_block", dst, src, nelems, pe)


def nvshmem_putmem_nbi_block(dst, src, nelems, pe):
    """TVM intrinsic to call nvshmem_putmem_nbi_block()

    Parameters
    ----------
    dst: Expr
        The pointer to the symmetric address of the destination data object.

    src: Expr
        The pointer to the symmetric address or host/device address of the data object to be copied.

    nelems: int
        The number of bytes to put per block.

    pe: int
        The PE number of the remote PE.

    Returns
    -------
    call : Expr
        The call expression.
    """

    return call_intrin("", "tirx.nvshmem.putmem_nbi_block", dst, src, nelems, pe)


def nvshmem_signal_op(sig_addr, signal, sig_op, pe):
    """TVM intrinsic to call nvshmem_signal_op()

    Parameters
    ----------
    sig_addr: Expr
        The pointer to the symmetric address of the signal word to be updated, must be uint64_t*.

    signal: uint64_t
        The value used to update sig_addr.

    sig_op: str
        Operation used to update sig_addr with signal, typical sig_op values are "set" and "add".

    pe: int
        The PE number of the remote PE.

    Returns
    -------
    call : Expr
        The call expression.
    """

    _choice("sig_op", sig_op, _NVSHMEM_SIG_OP)
    return call_intrin("", "tirx.nvshmem.signal_op", sig_addr, signal, sig_op, pe)


def nvshmem_wait_until(ivar, cmp, cmp_value, type="uint64_t"):
    """TVM intrinsic to call nvshmem_wait_until()

    Parameters
    ----------
    ivar: Expr
        The pointer to the symmetric address of a remotely accessible data object, must be TYPE*.

    cmp: str
        The compare operator that compares ivar with cmp_value.

    cmp_value: TYPE
        The value to be compared with ivar.

    type: str
        The TYPE of ivar and cmp_value.

    Returns
    -------
    call : Expr
        The call expression.
    """

    _choice("cmp", cmp, _NVSHMEM_CMP)
    return call_intrin("", "tirx.nvshmem.wait_until", ivar, cmp, cmp_value, type)


def nvshmem_quiet():
    """TVM intrinsic to call nvshmem_quiet()

    Returns
    -------
    call : Expr
        The call expression.
    """

    return call_intrin("", "tirx.nvshmem.quiet")


def nvshmem_putmem_signal_nbi(dst, src, nelems, sig_addr, signal, sig_op, pe):
    """TVM intrinsic to call nvshmem_putmem_signal_nbi()

    Parameters
    ----------
    dst: Expr
        The pointer to the symmetric address of the data object to be updated on the remote PE.

    src: Expr
        The pointer to the symmetric address or host/device address of data object containing the data to be copied.

    nelems: int
        The number of bytes to put per thread.

    sig_addr: Expr
        The pointer to the symmetric address of the signal data object to be updated on the remote PE as a signal, must be uint64_t*.

    signal: uint64_t
        The unsigned 64-bit value that is used for updating the remote sig_addr signal data object.

    sig_op: str
        Signal operator that represents the type of update to be performed on the remote sig_addr signal data object.

    pe: int
        The PE number of the remote PE.

    Returns
    -------
    call : Expr
        The call expression.
    """  # noqa: E501

    return call_intrin(
        "", "tirx.nvshmem.putmem_signal_nbi", dst, src, nelems, sig_addr, signal, sig_op, pe
    )


def nvshmem_putmem_signal_nbi_warp(dst, src, nelems, sig_addr, signal, sig_op, pe):
    """TVM intrinsic to call nvshmem_putmem_signal_nbi_warp()

    Parameters
    ----------
    dst: Expr
        The pointer to the symmetric address of the data object to be updated on the remote PE.

    src: Expr
        The pointer to the symmetric address or host/device address of data object containing the data to be copied.

    nelems: int
        The number of bytes to put per warp.

    sig_addr: Expr
        The pointer to the symmetric address of the signal data object to be updated on the remote PE as a signal, must be uint64_t*.

    signal: uint64_t
        The unsigned 64-bit value that is used for updating the remote sig_addr signal data object.

    sig_op: str
        Signal operator that represents the type of update to be performed on the remote sig_addr signal data object.

    pe: int
        The PE number of the remote PE.

    Returns
    -------
    call : Expr
        The call expression.
    """  # noqa: E501

    return call_intrin(
        "", "tirx.nvshmem.putmem_signal_nbi_warp", dst, src, nelems, sig_addr, signal, sig_op, pe
    )


def nvshmem_putmem_signal_nbi_block(dst, src, nelems, sig_addr, signal, sig_op, pe):
    """TVM intrinsic to call nvshmem_putmem_signal_nbi_block()

    Parameters
    ----------
    dst: Expr
        The pointer to the symmetric address of the data object to be updated on the remote PE.

    src: Expr
        The pointer to the symmetric address or host/device address of data object containing the data to be copied.

    nelems: int
        The number of bytes to put per block.

    sig_addr: Expr
        The pointer to the symmetric address of the signal data object to be updated on the remote PE as a signal, must be uint64_t*.

    signal: uint64_t
        The unsigned 64-bit value that is used for updating the remote sig_addr signal data object.

    sig_op: str
        Signal operator that represents the type of update to be performed on the remote sig_addr signal data object.

    pe: int
        The PE number of the remote PE.

    Returns
    -------
    call : Expr
        The call expression.
    """  # noqa: E501

    return call_intrin(
        "", "tirx.nvshmem.putmem_signal_nbi_block", dst, src, nelems, sig_addr, signal, sig_op, pe
    )


def nvshmem_fence():
    """TVM intrinsic to call nvshmem_fence()

    Returns
    -------
    call : Expr
        The call expression.
    """

    return call_intrin("", "tirx.nvshmem.fence")


def nvshmem_barrier_all():
    """TVM intrinsic to call nvshmem_barrier_all()

    Returns
    -------
    call : Expr
        The call expression.
    """

    return call_intrin("", "tirx.nvshmem.barrier_all")
