wehub-resource-sync
162 lines
6.3 KiB
Python
162 lines
6.3 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.
|
|
# pylint: disable=redefined-builtin, invalid-name
|
|
"""NVSHMEM intrinsics. Each backend call is one ``device_intrinsic(...)``."""
|
|
|
|
from ._schema import device_intrinsic
|
|
from .registry import CODEGEN_REGISTRY, register_codegen
|
|
|
|
_NVSHMEM = ("nvshmem",)
|
|
|
|
# =============================================================================
|
|
# No-arg helpers: PE queries, quiet, fence, barrier_all.
|
|
# =============================================================================
|
|
for _op, _call, _ret, _tvm_ret in [
|
|
("nvshmem_my_pe", "nvshmem_my_pe", "int32_t", "int32"),
|
|
("nvshmem_n_pes", "nvshmem_n_pes", "int32_t", "int32"),
|
|
("nvshmem_quiet", "nvshmem_quiet", "void", None),
|
|
("nvshmem_fence", "nvshmem_fence", "void", None),
|
|
("nvshmem_barrier_all", "nvshmem_barrier_all", "void", None),
|
|
]:
|
|
device_intrinsic(
|
|
_op,
|
|
body=(" " + (f"return {_call}();" if _ret != "void" else f"{_call}();")),
|
|
return_type=_ret,
|
|
tvm_return_type=_tvm_ret,
|
|
extra_deps=_NVSHMEM,
|
|
)
|
|
del _op, _call, _ret, _tvm_ret
|
|
|
|
|
|
# =============================================================================
|
|
# RMA get/put (thread/warp/block).
|
|
# =============================================================================
|
|
_RMA_SIG = "(void *dest, const void *source, size_t nelems, int pe)"
|
|
for _op, _backend_call in [
|
|
("nvshmem_getmem_nbi", "nvshmem_getmem_nbi"),
|
|
("nvshmem_putmem_nbi", "nvshmem_putmem_nbi"),
|
|
("nvshmem_getmem_nbi_warp", "nvshmemx_getmem_nbi_warp"),
|
|
("nvshmem_putmem_nbi_warp", "nvshmemx_putmem_nbi_warp"),
|
|
("nvshmem_getmem_nbi_block", "nvshmemx_getmem_nbi_block"),
|
|
("nvshmem_putmem_nbi_block", "nvshmemx_putmem_nbi_block"),
|
|
]:
|
|
device_intrinsic(
|
|
_op,
|
|
c_signature=_RMA_SIG,
|
|
body=f" {_backend_call}(dest, source, nelems, pe);",
|
|
extra_deps=_NVSHMEM,
|
|
)
|
|
del _op, _backend_call
|
|
|
|
|
|
# =============================================================================
|
|
# Signal / wait_until — each backend call is one device_intrinsic. String
|
|
# attrs (sig_op, cmp) are mapped to NVSHMEM integer constants in the
|
|
# user-facing dispatcher below.
|
|
# =============================================================================
|
|
|
|
_SIG_OP_VAL = {"set": 0, "add": 1}
|
|
_CMP_VAL = {"eq": 0, "ne": 1, "gt": 2, "ge": 3, "lt": 4, "le": 5}
|
|
|
|
|
|
def _resolve_attr(value, table, label):
|
|
s = value if isinstance(value, str) else value.value
|
|
if s not in table:
|
|
raise ValueError(f"Unsupported {label}: {s}")
|
|
return table[s]
|
|
|
|
|
|
device_intrinsic(
|
|
"_nvshmem_signal_op_impl",
|
|
helper_name="tvm_builtin_nvshmem_signal_op",
|
|
c_signature="(uint64_t* sig_addr, uint64_t signal, int sig_op, int pe)",
|
|
body=" nvshmemx_signal_op(sig_addr, signal, sig_op, pe);",
|
|
extra_deps=_NVSHMEM,
|
|
)
|
|
|
|
|
|
@register_codegen("nvshmem_signal_op")
|
|
def codegen_nvshmem_signal_op(sig_addr, signal, sig_op, pe):
|
|
"""Map ``sig_op`` (string) to its NVSHMEM int constant, then forward."""
|
|
sig_op_int = _resolve_attr(sig_op, _SIG_OP_VAL, "signal op")
|
|
result = CODEGEN_REGISTRY["tirx._nvshmem_signal_op_impl"]([sig_addr, signal, sig_op_int, pe])
|
|
return result
|
|
|
|
|
|
# nvshmem_<type>_wait_until — one device_intrinsic per supported type.
|
|
_WAIT_UNTIL_TYPES = {"uint64_t": "uint64", "uint64": "uint64"}
|
|
|
|
for _c_type, _suffix in [("uint64_t", "uint64")]:
|
|
device_intrinsic(
|
|
f"_nvshmem_{_suffix}_wait_until_impl",
|
|
helper_name=f"tvm_builtin_nvshmem_{_suffix}_wait_until",
|
|
c_signature=f"({_c_type}* ivar, int cmp, {_c_type} cmp_value)",
|
|
body=f" nvshmem_{_suffix}_wait_until(ivar, cmp, cmp_value);",
|
|
extra_deps=_NVSHMEM,
|
|
)
|
|
del _c_type, _suffix
|
|
|
|
|
|
@register_codegen("nvshmem_wait_until")
|
|
def codegen_nvshmem_wait_until(ivar, cmp, cmp_value, type):
|
|
"""Dispatch to the type-specific wait_until helper after mapping ``cmp``
|
|
(string) to its NVSHMEM int constant."""
|
|
type_str = type if isinstance(type, str) else type.value
|
|
if type_str not in _WAIT_UNTIL_TYPES:
|
|
raise ValueError(f"Unsupported type for nvshmem_wait_until: {type_str}")
|
|
suffix = _WAIT_UNTIL_TYPES[type_str]
|
|
cmp_int = _resolve_attr(cmp, _CMP_VAL, "cmp operation")
|
|
result = CODEGEN_REGISTRY[f"tirx._nvshmem_{suffix}_wait_until_impl"]([ivar, cmp_int, cmp_value])
|
|
return result
|
|
|
|
|
|
# putmem_signal_nbi (thread / warp / block) — three scope-specific helpers.
|
|
_PUTMEM_SIG_SIG = (
|
|
"(void* dest, const void* source, size_t nelems, "
|
|
"uint64_t* sig_addr, uint64_t signal, int sig_op, int pe)"
|
|
)
|
|
for _scope_suffix, _backend_call in [
|
|
("", "nvshmem_putmem_signal_nbi"),
|
|
("_warp", "nvshmemx_putmem_signal_nbi_warp"),
|
|
("_block", "nvshmemx_putmem_signal_nbi_block"),
|
|
]:
|
|
device_intrinsic(
|
|
f"_nvshmem_putmem_signal_nbi{_scope_suffix}_impl",
|
|
helper_name=f"tvm_builtin_nvshmem_putmem_signal_nbi{_scope_suffix}",
|
|
c_signature=_PUTMEM_SIG_SIG,
|
|
body=f" {_backend_call}(dest, source, nelems, sig_addr, signal, sig_op, pe);",
|
|
extra_deps=_NVSHMEM,
|
|
)
|
|
del _scope_suffix, _backend_call
|
|
|
|
|
|
def _make_putmem_signal_dispatcher(scope_suffix):
|
|
@register_codegen(f"nvshmem_putmem_signal_nbi{scope_suffix}")
|
|
def _codegen(dest, source, nelems, sig_addr, signal, sig_op, pe):
|
|
sig_op_int = _resolve_attr(sig_op, _SIG_OP_VAL, "signal op")
|
|
result = CODEGEN_REGISTRY[f"tirx._nvshmem_putmem_signal_nbi{scope_suffix}_impl"](
|
|
[dest, source, nelems, sig_addr, signal, sig_op_int, pe]
|
|
)
|
|
return result
|
|
|
|
return _codegen
|
|
|
|
|
|
for _suffix in ("", "_warp", "_block"):
|
|
_make_putmem_signal_dispatcher(_suffix)
|
|
del _suffix
|