wehub-resource-sync
574 lines
23 KiB
Python
574 lines
23 KiB
Python
# 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"]
|