wehub/apache--tvm
1
0
Fork 0
Code Issues Pull Requests Actions 6 Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
26446540fadaf1a5bdd5dfc29e594b312da09221
apache--tvm/python/tvm/backend/cuda/op.py
T
wehub-resource-sync 26446540fa
Lint / lint (push) Waiting to run
Details
CI / MacOS (push) Waiting to run
Details
CI / Windows (push) Waiting to run
Details
chore: import upstream snapshot with attribution
2026-07-13 13:36:25 +08:00

4628 lines
123 KiB
Python
Raw Blame History

# 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")
Reference in New Issue View Git Blame Copy Permalink