# 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. """CUDA TVMScript namespaces.""" from __future__ import annotations from collections.abc import Callable from typing import Any from tvm.backend.cuda import op as _cuda_op from tvm.tirx import Buffer from tvm.tirx import op as _tir_op from tvm.tirx.script.builder.ir import _dtype_forward, _op_wrapper # pylint: disable=protected-access def _ptx_ldg32(reg, guard, addr, local_addr): if isinstance(addr, Buffer): addr = addr[0] return _tir_op.call_intrin(reg.ty, "tirx.ptx.ldg32", reg, guard, addr, local_addr) _ptx_ldg32.__tir_op_name__ = "ptx.ldg32" class PTXNamespace: """The PTX instruction submodule.""" def __init__(self): self.ldg32 = _ptx_ldg32 self.ldmatrix = _dtype_forward(_cuda_op.ptx_ldmatrix) # Apache-compatible variant. Same lowered intrinsic as # ``ldmatrix`` but accepts the historical ``(trans, num, dtype, # local_ptr, local_offset, smem_ptr, smem_offset)`` form. Coexists # with the fork-native version so upstream-derived tests keep # working without rewriting their tirx code. self.ldmatrix_legacy = _dtype_forward(_cuda_op.ptx_ldmatrix_legacy) self.stmatrix = _op_wrapper(_cuda_op.ptx_stmatrix) self.setmaxnreg: Callable[..., Any] = _op_wrapper(_cuda_op.ptx_setmaxnreg) self.elect_sync: Callable[..., Any] = _op_wrapper(_cuda_op.ptx_elect_sync) self.clc_try_cancel = _op_wrapper(_cuda_op.ptx_clc_try_cancel) self.clc_query_cancel = _op_wrapper(_cuda_op.ptx_clc_query_cancel) self.fetch_register: Callable[..., Any] = _op_wrapper(_cuda_op.ptx_fetch_register) self.ld = _op_wrapper(_cuda_op.ptx_ld) self.ld_acquire = _op_wrapper(_cuda_op.ptx_ld_acquire) self.ld_relaxed = _op_wrapper(_cuda_op.ptx_ld_relaxed) self.ld_volatile = _op_wrapper(_cuda_op.ptx_ld_volatile) self.ld_mmio = _op_wrapper(_cuda_op.ptx_ld_mmio) self.ld_global_acquire = _op_wrapper(_cuda_op.ptx_ld_global_acquire) self.red_scalar = _op_wrapper(_cuda_op.ptx_red_scalar) self.atom_scalar = _op_wrapper(_cuda_op.ptx_atom_scalar) self.prefetch_tensormap = _op_wrapper(_cuda_op.ptx_prefetch_tensormap) self.mbarrier_test_wait_parity = _op_wrapper(_cuda_op.ptx_mbarrier_test_wait_parity) self.cp_async_bulk_g2s_cta = _op_wrapper(_cuda_op.ptx_cp_async_bulk_g2s_cta) self.cp_async_bulk_g2s_cluster = _op_wrapper(_cuda_op.ptx_cp_async_bulk_g2s_cluster) self.cp_async_bulk_s2s_cluster = _op_wrapper(_cuda_op.ptx_cp_async_bulk_s2s_cluster) self.cp_async_bulk_s2g = _op_wrapper(_cuda_op.ptx_cp_async_bulk_s2g) self.st = _op_wrapper(_cuda_op.ptx_st) self.st_relaxed = _op_wrapper(_cuda_op.ptx_st_relaxed) self.st_release = _op_wrapper(_cuda_op.ptx_st_release) self.st_volatile = _op_wrapper(_cuda_op.ptx_st_volatile) self.st_mmio = _op_wrapper(_cuda_op.ptx_st_mmio) self.st_bulk = _op_wrapper(_cuda_op.ptx_st_bulk) self.fns_b32 = _op_wrapper(_cuda_op.ptx_fns_b32) self.add_rn_f32_bf16 = _op_wrapper(_cuda_op.ptx_add_rn_f32_bf16) self.mapa = _op_wrapper(_cuda_op.ptx_mapa) self.map_shared_rank = _op_wrapper(_cuda_op.ptx_map_shared_rank) self.any_sync = _op_wrapper(_cuda_op.ptx_any_sync) # Math operations self.exp2 = _op_wrapper(_cuda_op.ptx_exp2) self.rcp = _op_wrapper(_cuda_op.ptx_rcp) self.reduce3_min_f32 = _op_wrapper(_cuda_op.ptx_reduce3_min_f32) self.reduce3_max_f32 = _op_wrapper(_cuda_op.ptx_reduce3_max_f32) # add/sub/mul/fma DPS form: (d_addr, a, b[, c], *, rounding, ftz[, sat]) self.add_f32 = _op_wrapper(_cuda_op.ptx_add_f32) self.add_f32x2 = _op_wrapper(_cuda_op.ptx_add_f32x2) self.add_f64 = _op_wrapper(_cuda_op.ptx_add_f64) self.sub_f32 = _op_wrapper(_cuda_op.ptx_sub_f32) self.sub_f32x2 = _op_wrapper(_cuda_op.ptx_sub_f32x2) self.sub_f64 = _op_wrapper(_cuda_op.ptx_sub_f64) self.mul_f32 = _op_wrapper(_cuda_op.ptx_mul_f32) self.mul_f32x2 = _op_wrapper(_cuda_op.ptx_mul_f32x2) self.mul_f64 = _op_wrapper(_cuda_op.ptx_mul_f64) self.fma_f32 = _op_wrapper(_cuda_op.ptx_fma_f32) self.fma_f32x2 = _op_wrapper(_cuda_op.ptx_fma_f32x2) self.fma_f64 = _op_wrapper(_cuda_op.ptx_fma_f64) self.max_f32 = _op_wrapper(_cuda_op.ptx_max_f32) self.mma = MmaNamespace() self.cp_async = CpAsyncNamespace() self.wgmma = WgmmaNamespace() self.mbarrier = MbarrierNamespace() self.tcgen05 = Tcgen05Namespace() self.bar = BarNamespace() self.barrier = BarrierNamespace() self.fence = FenceNamespace() self.griddepcontrol = GriddepcontrolNamespace() class MmaNamespace: """The MMA instruction submodule.""" def __init__(self): self.sp = _dtype_forward(_cuda_op.ptx_mma_sp) # Apache-compatible variant of ptx_mma. Coexists with the # fork-native ``__call__`` form (``T.ptx.mma(...)``). self.legacy = _dtype_forward(_cuda_op.ptx_mma_legacy) # __call__ corresponds to ptx_mma self.__tir_call_op_name__ = "ptx_mma" def __call__(self, *args, **kwds): return _dtype_forward(_cuda_op.ptx_mma)(*args, **kwds) class CpAsyncNamespace: """The CpAsync instruction submodule.""" def __init__(self): self.commit_group = _op_wrapper(_cuda_op.ptx_cp_async_commit_group) self.wait_group = _op_wrapper(_cuda_op.ptx_cp_async_wait_group) # Legacy variant: takes (dst_ptr, dst_offset, src_ptr, src_offset, # cp_size). Offsets are folded into the pointers; coexists with # the fork-native ``__call__`` form. self.legacy = _dtype_forward(_cuda_op.ptx_cp_async_legacy) self.bulk = CpAsyncBulkNamespace() self.mbarrier = CpAsyncMbarrierNamespace() def __call__(self, *args, **kwds): # Accept the legacy 6-arg form ``(elem_dtype, dst, dst_off, src, # src_off, cp_size)`` that the printer round-trips for the raw # ``tirx.ptx.cp_async`` Call emitted by # ``tvm.backend.cuda.transform.InjectPTXAsyncCopy``. The pass-emitted # Call has 5 args (no ``tvm_access_ptr`` fold) and a # per-element-dtype Call.dtype, so build it directly. if len(args) == 6 and isinstance(args[0], str) and "dtype" not in kwds: import tvm elem_dtype, dst, dst_off, src, src_off, cp_size = args return tvm.ir.Call( tvm.ir.Op.get("tirx.ptx.cp_async_raw"), [dst, dst_off, src, src_off, cp_size], ret_ty=tvm.ir.PrimType(elem_dtype), ) return _dtype_forward(_cuda_op.ptx_cp_async)(*args, **kwds) # __call__ corresponds to ptx_cp_async __tir_call_op_name__ = "ptx_cp_async" class CpAsyncBulkNamespace: """The CpAsyncBulk instruction submodule.""" def __init__(self): self.commit_group = _op_wrapper(_cuda_op.ptx_cp_async_bulk_commit_group) self.wait_group = _op_wrapper(_cuda_op.ptx_cp_async_bulk_wait_group) self.tensor = CpAsyncBulkTensorNamespace() self.s2c = _op_wrapper(_cuda_op.ptx_cp_async_bulk_shared_to_cluster) def __call__(self, *args, **kwds): return _dtype_forward(_cuda_op.ptx_cp_async_bulk)(*args, **kwds) # __call__ corresponds to ptx_cp_async_bulk __tir_call_op_name__ = "ptx_cp_async_bulk" class CpAsyncBulkTensorNamespace: """The CpAsyncBulkTensor instruction submodule.""" def __init__(self): self.g2c = _op_wrapper(_cuda_op.ptx_cp_async_bulk_tensor_global_to_cluster) self.g2c_tile_gather4 = _op_wrapper( _cuda_op.ptx_cp_async_bulk_tensor_tile_gather4_global_to_cluster ) self.s2g = _op_wrapper(_cuda_op.ptx_cp_async_bulk_tensor_shared_to_global) self.s2g_reduce = _op_wrapper(_cuda_op.ptx_cp_async_bulk_tensor_shared_to_global_reduce) self.g2c_prefetch = _op_wrapper( _cuda_op.ptx_cp_async_bulk_tensor_global_to_cluster_prefetch ) @staticmethod def g2c_bar_addr( dim, dst_ptr, bar_addr, tensormap_addr, cta_mask, cta_group, cache_hint, *coords, cache_policy=None, ): _cuda_op._choice("cta_group", cta_group, _cuda_op._TCGEN05_CTA_GROUP) cache_policy, has_cache_policy = _cuda_op._resolve_cache_policy(cache_hint, cache_policy) return _tir_op.call_intrin( "", "tirx.ptx.cp_async_bulk_tensor_global_to_cluster", dim, dst_ptr, bar_addr, tensormap_addr, cta_mask, cta_group, cache_policy, int(has_cache_policy), 1, *coords, ) @staticmethod def g2c_tile_gather4_bar_addr( dim, dst_ptr, bar_addr, tensormap_addr, cta_mask, cta_group, cache_hint, *coords, cache_policy=None, ): _cuda_op._choice("cta_group", cta_group, _cuda_op._TCGEN05_CTA_GROUP) cache_policy, has_cache_policy = _cuda_op._resolve_cache_policy(cache_hint, cache_policy) return _tir_op.call_intrin( "", "tirx.ptx.cp_async_bulk_tensor_tile_gather4_global_to_cluster", dim, dst_ptr, bar_addr, tensormap_addr, cta_mask, cta_group, cache_policy, int(has_cache_policy), 1, *coords, ) class CpAsyncMbarrierNamespace: """The CpAsyncMbarrier instruction submodule.""" def __init__(self): self.arrive = _op_wrapper(_cuda_op.ptx_cp_async_mbarrier_arrive) class WgmmaNamespace: """The WGMMA instruction submodule.""" def __init__(self): self.fence: Callable[..., Any] = _op_wrapper(_cuda_op.ptx_wgmma_fence) self.commit_group = _op_wrapper(_cuda_op.ptx_wgmma_commit_group) self.wait_group = _op_wrapper(_cuda_op.ptx_wgmma_wait_group) self.noop_barrier = _op_wrapper(_cuda_op.ptx_wgmma_noop_barrier) self.mma_async = WgmmaMmaAsyncNamespace() self.encode_matrix_descriptor = _op_wrapper(_cuda_op.ptx_wgmma_encode_matrix_descriptor) class WgmmaMmaAsyncNamespace: """The WGMMA MMAAsync instruction submodule.""" def __init__(self): self.ss = _op_wrapper(_cuda_op.ptx_wgmma_mma_async_ss) self.rs = _op_wrapper(_cuda_op.ptx_wgmma_mma_async_rs) class MbarrierNamespace: """The Mbarrier instruction submodule.""" def __init__(self): self.init = _op_wrapper(_cuda_op.ptx_mbarrier_init) self.try_wait = _op_wrapper(_cuda_op.ptx_mbarrier_try_wait) self.try_wait_once = _op_wrapper(_cuda_op.ptx_mbarrier_try_wait_once) self.try_wait_acquire_cluster = _op_wrapper(_cuda_op.ptx_mbarrier_try_wait_acquire_cluster) self.arrive = MbarrierArriveNamespace() class MbarrierArriveNamespace: """The Mbarrier Arrive instruction submodule.""" def __init__(self): self.expect_tx = _op_wrapper(_cuda_op.ptx_mbarrier_arrive_expect_tx) self.cluster_count = _op_wrapper(_cuda_op.ptx_mbarrier_arrive_cluster_count) def __call__(self, *args, **kwds): return _op_wrapper(_cuda_op.ptx_mbarrier_arrive)(*args, **kwds) # __call__ corresponds to ptx_mbarrier_arrive __tir_call_op_name__ = "ptx_mbarrier_arrive" class Tcgen05Namespace: """The Tcgen05 instruction submodule.""" def __init__(self): self.alloc = _op_wrapper(_cuda_op.ptx_tcgen05_alloc) self.dealloc = _op_wrapper(_cuda_op.ptx_tcgen05_dealloc) self.relinquish_alloc_permit = _op_wrapper(_cuda_op.ptx_tcgen05_relinquish_alloc_permit) self.encode_matrix_descriptor = _op_wrapper(_cuda_op.ptx_tcgen05_encode_matrix_descriptor) self.encode_instr_descriptor = _op_wrapper(_cuda_op.ptx_tcgen05_encode_instr_descriptor) self.encode_instr_descriptor_block_scaled = _op_wrapper( _cuda_op.ptx_tcgen05_encode_instr_descriptor_block_scaled ) self.ld = _op_wrapper(_cuda_op.ptx_tcgen05_ld) self.st = _op_wrapper(_cuda_op.ptx_tcgen05_st) self.cp = _op_wrapper(_cuda_op.ptx_tcgen05_cp) self.shift = _op_wrapper(_cuda_op.ptx_tcgen05_shift) self.commit = _op_wrapper(_cuda_op.ptx_tcgen05_commit) self.wait = Tcgen05WaitNamespace() self.mma = Tcgen05MmaNamespace() self.fence = Tcgen05FenceNamespace() class Tcgen05FenceNamespace: """The Tcgen05 Fence instruction submodule.""" def __init__(self): self.before_thread_sync = _op_wrapper(_cuda_op.ptx_tcgen05_fence_before_thread_sync) self.after_thread_sync = _op_wrapper(_cuda_op.ptx_tcgen05_fence_after_thread_sync) class Tcgen05MmaNamespace: """The Tcgen05 MMA instruction submodule.""" def __init__(self): self.block_scale = _op_wrapper(_cuda_op.ptx_tcgen05_mma_block_scale) self.sp = Tcgen05MmaSpNamespace() def __call__(self, *args, **kwds): return _op_wrapper(_cuda_op.ptx_tcgen05_mma)(*args, **kwds) # __call__ corresponds to ptx_tcgen05_mma __tir_call_op_name__ = "ptx_tcgen05_mma" class Tcgen05MmaSpNamespace: """Tcgen05 Sparse MMA instruction submodule.""" def __init__(self): self.block_scale = _op_wrapper(_cuda_op.ptx_tcgen05_mma_sp_block_scale) def __call__(self, *args, **kwds): return _op_wrapper(_cuda_op.ptx_tcgen05_mma_sp)(*args, **kwds) # __call__ corresponds to ptx_tcgen05_mma_sp __tir_call_op_name__ = "ptx_tcgen05_mma_sp" class Tcgen05WaitNamespace: """The Tcgen05 Wait instruction submodule.""" def __init__(self): self.ld = _op_wrapper(_cuda_op.ptx_tcgen05_wait_ld) self.st = _op_wrapper(_cuda_op.ptx_tcgen05_wait_st) class BarNamespace: """The Bar instruction submodule.""" def __init__(self): self.arrive = _op_wrapper(_cuda_op.ptx_bar_arrive) self.sync = _op_wrapper(_cuda_op.ptx_bar_sync) class BarrierNamespace: """The Barrier instruction submodule.""" def __init__(self): self.cluster = BarrierClusterNamespace() class BarrierClusterNamespace: """The BarrierCluster instruction submodule.""" def __init__(self): self.arrive = _op_wrapper(_cuda_op.ptx_barrier_cluster_arrive) self.wait = _op_wrapper(_cuda_op.ptx_barrier_cluster_wait) class FenceNamespace: """PTX fence instruction submodule.""" def __init__(self): self.proxy_async = _op_wrapper(_cuda_op.ptx_fence_proxy_async) self.mbarrier_init = _op_wrapper(_cuda_op.ptx_fence_mbarrier_init) def __call__(self, *args, **kwds): return _op_wrapper(_cuda_op.ptx_fence)(*args, **kwds) __tir_call_op_name__ = "ptx_fence" class GriddepcontrolNamespace: """PTX griddepcontrol instruction submodule (sm_90+).""" def __init__(self): self.wait = _op_wrapper(_cuda_op.ptx_griddepcontrol_wait) self.launch_dependents = _op_wrapper(_cuda_op.ptx_griddepcontrol_launch_dependents) class CUDANamespace: """The CUDA intrinsics submodule.""" def __init__(self): self.atomic_add = _op_wrapper(_cuda_op.cuda_atomic_add) self.thread_fence = _op_wrapper(_cuda_op.cuda_thread_fence) self.warpgroup_sync = _op_wrapper(_cuda_op.cuda_warpgroup_sync) self.warp_sync = _op_wrapper(_cuda_op.cuda_warp_sync) self.warp_reduce = _op_wrapper(_cuda_op.cuda_warp_reduce) self.warp_sum = _op_wrapper(_cuda_op.cuda_warp_sum) self.warp_max = _op_wrapper(_cuda_op.cuda_warp_max) self.warp_min = _op_wrapper(_cuda_op.cuda_warp_min) self.cta_reduce = _op_wrapper(_cuda_op.cuda_cta_reduce) self.cta_sum = _op_wrapper(_cuda_op.cuda_cta_sum) self.cta_max = _op_wrapper(_cuda_op.cuda_cta_max) self.cta_min = _op_wrapper(_cuda_op.cuda_cta_min) self.cta_sync = _op_wrapper(_cuda_op.cuda_cta_sync) self.grid_sync = _op_wrapper(_cuda_op.cuda_grid_sync) self.cluster_sync = _op_wrapper(_cuda_op.cuda_cluster_sync) self.thread_rank = _op_wrapper(_cuda_op.cuda_thread_rank) self.trap_when_assert_failed = _op_wrapper(_cuda_op.cuda_trap_when_assert_failed) self.runtime_instr_desc = _op_wrapper(_cuda_op.cuda_runtime_instr_desc) self.half2float = _op_wrapper(_cuda_op.cuda_half2float) self.bfloat162float = _op_wrapper(_cuda_op.cuda_bfloat162float) self.float22half2 = _op_wrapper(_cuda_op.cuda_float22half2) self.half8tofloat8 = _op_wrapper(_cuda_op.cuda_half8tofloat8) self.float8tohalf8 = _op_wrapper(_cuda_op.cuda_float8tohalf8) self.syncthreads_and = _op_wrapper(_cuda_op.cuda_syncthreads_and) self.syncthreads_or = _op_wrapper(_cuda_op.cuda_syncthreads_or) self.nano_sleep = _op_wrapper(_cuda_op.cuda_nano_sleep) self.atomic_cas = _op_wrapper(_cuda_op.cuda_atomic_cas) self.func_call = _op_wrapper(_cuda_op.cuda_func_call) self.printf = _op_wrapper(_cuda_op.cuda_printf) self.ldg = _op_wrapper(_cuda_op.cuda_ldg) self.get_tmem_addr = _op_wrapper(_cuda_op.cuda_get_tmem_addr) self.cvta_generic_to_shared = _op_wrapper(_cuda_op.cuda_cvta_generic_to_shared) self.smem_addr_from_uint64 = _op_wrapper(_cuda_op.cuda_smem_addr_from_uint64) self.sm100_tma_2sm_mbarrier_addr = _op_wrapper(_cuda_op.cuda_sm100_tma_2sm_mbarrier_addr) self.uint_as_float = _op_wrapper(_cuda_op.cuda_uint_as_float) self.float_as_uint = _op_wrapper(_cuda_op.cuda_float_as_uint) self.ballot_sync = _op_wrapper(_cuda_op.cuda_ballot_sync) self.ffs_u32 = _op_wrapper(_cuda_op.cuda_ffs_u32) self.reduce_add_sync_u32 = _op_wrapper(_cuda_op.cuda_reduce_add_sync_u32) self.reduce_min_sync_u32 = _op_wrapper(_cuda_op.cuda_reduce_min_sync_u32) self.clock64 = _op_wrapper(_cuda_op.cuda_clock64) self.make_float2 = _op_wrapper(_cuda_op.cuda_make_float2) self.float2_x = _op_wrapper(_cuda_op.cuda_float2_x) self.float2_y = _op_wrapper(_cuda_op.cuda_float2_y) self.fmul2_rn = _op_wrapper(_cuda_op.cuda_fmul2_rn) self.fadd2_rn = _op_wrapper(_cuda_op.cuda_fadd2_rn) self.float22bfloat162_rn = _op_wrapper(_cuda_op.cuda_float22bfloat162_rn) self.float22bfloat162_rn_from_float2 = _op_wrapper( _cuda_op.cuda_float22bfloat162_rn_from_float2 ) self.bfloat1622float2 = _op_wrapper(_cuda_op.cuda_bfloat1622float2) self.hmin2 = _op_wrapper(_cuda_op.cuda_hmin2) self.hmax2 = _op_wrapper(_cuda_op.cuda_hmax2) self.fp8x4_e4m3_from_float4 = _op_wrapper(_cuda_op.cuda_fp8x4_e4m3_from_float4) self.timer_init = _op_wrapper(_cuda_op.timer_init_cuda) self.timer_start = _op_wrapper(_cuda_op.timer_start_cuda) self.timer_end = _op_wrapper(_cuda_op.timer_end_cuda) self.timer_finalize = _op_wrapper(_cuda_op.timer_finalize_cuda) self.mma_store = _dtype_forward(_cuda_op.mma_store) self.mma_fill = _dtype_forward(_cuda_op.mma_fill) self.mma_store_legacy = _dtype_forward(_cuda_op.mma_store_legacy) self.mma_fill_legacy = _dtype_forward(_cuda_op.mma_fill_legacy) setattr(self, "__shfl_sync", self._shfl_sync) setattr(self, "__shfl_up_sync", self._shfl_up_sync) setattr(self, "__shfl_down_sync", self._shfl_down_sync) setattr(self, "__shfl_xor_sync", self._shfl_xor_sync) setattr(self, "__activemask", self._activemask) @staticmethod def _shfl_sync(mask, var, lane, width): if isinstance(var, Buffer): var = var[0] return _tir_op.call_intrin(var.ty, "tirx.cuda.__shfl_sync", mask, var, lane, width) @staticmethod def _shfl_up_sync(mask, var, delta, width): if isinstance(var, Buffer): var = var[0] return _tir_op.call_intrin(var.ty, "tirx.cuda.__shfl_up_sync", mask, var, delta, width) @staticmethod def _shfl_down_sync(mask, var, delta, width): if isinstance(var, Buffer): var = var[0] return _tir_op.call_intrin(var.ty, "tirx.cuda.__shfl_down_sync", mask, var, delta, width) @staticmethod def _shfl_xor_sync(mask, var, lane_mask, width): if isinstance(var, Buffer): var = var[0] return _tir_op.call_intrin(var.ty, "tirx.cuda.__shfl_xor_sync", mask, var, lane_mask, width) @staticmethod def _activemask(): return _tir_op.call_intrin("uint32", "tirx.cuda.__activemask") class NVSHMEMNamespace: """The NVSHMEM intrinsics submodule.""" def __init__(self): self.my_pe = _op_wrapper(_cuda_op.nvshmem_my_pe) self.n_pes = _op_wrapper(_cuda_op.nvshmem_n_pes) self.signal_op = _op_wrapper(_cuda_op.nvshmem_signal_op) self.wait_until = _op_wrapper(_cuda_op.nvshmem_wait_until) self.quiet = _op_wrapper(_cuda_op.nvshmem_quiet) self.fence = _op_wrapper(_cuda_op.nvshmem_fence) self.barrier_all = _op_wrapper(_cuda_op.nvshmem_barrier_all) self.getmem_nbi = NVSHMEMGetMemNBINamespace() self.putmem_nbi = NVSHMEMPutMemNBINamespace() self.putmem_signal_nbi = NVSHMEMPutMemSignalNBINamespace() class NVSHMEMGetMemNBINamespace: """The NVSHMEM GetMemNBI intrinsics submodule.""" def __init__(self): self.warp = _op_wrapper(_cuda_op.nvshmem_getmem_nbi_warp) self.block = _op_wrapper(_cuda_op.nvshmem_getmem_nbi_block) def __call__(self, *args, **kwds): return _op_wrapper(_cuda_op.nvshmem_getmem_nbi)(*args, **kwds) # __call__ corresponds to nvshmem_getmem_nbi __tir_call_op_name__ = "nvshmem_getmem_nbi" class NVSHMEMPutMemNBINamespace: """The NVSHMEM PutMemNBI intrinsics submodule.""" def __init__(self): self.warp = _op_wrapper(_cuda_op.nvshmem_putmem_nbi_warp) self.block = _op_wrapper(_cuda_op.nvshmem_putmem_nbi_block) def __call__(self, *args, **kwds): return _op_wrapper(_cuda_op.nvshmem_putmem_nbi)(*args, **kwds) # __call__ corresponds to nvshmem_putmem_nbi __tir_call_op_name__ = "nvshmem_putmem_nbi" class NVSHMEMPutMemSignalNBINamespace: """The NVSHMEM PutMemSignalNBI intrinsics submodule.""" def __init__(self): self.warp = _op_wrapper(_cuda_op.nvshmem_putmem_signal_nbi_warp) self.block = _op_wrapper(_cuda_op.nvshmem_putmem_signal_nbi_block) def __call__(self, *args, **kwds): return _op_wrapper(_cuda_op.nvshmem_putmem_signal_nbi)(*args, **kwds) # __call__ corresponds to nvshmem_putmem_signal_nbi __tir_call_op_name__ = "nvshmem_putmem_signal_nbi" __all__ = ["CUDANamespace", "NVSHMEMNamespace", "PTXNamespace"]