chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,49 @@
|
||||
# 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=unused-import
|
||||
"""CUDA HW intrinsic codegens, grouped by feature domain.
|
||||
|
||||
- ``mma`` / ``wgmma`` / ``tcgen05`` — matrix-multiply hardware (Volta+/Hopper/Blackwell).
|
||||
- ``cp_async`` — cp.async + cp.async.bulk + cp.async.bulk.tensor (TMA), incl. TMA address helpers.
|
||||
- ``sync`` — barriers, fences, mbarrier, cluster.barrier, warp vote, elect, sync helpers.
|
||||
- ``math`` — packed-f32x2 arithmetic, exp2/rcp/reduce3, warp/CTA reductions.
|
||||
- ``memory`` — typed copies, ldg, ld.global.acquire, atomics, type conversions, address casts.
|
||||
- ``nvshmem`` — NVSHMEM RMA / signal / collective.
|
||||
- ``misc`` — register-allocation control, profiler timer, debug helpers (printf / trap).
|
||||
|
||||
Plus the support modules:
|
||||
|
||||
- ``header`` — CUDA header generator and helper-tag table.
|
||||
- ``registry`` — codegen registry.
|
||||
- ``types`` — PTX dtype enum.
|
||||
- ``utils`` — small parsing / validation helpers.
|
||||
"""
|
||||
|
||||
# Import op modules to register their codegen functions.
|
||||
from . import cp_async, math, memory, misc, mma, nvshmem, sync, tcgen05, wgmma
|
||||
from .header import TAGS, header_generator
|
||||
from .registry import CODEGEN_REGISTRY, get_codegen, register_codegen
|
||||
from .types import PTXDataType
|
||||
|
||||
__all__ = [
|
||||
"CODEGEN_REGISTRY",
|
||||
"TAGS",
|
||||
"PTXDataType",
|
||||
"get_codegen",
|
||||
"header_generator",
|
||||
"register_codegen",
|
||||
]
|
||||
@@ -0,0 +1,144 @@
|
||||
# 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
|
||||
"""Thin device-helper registration for TIRx intrinsic codegens.
|
||||
|
||||
Exposes one entry point: :func:`device_intrinsic`. Given an op name plus
|
||||
``(helper_name, c_signature, body)`` (each a string or a ``(*args) -> str``
|
||||
callable), it:
|
||||
|
||||
* wraps the body in
|
||||
``__forceinline__ __device__ <ret> <name><sig> { <body> }``,
|
||||
* registers a codegen function under the op name so
|
||||
``call_intrin("", "tirx.<op_name>", *args)`` resolves to a call to that
|
||||
helper. TVM Op registration is static C++ only.
|
||||
|
||||
Args passed to the codegen are split into ``(forward_args, attr_args)``:
|
||||
the trailing ``n_attrs`` are attrs (consumed by the ``helper_name`` /
|
||||
``c_signature`` / ``body`` callables but **not** forwarded to the helper),
|
||||
and the rest are operand args (forwarded). The default ``n_attrs=0`` means
|
||||
every arg is forwarded — appropriate for fixed-arity ops with literal
|
||||
``c_signature`` and ``helper_name``.
|
||||
|
||||
Coerce / validate attrs explicitly inside the callables — there is no
|
||||
``Choice`` / ``Bool`` / ``IntAttr`` machinery; just call ``parse_str`` /
|
||||
``int`` / ``bool`` on the raw arg as needed.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
|
||||
from tvm.backend.cuda.op import cuda_func_call
|
||||
from tvm.backend.cuda.operator.intrinsics.registry import register_codegen
|
||||
|
||||
# C primitive type → TVM dtype string. Used when the caller specifies a
|
||||
# non-void ``return_type`` but no explicit ``tvm_return_type`` — the helper
|
||||
# knows the TVM-side dtype from the C return type.
|
||||
_C_TO_TVM_DTYPE = {
|
||||
"float": "float32",
|
||||
"double": "float64",
|
||||
"uint32_t": "uint32",
|
||||
"int32_t": "int32",
|
||||
"uint64_t": "uint64",
|
||||
"int64_t": "int64",
|
||||
"uint16_t": "uint16",
|
||||
"int16_t": "int16",
|
||||
"unsigned long long": "uint64",
|
||||
"long long": "int64",
|
||||
"unsigned short": "uint16",
|
||||
"bool": "bool",
|
||||
"unsigned int": "uint32",
|
||||
"int": "int32",
|
||||
}
|
||||
|
||||
|
||||
def device_intrinsic(
|
||||
op_name: str,
|
||||
*,
|
||||
helper_name: str | Callable | None = None,
|
||||
c_signature: str | Callable = "()",
|
||||
body: str | Callable,
|
||||
n_attrs: int = 0,
|
||||
return_type: str | Callable = "void",
|
||||
tvm_return_type: str | Callable | None = None,
|
||||
templated: bool = False,
|
||||
extra_deps: tuple = (),
|
||||
) -> None:
|
||||
"""Register a CUDA device-helper intrinsic.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
op_name :
|
||||
Registry key — ``call_intrin("", "tirx.<op_name>", ...)`` resolves
|
||||
here. Also used as the default helper name (``tvm_builtin_<op_name>``)
|
||||
when ``helper_name`` is not provided.
|
||||
helper_name :
|
||||
Literal C function name, OR ``(*args) -> str`` to compute it from
|
||||
attr values. Defaults to ``f"tvm_builtin_{op_name}"``.
|
||||
c_signature :
|
||||
Literal C parameter list including outer parens (``"(int x, int y)"``),
|
||||
OR ``(*args) -> str`` to compute it from attr values. Defaults to
|
||||
``"()"``.
|
||||
body :
|
||||
Literal C body string (already indented), OR ``(*args) -> str``.
|
||||
n_attrs :
|
||||
Number of trailing args that are attrs (consumed by ``helper_name``
|
||||
/ ``c_signature`` / ``body`` callables, NOT forwarded to the helper
|
||||
as call arguments). The first ``len(args) - n_attrs`` args are the
|
||||
operand args forwarded to the helper.
|
||||
return_type :
|
||||
C return type. Default ``"void"``. Either a literal string or
|
||||
``(*args) -> str`` when the helper return type depends on attrs.
|
||||
tvm_return_type :
|
||||
TVM dtype for the call result, when the helper has a non-void
|
||||
return. Either a literal string (``"int32"``) or ``(*args) -> str``.
|
||||
If omitted and ``return_type`` is non-void, it is auto-derived from
|
||||
the ``_C_TO_TVM_DTYPE`` table.
|
||||
templated :
|
||||
Prefix the helper with ``template <typename T>``.
|
||||
extra_deps :
|
||||
Helper-tag list (e.g. ``("get_tmem_addr",)``) forwarded as the second
|
||||
element of the codegen result so the header generator emits the
|
||||
prerequisite snippets.
|
||||
"""
|
||||
if helper_name is None:
|
||||
helper_name = f"tvm_builtin_{op_name}"
|
||||
extra_deps = tuple(extra_deps)
|
||||
|
||||
def codegen(*args):
|
||||
forward = args if n_attrs == 0 else args[:-n_attrs]
|
||||
name = helper_name(*args) if callable(helper_name) else helper_name
|
||||
sig = c_signature(*args) if callable(c_signature) else c_signature
|
||||
body_str = body(*args) if callable(body) else body
|
||||
ret_type = return_type(*args) if callable(return_type) else return_type
|
||||
prefix = "template <typename T>\n" if templated else ""
|
||||
source_code = (
|
||||
f"\n{prefix}__forceinline__ __device__ {ret_type} {name}{sig} {{\n{body_str}\n}}\n"
|
||||
)
|
||||
kwargs = {"source_code": source_code}
|
||||
if tvm_return_type is not None:
|
||||
kwargs["return_type"] = (
|
||||
tvm_return_type(*args) if callable(tvm_return_type) else tvm_return_type
|
||||
)
|
||||
elif ret_type != "void":
|
||||
kwargs["return_type"] = _C_TO_TVM_DTYPE.get(ret_type, ret_type)
|
||||
result = cuda_func_call(name, *forward, **kwargs)
|
||||
return (result, list(extra_deps)) if extra_deps else result
|
||||
|
||||
codegen.__name__ = f"codegen_{op_name}"
|
||||
register_codegen(op_name)(codegen)
|
||||
@@ -0,0 +1,924 @@
|
||||
# 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, too-many-arguments, too-many-locals, too-many-positional-arguments
|
||||
"""PTX cp.async / cp.async.bulk / cp.async.bulk.tensor intrinsics.
|
||||
|
||||
Each PTX form table entry is registered as one ``device_intrinsic``.
|
||||
User-facing wrappers in ``tvm.tirx.op`` keep their v1 signatures;
|
||||
``register_codegen`` dispatchers below decode the (cp_size, fill_mode,
|
||||
predicate) / (dim, cta_mask, tile_mode) arguments to pick the right form.
|
||||
Bodies are hand-written ``asm volatile(...)`` strings. The file is grouped
|
||||
as cp.async, cp.async.bulk.tensor, cp.async.bulk non-TMA, and CUDA
|
||||
compatibility helpers.
|
||||
"""
|
||||
|
||||
import tvm
|
||||
from tvm.backend.cuda.op import cuda_func_call
|
||||
|
||||
from ._schema import device_intrinsic
|
||||
from .registry import CODEGEN_REGISTRY, register_codegen
|
||||
from .utils import parse_str
|
||||
|
||||
_PREFETCH_CHOICES = ("", "64", "128", "256")
|
||||
_DIM_CHOICES = (1, 2, 3, 4, 5)
|
||||
_TILE_MODE_CHOICES = ("tile", "tile_gather4")
|
||||
|
||||
|
||||
def _safe(s):
|
||||
return s.replace("::", "_").replace(".", "_")
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# cp.async forms from the PTX Syntax block.
|
||||
#
|
||||
# Includes commit/wait plus the non-bulk shared/global copy forms.
|
||||
# =============================================================================
|
||||
device_intrinsic(
|
||||
"ptx_cp_async_commit_group",
|
||||
helper_name="tvm_builtin_ptx_cp_async_commit_group",
|
||||
body=' asm volatile("cp.async.commit_group;");',
|
||||
)
|
||||
device_intrinsic(
|
||||
"ptx_cp_async_wait_group",
|
||||
n_attrs=1,
|
||||
helper_name=lambda n: f"tvm_builtin_ptx_cp_async_wait_group_{int(n)}",
|
||||
body=lambda n: f' asm volatile("cp.async.wait_group {int(n)};");',
|
||||
)
|
||||
|
||||
|
||||
# cp.async non-bulk copy forms:
|
||||
# Form 1: cp.async.ca.shared.global ... [dst], [src], cp-size{, src-size}{, cache-policy}
|
||||
# Form 2: cp.async.cg.shared.global ... [dst], [src], 16{, src-size}{, cache-policy}
|
||||
# Form 3: cp.async.ca.shared.global ... [dst], [src], cp-size{, ignore-src}{, cache-policy}
|
||||
# Form 4: cp.async.cg.shared.global ... [dst], [src], 16{, ignore-src}{, cache-policy}
|
||||
|
||||
|
||||
def _cp_async_modifier_str(has_cache_hint, prefetch_size):
|
||||
s = ""
|
||||
if has_cache_hint:
|
||||
s += ".L2::cache_hint"
|
||||
if prefetch_size:
|
||||
s += f".L2::{prefetch_size}B"
|
||||
return s
|
||||
|
||||
|
||||
def _make_form_parts(ca_or_cg, fixed_cp_size, extra):
|
||||
"""Build a parts callable for one of the cp.async PTX forms.
|
||||
|
||||
Args layout: (dst, src [, extra_int], cache_policy, has_cache, prefetch_size [, cp_size_attr])
|
||||
Forwarded operands: dst, src [, extra_int], cache_policy.
|
||||
Trailing attrs: has_cache, prefetch_size [, cp_size if .ca].
|
||||
"""
|
||||
n_op = 3 if extra is not None else 2
|
||||
n_attrs = 2 if fixed_cp_size is not None else 3
|
||||
extra_in_name = f"_with_{extra}" if extra is not None else ""
|
||||
|
||||
def _parts(*args):
|
||||
# Operand args (forwarded) come first, then attr args.
|
||||
attr_args = args[-n_attrs:]
|
||||
has_cache = _bool_attr(attr_args[0])
|
||||
prefetch_size = parse_str(attr_args[1])
|
||||
cp_size = fixed_cp_size if fixed_cp_size is not None else int(attr_args[2])
|
||||
modifier = _cp_async_modifier_str(has_cache, prefetch_size)
|
||||
cache_operand = ', "l"(cache_policy)' if has_cache else ""
|
||||
# name parts
|
||||
name_cache = "_cache_hint" if has_cache else ""
|
||||
name_prefetch = f"_prefetch_{prefetch_size}" if prefetch_size else ""
|
||||
name = (
|
||||
f"tvm_builtin_ptx_cp_async_{ca_or_cg}_{cp_size}"
|
||||
f"{name_cache}{name_prefetch}{extra_in_name}"
|
||||
)
|
||||
sig = (
|
||||
"(void* dst, void* src"
|
||||
+ (f", int {extra}" if extra else "")
|
||||
+ ", unsigned long long cache_policy)"
|
||||
)
|
||||
instr_base = f"cp.async.{ca_or_cg}.shared.global{modifier}"
|
||||
if extra is None:
|
||||
cache_arg = ", %2" if has_cache else ""
|
||||
body = (
|
||||
" unsigned int dst_addr = __cvta_generic_to_shared(dst);\n"
|
||||
f' asm volatile("{instr_base} [%0], [%1], {cp_size}{cache_arg};\\n"\n'
|
||||
f' :: "r"(dst_addr), "l"(src){cache_operand} : "memory");'
|
||||
)
|
||||
else:
|
||||
cache_arg = ", %3" if has_cache else ""
|
||||
body = (
|
||||
" unsigned int dst_addr = __cvta_generic_to_shared(dst);\n"
|
||||
f' asm volatile("{instr_base} [%0], [%1], {cp_size}, %2{cache_arg};\\n"\n'
|
||||
f' :: "r"(dst_addr), "l"(src), "r"({extra})'
|
||||
f'{cache_operand} : "memory");'
|
||||
)
|
||||
return name, sig, body
|
||||
|
||||
return _parts, n_op + n_attrs - n_op # n_attrs
|
||||
|
||||
|
||||
def _register_nb_form(op_name, ca_or_cg, fixed_cp_size, extra):
|
||||
parts_fn, n_attrs = _make_form_parts(ca_or_cg, fixed_cp_size, extra)
|
||||
n_op = 3 if extra is not None else 2
|
||||
sig_static = (
|
||||
"(void* dst, void* src"
|
||||
+ (f", int {extra}" if extra else "")
|
||||
+ ", unsigned long long cache_policy)"
|
||||
)
|
||||
device_intrinsic(
|
||||
f"ptx_cp_async_{op_name}",
|
||||
n_attrs=n_attrs,
|
||||
c_signature=sig_static, # static — depends on `extra` not on attrs
|
||||
helper_name=lambda *a, fn=parts_fn: fn(*a)[0],
|
||||
body=lambda *a, fn=parts_fn: fn(*a)[2],
|
||||
)
|
||||
return n_op
|
||||
|
||||
|
||||
# Form 1: .ca + src-size (cp-size ∈ {4, 8}). src-size is required when present.
|
||||
_register_nb_form("ca_src_size", "ca", fixed_cp_size=None, extra="src_size")
|
||||
# Form 2: .cg + src-size (cp-size = 16).
|
||||
_register_nb_form("cg_src_size", "cg", fixed_cp_size=16, extra="src_size")
|
||||
# Form 3: .ca + ignore-src.
|
||||
_register_nb_form("ca_ignore_src", "ca", fixed_cp_size=None, extra="ignore_src")
|
||||
# Form 4: .cg + ignore-src.
|
||||
_register_nb_form("cg_ignore_src", "cg", fixed_cp_size=16, extra="ignore_src")
|
||||
# Plain degenerate of forms 1+2 with optional src-size omitted.
|
||||
_register_nb_form("ca", "ca", fixed_cp_size=None, extra=None)
|
||||
_register_nb_form("cg", "cg", fixed_cp_size=16, extra=None)
|
||||
|
||||
|
||||
def _make_setp_at_p_helper(ca_or_cg, cp_size, has_cache, prefetch):
|
||||
"""Wrapper convenience: ``setp+@p`` around a form 1/2 cp.async (predicate-
|
||||
gated skip with dst untouched on false). Not a PTX form — emitted directly
|
||||
here as a one-off helper rather than a separate device_intrinsic."""
|
||||
modifier = _cp_async_modifier_str(has_cache, prefetch)
|
||||
cache_arg = ", %4" if has_cache else ""
|
||||
cache_operand = ', "l"(cache_policy)' if has_cache else ""
|
||||
func_name = (
|
||||
f"tvm_builtin_ptx_cp_async_{cp_size}"
|
||||
+ ("_cache_hint" if has_cache else "")
|
||||
+ (f"_prefetch_{prefetch}" if prefetch else "")
|
||||
+ "_predicate"
|
||||
)
|
||||
body = (
|
||||
" unsigned int dst_addr = __cvta_generic_to_shared(dst);\n"
|
||||
" __asm__ __volatile__(\n"
|
||||
' "{\\n"\n'
|
||||
' " .reg .pred p;\\n"\n'
|
||||
' " setp.eq.u32 p, %3, 1;\\n"\n'
|
||||
f' " @p cp.async.{ca_or_cg}.shared.global{modifier}'
|
||||
f' [%0], [%1], %2{cache_arg};\\n"\n'
|
||||
' "}\\n"\n'
|
||||
f' :: "r"(dst_addr), "l"(src), "n"({cp_size}), "r"(predicate){cache_operand}\n'
|
||||
" );"
|
||||
)
|
||||
source_code = (
|
||||
f"\n__forceinline__ __device__ void {func_name}"
|
||||
"(void* dst, void* src, int predicate, unsigned long long cache_policy) {\n"
|
||||
f"{body}\n"
|
||||
"}\n"
|
||||
)
|
||||
return func_name, source_code
|
||||
|
||||
|
||||
@register_codegen("ptx_cp_async")
|
||||
def codegen_ptx_cp_async(*args):
|
||||
"""Map the wrapper API to the 4 PTX form table entries.
|
||||
|
||||
Accepts three call shapes (sorted by arity):
|
||||
|
||||
* 5 args ``(dst_ptr, dst_offset, src_ptr, src_offset, cp_size)`` —
|
||||
the legacy form emitted by
|
||||
``tvm.backend.cuda.transform.InjectPTXAsyncCopy``.
|
||||
Offsets are folded into the pointers via ``tvm_access_ptr`` (in
|
||||
bytes; offsets are pre-scaled by the pass) and the call is
|
||||
forwarded with default cache / predicate / fill_mode.
|
||||
* 6 args ``(dst_ptr, dst_offset, src_ptr, src_offset, cp_size,
|
||||
predicate)`` — same as 5-arg form with an explicit predicate,
|
||||
zero-filling the destination when the predicate is false.
|
||||
* 8 args ``(dst_ptr, src_ptr, cp_size, cache_policy, has_cache_hint,
|
||||
prefetch_size, predicate, fill_mode)`` — the fork-native wrapper
|
||||
API.
|
||||
|
||||
The three resulting form_kinds:
|
||||
|
||||
* ``fill_mode == "zero"`` -> form 1/2 (src-size = predicate ? cp_size : 0)
|
||||
* ``predicate != -1`` and no fill_mode -> form 1/2 wrapped in setp+@p
|
||||
(wrapper convenience; not a PTX form)
|
||||
* else -> form 1/2 with src-size omitted (the "plain" degenerate)
|
||||
"""
|
||||
from tvm.tirx.op import if_then_else
|
||||
|
||||
if len(args) in (5, 6):
|
||||
# Legacy InjectPTXAsyncCopy emission: (dst_ptr, dst_off, src_ptr,
|
||||
# src_off, cp_size [, predicate]). Offsets are element indices into
|
||||
# the typed buffers (the pass uses index_factor=1 except for the
|
||||
# shared.dyn-merged byte-buffer path). Emit a C helper that scales
|
||||
# the offset by the buffer element size, then runs cp.async.
|
||||
#
|
||||
# PTX plain form for both .ca and .cg is just
|
||||
# ``cp.async.<v>.shared.global [dst], [src], cp_size;`` — three
|
||||
# operands, no trailing src-size / cache-policy.
|
||||
from tvm import DataType
|
||||
|
||||
dst_ptr_in, dst_offset, src_ptr_in, src_offset, cp_size = args[:5]
|
||||
predicate = args[5] if len(args) == 6 else -1
|
||||
cp_size_v = int(cp_size)
|
||||
ca_or_cg = "cg" if cp_size_v == 16 else "ca"
|
||||
|
||||
# Recover the per-side element dtype from each pointer's type
|
||||
# type (Var has ty = PointerType(PrimType(dtype))).
|
||||
# InjectPTXAsyncCopy emits offsets in element-units of each side's
|
||||
# buffer dtype (dst gets dst_offset * src_elem_size only when dst is a
|
||||
# merged shared.dyn byte buffer, in which case dst_elem_dtype is uint8
|
||||
# and the resulting scale-by-1 is a no-op).
|
||||
def _elem_bytes(ptr):
|
||||
ta = getattr(ptr, "ty", None)
|
||||
if ta is None or getattr(ta, "element_type", None) is None:
|
||||
return 1
|
||||
et = ta.element_type
|
||||
if not hasattr(et, "dtype"):
|
||||
return 1
|
||||
bits = DataType(str(et.dtype)).bits
|
||||
assert bits % 8 == 0, f"non-byte element dtype: {et.dtype}"
|
||||
return bits // 8
|
||||
|
||||
dst_elem_bytes = _elem_bytes(dst_ptr_in)
|
||||
src_elem_bytes = _elem_bytes(src_ptr_in)
|
||||
has_predicate = not (
|
||||
(isinstance(predicate, int) and predicate == -1)
|
||||
or (hasattr(predicate, "value") and int(predicate.value) == -1)
|
||||
)
|
||||
|
||||
def _scale(n):
|
||||
return "" if n == 1 else f" * {n}"
|
||||
|
||||
dst_scale = _scale(dst_elem_bytes)
|
||||
src_scale = _scale(src_elem_bytes)
|
||||
if has_predicate:
|
||||
func_name = (
|
||||
f"ptx_cp_async_legacy_pred_{ca_or_cg}_{cp_size_v}_{dst_elem_bytes}_{src_elem_bytes}"
|
||||
)
|
||||
if cp_size_v == 4:
|
||||
zero_fill = ' " @!p st.shared.u32 [%0], {%4};\\n"\n'
|
||||
elif cp_size_v == 8:
|
||||
zero_fill = ' " @!p st.shared.v2.u32 [%0], {%4, %4};\\n"\n'
|
||||
elif cp_size_v == 16:
|
||||
zero_fill = ' " @!p st.shared.v4.u32 [%0], {%4, %4, %4, %4};\\n"\n'
|
||||
else:
|
||||
raise ValueError(f"unsupported legacy predicated cp.async size: {cp_size_v}")
|
||||
body = (
|
||||
f" uint8_t* dst_p = (uint8_t*)dst + dst_off{dst_scale};\n"
|
||||
f" uint8_t* src_p = (uint8_t*)src + src_off{src_scale};\n"
|
||||
" unsigned int dst_addr = __cvta_generic_to_shared(dst_p);\n"
|
||||
" __asm__ __volatile__(\n"
|
||||
' "{\\n"\n'
|
||||
' " .reg .pred p;\\n"\n'
|
||||
' " setp.eq.u32 p, %3, 1;\\n"\n'
|
||||
f' " @p cp.async.{ca_or_cg}.shared.global'
|
||||
' [%0], [%1], %2;\\n"\n'
|
||||
f"{zero_fill}"
|
||||
' "}\\n"\n'
|
||||
f' :: "r"(dst_addr), "l"(src_p), "n"({cp_size_v}), "r"(predicate), "r"(0)\n'
|
||||
" );"
|
||||
)
|
||||
source_code = (
|
||||
f"\n__forceinline__ __device__ void {func_name}"
|
||||
"(void* dst, int dst_off, void* src, int src_off, int predicate) {\n"
|
||||
f"{body}\n"
|
||||
"}\n"
|
||||
)
|
||||
return cuda_func_call(
|
||||
func_name,
|
||||
dst_ptr_in,
|
||||
dst_offset,
|
||||
src_ptr_in,
|
||||
src_offset,
|
||||
predicate,
|
||||
source_code=source_code,
|
||||
)
|
||||
# No predicate — plain cp.async.
|
||||
func_name = f"ptx_cp_async_legacy_{ca_or_cg}_{cp_size_v}_{dst_elem_bytes}_{src_elem_bytes}"
|
||||
body = (
|
||||
f" uint8_t* dst_p = (uint8_t*)dst + dst_off{dst_scale};\n"
|
||||
f" uint8_t* src_p = (uint8_t*)src + src_off{src_scale};\n"
|
||||
" unsigned int dst_addr = __cvta_generic_to_shared(dst_p);\n"
|
||||
f' asm volatile("cp.async.{ca_or_cg}.shared.global'
|
||||
' [%0], [%1], %2;"\n'
|
||||
f' :: "r"(dst_addr), "l"(src_p), "n"({cp_size_v}));'
|
||||
)
|
||||
source_code = (
|
||||
f"\n__forceinline__ __device__ void {func_name}"
|
||||
"(void* dst, int dst_off, void* src, int src_off) {\n"
|
||||
f"{body}\n"
|
||||
"}\n"
|
||||
)
|
||||
return cuda_func_call(
|
||||
func_name,
|
||||
dst_ptr_in,
|
||||
dst_offset,
|
||||
src_ptr_in,
|
||||
src_offset,
|
||||
source_code=source_code,
|
||||
)
|
||||
elif len(args) == 8:
|
||||
(
|
||||
dst_ptr,
|
||||
src_ptr,
|
||||
cp_size,
|
||||
cache_policy,
|
||||
has_cache_hint,
|
||||
prefetch_size,
|
||||
predicate,
|
||||
fill_mode,
|
||||
) = args
|
||||
else:
|
||||
raise ValueError(f"ptx_cp_async codegen expects 5/6/8 args, got {len(args)}")
|
||||
|
||||
cp_size_v = int(cp_size)
|
||||
ca_or_cg = "cg" if cp_size_v == 16 else "ca"
|
||||
pref = "" if int(prefetch_size) == -1 else str(int(prefetch_size))
|
||||
fill = parse_str(fill_mode)
|
||||
has_cache = _bool_attr(has_cache_hint)
|
||||
has_predicate = not (
|
||||
(isinstance(predicate, int) and predicate == -1)
|
||||
or (hasattr(predicate, "value") and int(predicate.value) == -1)
|
||||
)
|
||||
|
||||
if fill == "zero":
|
||||
src_size = if_then_else(predicate != 0, cp_size_v, 0)
|
||||
op = f"tirx.ptx_cp_async_{ca_or_cg}_src_size"
|
||||
if cp_size_v == 16:
|
||||
args = [dst_ptr, src_ptr, src_size, cache_policy, has_cache, pref]
|
||||
else:
|
||||
args = [dst_ptr, src_ptr, src_size, cache_policy, has_cache, pref, cp_size_v]
|
||||
result = CODEGEN_REGISTRY[op](args)
|
||||
return result[0] if isinstance(result, tuple) else result
|
||||
|
||||
if has_predicate:
|
||||
func_name, source_code = _make_setp_at_p_helper(ca_or_cg, cp_size_v, has_cache, pref)
|
||||
return cuda_func_call(
|
||||
func_name, dst_ptr, src_ptr, predicate, cache_policy, source_code=source_code
|
||||
)
|
||||
|
||||
# Plain — form 1/2 with src-size omitted.
|
||||
op = f"tirx.ptx_cp_async_{ca_or_cg}"
|
||||
if cp_size_v == 16:
|
||||
args = [dst_ptr, src_ptr, cache_policy, has_cache, pref]
|
||||
else:
|
||||
args = [dst_ptr, src_ptr, cache_policy, has_cache, pref, cp_size_v]
|
||||
result = CODEGEN_REGISTRY[op](args)
|
||||
return result[0] if isinstance(result, tuple) else result
|
||||
|
||||
|
||||
CODEGEN_REGISTRY["tirx.ptx.cp_async_raw"] = CODEGEN_REGISTRY["tirx.ptx.cp_async"]
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# cp.async.bulk.tensor (TMA) — one device_intrinsic per arity variant of each
|
||||
# PTX form. Per-dim coord operands materialise via the ``c_signature`` callable.
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def _is_sm100_or_higher():
|
||||
target = tvm.target.Target.current()
|
||||
if target is None:
|
||||
return False
|
||||
arch = target.arch[3:]
|
||||
if not arch[-1].isdigit():
|
||||
arch = arch[:-1]
|
||||
return int(arch) >= 100
|
||||
|
||||
|
||||
def _resolve_cta_group_str(cta_group):
|
||||
if cta_group == 2 or (cta_group != -1 and _is_sm100_or_higher()):
|
||||
return f".cta_group::{cta_group}"
|
||||
return ""
|
||||
|
||||
|
||||
def _coord_template(coord_count, start_slot):
|
||||
inner = ", ".join(f"%{start_slot + i}" for i in range(coord_count))
|
||||
return f"{{{inner}}}"
|
||||
|
||||
|
||||
def _coord_constraints(coord_count):
|
||||
return ", ".join(f'"r"(coord{i})' for i in range(coord_count))
|
||||
|
||||
|
||||
def _coord_sig(n):
|
||||
return ", ".join(f"int coord{i}" for i in range(n))
|
||||
|
||||
|
||||
# PTX cp.async.bulk.tensor global -> shared::cluster form:
|
||||
# cp.async.bulk.tensor.dim.dst.src{.load_mode}.completion_mechanism
|
||||
# {.multicast}{.cta_group}{.level::cache_hint}
|
||||
# [dstMem], [tensorMap, tensorCoords], [mbar]{, im2colInfo}
|
||||
# {, ctaMask} {, cache-policy}
|
||||
# .dst = {.shared::cluster}; .src = {.global}
|
||||
# .completion_mechanism = {.mbarrier::complete_tx::bytes}
|
||||
# .multicast = {.multicast::cluster}
|
||||
# .cta_group = {.cta_group::1, .cta_group::2}
|
||||
# .load_mode = {.tile, .tile::gather4, .im2col, .im2col::w, .im2col::w::128}
|
||||
# .level::cache_hint = {.L2::cache_hint}
|
||||
# This registration supports tile/tile::gather4 modes; ctaMask is only used
|
||||
# when the optional ``.multicast::cluster`` modifier is enabled.
|
||||
def _g2cluster_parts(*args):
|
||||
attrs = args[-6:]
|
||||
dim = int(attrs[0])
|
||||
cta_group = int(attrs[1])
|
||||
has_cache = _bool_attr(attrs[2])
|
||||
tile_mode = parse_str(attrs[3])
|
||||
bar_is_addr = _bool_attr(attrs[4])
|
||||
multicast = _bool_attr(attrs[5])
|
||||
coord_count = 5 if tile_mode == "tile_gather4" else dim
|
||||
bar_type = "unsigned int bar_addr" if bar_is_addr else "void* bar"
|
||||
sig = (
|
||||
f"(void* dst, {bar_type}, unsigned long long tensormap_addr, "
|
||||
"uint16_t cta_mask, unsigned long long cache_policy"
|
||||
+ (", " + _coord_sig(coord_count) if coord_count else "")
|
||||
+ ")"
|
||||
)
|
||||
name = (
|
||||
f"ptx_cp_async_bulk_tensor_g2cluster_{tile_mode}_{dim}d"
|
||||
f"{'_multicast' if multicast else ''}"
|
||||
f"{'_cache_hint' if has_cache else ''}{'_bar_addr' if bar_is_addr else ''}"
|
||||
)
|
||||
tile_modifier = ".tile::gather4" if tile_mode == "tile_gather4" else ""
|
||||
cta_group_str = _resolve_cta_group_str(cta_group)
|
||||
multicast_inst = ".multicast::cluster" if multicast else ""
|
||||
cache_inst = ".L2::cache_hint" if has_cache else ""
|
||||
mask_arg = ',\n "h"(cta_mask)' if multicast else ""
|
||||
cache_arg = ',\n "l"(cache_policy)' if has_cache else ""
|
||||
mask_slot = ", %3" if multicast else ""
|
||||
cache_slot = ", %4" if multicast and has_cache else ", %3" if has_cache else ""
|
||||
coord_start = 5 if multicast and has_cache else 4 if multicast or has_cache else 3
|
||||
coord_tpl = _coord_template(coord_count, coord_start)
|
||||
instr = (
|
||||
f"cp.async.bulk.tensor.{dim}d.shared::cluster.global{tile_modifier}"
|
||||
f".mbarrier::complete_tx::bytes{multicast_inst}"
|
||||
f"{cta_group_str}{cache_inst}"
|
||||
)
|
||||
bar_addr_decl = (
|
||||
"" if bar_is_addr else " unsigned int bar_addr = __cvta_generic_to_shared(bar);\n"
|
||||
)
|
||||
body = (
|
||||
" unsigned int dst_addr = __cvta_generic_to_shared(dst);\n"
|
||||
f"{bar_addr_decl}"
|
||||
" asm volatile(\n"
|
||||
f' "{instr} [%0], [%1, {coord_tpl}], [%2]{mask_slot}{cache_slot};"\n'
|
||||
" :\n"
|
||||
f' : "r"(dst_addr), "l"(tensormap_addr), "r"(bar_addr){mask_arg}{cache_arg},\n'
|
||||
f" {_coord_constraints(coord_count)}\n"
|
||||
' : "memory"\n'
|
||||
" );"
|
||||
)
|
||||
return name, sig, body
|
||||
|
||||
|
||||
device_intrinsic(
|
||||
"ptx_cp_async_bulk_tensor_g2cluster",
|
||||
n_attrs=6,
|
||||
helper_name=lambda *a: _g2cluster_parts(*a)[0],
|
||||
c_signature=lambda *a: _g2cluster_parts(*a)[1],
|
||||
body=lambda *a: _g2cluster_parts(*a)[2],
|
||||
)
|
||||
|
||||
|
||||
# PTX cp.async.bulk.tensor shared::cta -> global form:
|
||||
# cp.async.bulk.tensor.dim.dst.src{.load_mode}.completion_mechanism
|
||||
# {.level::cache_hint}
|
||||
# [tensorMap, tensorCoords], [srcMem] {, cache-policy}
|
||||
# .dst = {.global}; .src = {.shared::cta}
|
||||
# .completion_mechanism = {.bulk_group}
|
||||
# .load_mode = {.tile, .tile::scatter4, .im2col_no_offs}
|
||||
# .level::cache_hint = {.L2::cache_hint}
|
||||
# This registration supports tile mode; cache-policy is a real operand.
|
||||
def _s2g_parts(*args):
|
||||
attrs = args[-2:]
|
||||
dim = int(attrs[0])
|
||||
has_cache = _bool_attr(attrs[1])
|
||||
sig = (
|
||||
"(void* src, unsigned long long tensormap_addr, unsigned long long cache_policy"
|
||||
+ (", " + _coord_sig(dim) if dim else "")
|
||||
+ ")"
|
||||
)
|
||||
name = f"ptx_cp_async_bulk_tensor_shared_to_global_{dim}d{'_cache_hint' if has_cache else ''}"
|
||||
cache_inst = ".L2::cache_hint" if has_cache else ""
|
||||
cache_arg = ', "l"(cache_policy)' if has_cache else ""
|
||||
cache_slot = ", %2" if has_cache else ""
|
||||
coord_start = 3 if has_cache else 2
|
||||
coord_tpl = _coord_template(dim, coord_start)
|
||||
instr = f"cp.async.bulk.tensor.{dim}d.global.shared::cta.tile.bulk_group{cache_inst}"
|
||||
body = (
|
||||
" unsigned int src_addr = __cvta_generic_to_shared(src);\n"
|
||||
" asm volatile(\n"
|
||||
f' "{instr} [%0, {coord_tpl}], [%1]{cache_slot};"\n'
|
||||
" :\n"
|
||||
f' : "l"(tensormap_addr), "r"(src_addr){cache_arg},\n'
|
||||
f" {_coord_constraints(dim)}\n"
|
||||
' : "memory"\n'
|
||||
" );"
|
||||
)
|
||||
return name, sig, body
|
||||
|
||||
|
||||
device_intrinsic(
|
||||
"ptx_cp_async_bulk_tensor_s2g",
|
||||
n_attrs=2,
|
||||
helper_name=lambda *a: _s2g_parts(*a)[0],
|
||||
c_signature=lambda *a: _s2g_parts(*a)[1],
|
||||
body=lambda *a: _s2g_parts(*a)[2],
|
||||
)
|
||||
|
||||
|
||||
# PTX cp.async.bulk.prefetch.tensor form:
|
||||
# cp.async.bulk.prefetch.tensor.dim.L2.src{.load_mode}{.level::cache_hint}
|
||||
# [tensorMap, tensorCoords] {, im2colInfo} {, cache-policy}
|
||||
# .src = {.global}
|
||||
# .load_mode = {.tile, .tile::gather4, .im2col, .im2col::w, .im2col::w::128}
|
||||
# .level::cache_hint = {.L2::cache_hint}
|
||||
# This registration supports tile mode; cache-policy is a real operand.
|
||||
def _prefetch_parts(*args):
|
||||
attrs = args[-2:]
|
||||
dim = int(attrs[0])
|
||||
has_cache = _bool_attr(attrs[1])
|
||||
sig = (
|
||||
"(unsigned long long tensormap_addr, unsigned long long cache_policy"
|
||||
+ (", " + _coord_sig(dim) if dim else "")
|
||||
+ ")"
|
||||
)
|
||||
name = (
|
||||
f"ptx_cp_async_bulk_tensor_global_to_cluster_prefetch_{dim}d"
|
||||
f"{'_cache_hint' if has_cache else ''}"
|
||||
)
|
||||
cache_inst = ".L2::cache_hint" if has_cache else ""
|
||||
cache_arg = ', "l"(cache_policy)' if has_cache else ""
|
||||
cache_slot = ", %1" if has_cache else ""
|
||||
coord_start = 2 if has_cache else 1
|
||||
coord_tpl = _coord_template(dim, coord_start)
|
||||
instr = f"cp.async.bulk.prefetch.tensor.{dim}d.L2.global.tile{cache_inst}"
|
||||
body = (
|
||||
" asm volatile(\n"
|
||||
f' "{instr} [%0, {coord_tpl}]{cache_slot};"\n'
|
||||
" :\n"
|
||||
f' : "l"(tensormap_addr){cache_arg},\n'
|
||||
f" {_coord_constraints(dim)}\n"
|
||||
' : "memory"\n'
|
||||
" );"
|
||||
)
|
||||
return name, sig, body
|
||||
|
||||
|
||||
device_intrinsic(
|
||||
"ptx_cp_async_bulk_tensor_prefetch",
|
||||
n_attrs=2,
|
||||
helper_name=lambda *a: _prefetch_parts(*a)[0],
|
||||
c_signature=lambda *a: _prefetch_parts(*a)[1],
|
||||
body=lambda *a: _prefetch_parts(*a)[2],
|
||||
)
|
||||
|
||||
|
||||
# PTX cp.reduce.async.bulk.tensor shared::cta -> global form:
|
||||
# cp.reduce.async.bulk.tensor.dim.dst.src.redOp{.load_mode}.completion_mechanism
|
||||
# {.level::cache_hint}
|
||||
# [tensorMap, tensorCoords], [srcMem] {, cache-policy}
|
||||
# .dst = {.global}; .src = {.shared::cta}
|
||||
# .completion_mechanism = {.bulk_group}
|
||||
# .redOp = {.add, .min, .max, .inc, .dec, .and, .or, .xor}
|
||||
# .level::cache_hint = {.L2::cache_hint}
|
||||
# This registration supports tile mode; redOp is syntax, cache-policy is an operand.
|
||||
def _reduce_parts(*args):
|
||||
attrs = args[-3:]
|
||||
dim = int(attrs[0])
|
||||
has_cache = _bool_attr(attrs[1])
|
||||
red_op = parse_str(attrs[2])
|
||||
sig = (
|
||||
"(void* src, unsigned long long tensormap_addr, unsigned long long cache_policy"
|
||||
+ (", " + _coord_sig(dim) if dim else "")
|
||||
+ ")"
|
||||
)
|
||||
name = (
|
||||
f"ptx_cp_async_bulk_tensor_shared_to_global_reduce_{dim}d"
|
||||
f"{'_cache_hint' if has_cache else ''}"
|
||||
)
|
||||
cache_inst = ".L2::cache_hint" if has_cache else ""
|
||||
cache_arg = ', "l"(cache_policy)' if has_cache else ""
|
||||
cache_slot = ", %2" if has_cache else ""
|
||||
coord_start = 3 if has_cache else 2
|
||||
coord_tpl = _coord_template(dim, coord_start)
|
||||
instr = (
|
||||
f"cp.reduce.async.bulk.tensor.{dim}d.global.shared::cta"
|
||||
f".{red_op}.tile.bulk_group{cache_inst}"
|
||||
)
|
||||
body = (
|
||||
" unsigned int src_addr = __cvta_generic_to_shared(src);\n"
|
||||
" asm volatile(\n"
|
||||
f' "{instr} [%0, {coord_tpl}], [%1]{cache_slot};"\n'
|
||||
" :\n"
|
||||
f' : "l"(tensormap_addr), "r"(src_addr){cache_arg},\n'
|
||||
f" {_coord_constraints(dim)}\n"
|
||||
' : "memory"\n'
|
||||
" );"
|
||||
)
|
||||
return name, sig, body
|
||||
|
||||
|
||||
device_intrinsic(
|
||||
"ptx_cp_async_bulk_tensor_reduce",
|
||||
n_attrs=3,
|
||||
helper_name=lambda *a: _reduce_parts(*a)[0],
|
||||
c_signature=lambda *a: _reduce_parts(*a)[1],
|
||||
body=lambda *a: _reduce_parts(*a)[2],
|
||||
)
|
||||
|
||||
|
||||
# User-facing dispatchers for tensor global -> shared::cluster. The same
|
||||
# backend root handles the optional ``.multicast::cluster`` modifier.
|
||||
|
||||
|
||||
def _g2c_dispatch(dim, dst_ptr, bar, tensormap, *args, tile_mode):
|
||||
cta_mask, cta_group, cache_policy, has_cache, *rest = args
|
||||
coord_count = 5 if tile_mode == "tile_gather4" else int(dim)
|
||||
if len(rest) == coord_count + 1:
|
||||
bar_is_addr = _bool_attr(rest[0])
|
||||
coords = rest[1:]
|
||||
else:
|
||||
bar_is_addr = False
|
||||
coords = rest
|
||||
is_unicast = isinstance(cta_mask, tvm.tirx.IntImm) and bin(int(cta_mask)).count("1") <= 1
|
||||
cg = int(cta_group)
|
||||
op = "tirx.ptx_cp_async_bulk_tensor_g2cluster"
|
||||
call_args = [
|
||||
dst_ptr,
|
||||
bar,
|
||||
tensormap,
|
||||
cta_mask,
|
||||
cache_policy,
|
||||
*coords,
|
||||
int(dim),
|
||||
cg,
|
||||
has_cache,
|
||||
tile_mode,
|
||||
bar_is_addr,
|
||||
int(not is_unicast),
|
||||
]
|
||||
result = CODEGEN_REGISTRY[op](call_args)
|
||||
return result[0] if isinstance(result, tuple) else result
|
||||
|
||||
|
||||
@register_codegen("ptx_cp_async_bulk_tensor_global_to_cluster")
|
||||
def codegen_g2c(dim, dst_ptr, bar, tensormap, *args):
|
||||
return _g2c_dispatch(dim, dst_ptr, bar, tensormap, *args, tile_mode="tile")
|
||||
|
||||
|
||||
@register_codegen("ptx_cp_async_bulk_tensor_tile_gather4_global_to_cluster")
|
||||
def codegen_g2c_gather4(dim, dst_ptr, bar, tensormap, *args):
|
||||
return _g2c_dispatch(dim, dst_ptr, bar, tensormap, *args, tile_mode="tile_gather4")
|
||||
|
||||
|
||||
@register_codegen("ptx_cp_async_bulk_tensor_shared_to_global")
|
||||
def codegen_s2g(dim, src_ptr, tensormap, *args):
|
||||
cache_policy, has_cache, *coords = args
|
||||
result = CODEGEN_REGISTRY["tirx.ptx_cp_async_bulk_tensor_s2g"](
|
||||
[src_ptr, tensormap, cache_policy, *coords, int(dim), has_cache]
|
||||
)
|
||||
return result[0] if isinstance(result, tuple) else result
|
||||
|
||||
|
||||
@register_codegen("ptx_cp_async_bulk_tensor_global_to_cluster_prefetch")
|
||||
def codegen_prefetch(dim, tensormap, *args):
|
||||
cache_policy, has_cache, *coords = args
|
||||
result = CODEGEN_REGISTRY["tirx.ptx_cp_async_bulk_tensor_prefetch"](
|
||||
[tensormap, cache_policy, *coords, int(dim), has_cache]
|
||||
)
|
||||
return result[0] if isinstance(result, tuple) else result
|
||||
|
||||
|
||||
@register_codegen("ptx_cp_async_bulk_tensor_shared_to_global_reduce")
|
||||
def codegen_reduce(dim, src_ptr, tensormap, *args):
|
||||
cache_policy, has_cache, red_op, *coords = args
|
||||
result = CODEGEN_REGISTRY["tirx.ptx_cp_async_bulk_tensor_reduce"](
|
||||
[src_ptr, tensormap, cache_policy, *coords, int(dim), has_cache, red_op]
|
||||
)
|
||||
return result[0] if isinstance(result, tuple) else result
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# cp.async.bulk non-TMA forms from the PTX Syntax block. Each form is one
|
||||
# device_intrinsic; optional PTX modifiers are attrs, not separate fixed ops.
|
||||
# =============================================================================
|
||||
device_intrinsic(
|
||||
"ptx_cp_async_bulk_commit_group",
|
||||
helper_name="ptx_cp_async_bulk_tensor_commit_group",
|
||||
body=' asm volatile("cp.async.bulk.commit_group;");',
|
||||
)
|
||||
|
||||
|
||||
def _ptx_cp_async_bulk_wait_group_parts(n, read):
|
||||
n = int(n)
|
||||
read_b = bool(int(read)) if hasattr(read, "value") else bool(read)
|
||||
return (
|
||||
f"ptx_cp_async_bulk_wait_group{'_read' if read_b else ''}_{n}",
|
||||
f' asm volatile("cp.async.bulk.wait_group{".read" if read_b else ""} {n};");',
|
||||
)
|
||||
|
||||
|
||||
device_intrinsic(
|
||||
"ptx_cp_async_bulk_wait_group",
|
||||
n_attrs=2,
|
||||
helper_name=lambda n, read: _ptx_cp_async_bulk_wait_group_parts(n, read)[0],
|
||||
body=lambda n, read: _ptx_cp_async_bulk_wait_group_parts(n, read)[1],
|
||||
)
|
||||
|
||||
|
||||
def _bool_attr(value):
|
||||
return bool(int(value)) if hasattr(value, "value") else bool(value)
|
||||
|
||||
|
||||
def _bulk_cache_operand_constraint(has_cache):
|
||||
return ', "l"(cache_policy)' if has_cache else ""
|
||||
|
||||
|
||||
def _bulk_cache_operand_suffix(has_cache):
|
||||
return ".L2::cache_hint" if has_cache else ""
|
||||
|
||||
|
||||
# PTX cp.async.bulk global -> shared::cta form:
|
||||
# cp.async.bulk.dst.src.completion_mechanism{.level::cache_hint}{.ignore_oob}
|
||||
# [dstMem], [srcMem], size{, ignoreBytesLeft, ignoreBytesRight}, [mbar] {, cache-policy}
|
||||
# .dst = {.shared::cta}; .src = {.global}
|
||||
# .completion_mechanism = {.mbarrier::complete_tx::bytes}
|
||||
# .level::cache_hint = {.L2::cache_hint}
|
||||
def _bulk_g2s_cta_parts(*args):
|
||||
has_cache = _bool_attr(args[-2])
|
||||
ignore_oob = _bool_attr(args[-1])
|
||||
instr = (
|
||||
"cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes"
|
||||
f"{_bulk_cache_operand_suffix(has_cache)}{'.ignore_oob' if ignore_oob else ''}"
|
||||
)
|
||||
if ignore_oob:
|
||||
asm_args = (
|
||||
'"r"(dst), "l"(src_ptr), "r"(num_bytes), "r"(ignore_bytes_left), '
|
||||
'"r"(ignore_bytes_right), "r"(mbarrier)'
|
||||
)
|
||||
operands = "%2, %3, %4, [%5]"
|
||||
cache_slot = ", %6" if has_cache else ""
|
||||
else:
|
||||
asm_args = '"r"(dst), "l"(src_ptr), "r"(num_bytes), "r"(mbarrier)'
|
||||
operands = "%2, [%3]"
|
||||
cache_slot = ", %4" if has_cache else ""
|
||||
body = (
|
||||
" unsigned int dst = (unsigned int)__cvta_generic_to_shared(dst_ptr);\n"
|
||||
" unsigned int mbarrier = (unsigned int)__cvta_generic_to_shared(mbarrier_ptr);\n"
|
||||
f' asm volatile("{instr} [%0], [%1], {operands}{cache_slot};"\n'
|
||||
" :\n"
|
||||
f" : {asm_args}{_bulk_cache_operand_constraint(has_cache)}\n"
|
||||
' : "memory");'
|
||||
)
|
||||
name = (
|
||||
"tvm_builtin_ptx_cp_async_bulk_g2s_cta"
|
||||
f"{'_cache_hint' if has_cache else ''}{'_ignore_oob' if ignore_oob else ''}"
|
||||
)
|
||||
return name, body
|
||||
|
||||
|
||||
device_intrinsic(
|
||||
"ptx_cp_async_bulk_g2s_cta",
|
||||
n_attrs=2,
|
||||
helper_name=lambda *a: _bulk_g2s_cta_parts(*a)[0],
|
||||
c_signature=(
|
||||
"(void* dst_ptr, void* src_ptr, unsigned int num_bytes, "
|
||||
"unsigned int ignore_bytes_left, unsigned int ignore_bytes_right, "
|
||||
"void* mbarrier_ptr, unsigned long long cache_policy)"
|
||||
),
|
||||
body=lambda *a: _bulk_g2s_cta_parts(*a)[1],
|
||||
)
|
||||
|
||||
|
||||
# PTX cp.async.bulk global -> shared::cluster form:
|
||||
# cp.async.bulk.dst.src.completion_mechanism{.multicast}{.level::cache_hint}
|
||||
# [dstMem], [srcMem], size, [mbar] {, ctaMask} {, cache-policy}
|
||||
# .dst = {.shared::cluster}; .src = {.global}
|
||||
# .completion_mechanism = {.mbarrier::complete_tx::bytes}
|
||||
# .level::cache_hint = {.L2::cache_hint}
|
||||
# .multicast = {.multicast::cluster}
|
||||
def _bulk_g2s_cluster_parts(*args):
|
||||
has_cache = _bool_attr(args[-2])
|
||||
multicast = _bool_attr(args[-1])
|
||||
instr = (
|
||||
"cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes"
|
||||
f"{'.multicast::cluster' if multicast else ''}{_bulk_cache_operand_suffix(has_cache)}"
|
||||
)
|
||||
cta_constraint = ', "h"(cta_mask)' if multicast else ""
|
||||
mask_slot = ", %4" if multicast else ""
|
||||
cache_slot = ", %5" if multicast and has_cache else ", %4" if has_cache else ""
|
||||
body = (
|
||||
" unsigned int dst = (unsigned int)__cvta_generic_to_shared(dst_ptr);\n"
|
||||
" unsigned int mbarrier = (unsigned int)__cvta_generic_to_shared(mbarrier_ptr);\n"
|
||||
f' asm volatile("{instr} [%0], [%1], %2, [%3]'
|
||||
f'{mask_slot}{cache_slot};"\n'
|
||||
" :\n"
|
||||
' : "r"(dst), "l"(src_ptr), "r"(num_bytes), "r"(mbarrier)'
|
||||
f"{cta_constraint}{_bulk_cache_operand_constraint(has_cache)}\n"
|
||||
' : "memory");'
|
||||
)
|
||||
name = (
|
||||
"tvm_builtin_ptx_cp_async_bulk_g2s_cluster"
|
||||
f"{'_multicast' if multicast else ''}{'_cache_hint' if has_cache else ''}"
|
||||
)
|
||||
return name, body
|
||||
|
||||
|
||||
device_intrinsic(
|
||||
"ptx_cp_async_bulk_g2s_cluster",
|
||||
n_attrs=2,
|
||||
helper_name=lambda *a: _bulk_g2s_cluster_parts(*a)[0],
|
||||
c_signature=(
|
||||
"(void* dst_ptr, void* src_ptr, unsigned int num_bytes, "
|
||||
"void* mbarrier_ptr, unsigned short cta_mask, unsigned long long cache_policy)"
|
||||
),
|
||||
body=lambda *a: _bulk_g2s_cluster_parts(*a)[1],
|
||||
)
|
||||
|
||||
|
||||
# PTX cp.async.bulk shared::cta -> shared::cluster form:
|
||||
# cp.async.bulk.dst.src.completion_mechanism [dstMem], [srcMem], size, [mbar]
|
||||
# .dst = {.shared::cluster}; .src = {.shared::cta}
|
||||
# .completion_mechanism = {.mbarrier::complete_tx::bytes}
|
||||
device_intrinsic(
|
||||
"ptx_cp_async_bulk_s2s_cluster",
|
||||
helper_name="tvm_builtin_ptx_cp_async_bulk_s2s_cluster",
|
||||
c_signature="(uint64_t dst, void* src, int size, uint64_t mbar)",
|
||||
body=r""" unsigned int dst_addr = static_cast<unsigned int>(dst);
|
||||
unsigned int src_addr = __cvta_generic_to_shared(src);
|
||||
unsigned int mbar_addr = static_cast<unsigned int>(mbar);
|
||||
asm volatile(
|
||||
"cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes"
|
||||
" [%0], [%1], %2, [%3];"
|
||||
:
|
||||
: "r"(dst_addr), "r"(src_addr), "r"(size), "r"(mbar_addr)
|
||||
: "memory");""",
|
||||
)
|
||||
|
||||
|
||||
@register_codegen("ptx_cp_async_bulk_shared_to_cluster")
|
||||
def codegen_ptx_cp_async_bulk_shared_to_cluster(dst_ptr, src_ptr, size, mbar):
|
||||
result = CODEGEN_REGISTRY["tirx.ptx_cp_async_bulk_s2s_cluster"]([dst_ptr, src_ptr, size, mbar])
|
||||
return result[0] if isinstance(result, tuple) else result
|
||||
|
||||
|
||||
# PTX cp.async.bulk shared::cta -> global form:
|
||||
# cp.async.bulk.dst.src.completion_mechanism{.level::cache_hint}{.cp_mask}
|
||||
# [dstMem], [srcMem], size {, cache-policy} {, byteMask}
|
||||
# .dst = {.global}; .src = {.shared::cta}
|
||||
# .completion_mechanism = {.bulk_group}
|
||||
# .level::cache_hint = {.L2::cache_hint}
|
||||
def _bulk_s2g_parts(*args):
|
||||
has_cache = _bool_attr(args[-2])
|
||||
cp_mask = _bool_attr(args[-1])
|
||||
if cp_mask and not has_cache:
|
||||
raise ValueError("cp.async.bulk shared::cta -> global .cp_mask requires .L2::cache_hint")
|
||||
instr = f"cp.async.bulk.global.shared::cta.bulk_group{_bulk_cache_operand_suffix(has_cache)}"
|
||||
if cp_mask:
|
||||
instr += ".cp_mask"
|
||||
cache_slot = ", %3" if has_cache else ""
|
||||
mask_slot = ", %4" if cp_mask else ""
|
||||
mask_constraint = ', "r"(byte_mask)' if cp_mask else ""
|
||||
body = (
|
||||
" unsigned int src = (unsigned int)__cvta_generic_to_shared(src_ptr);\n"
|
||||
f' asm volatile("{instr} [%0], [%1], %2'
|
||||
f'{cache_slot}{mask_slot};"\n'
|
||||
" :\n"
|
||||
' : "l"(dst_ptr), "r"(src), "r"(num_bytes)'
|
||||
f"{_bulk_cache_operand_constraint(has_cache)}{mask_constraint}\n"
|
||||
' : "memory");'
|
||||
)
|
||||
name = (
|
||||
"tvm_builtin_ptx_cp_async_bulk_s2g"
|
||||
f"{'_cache_hint' if has_cache else ''}{'_cp_mask' if cp_mask else ''}"
|
||||
)
|
||||
return name, body
|
||||
|
||||
|
||||
device_intrinsic(
|
||||
"ptx_cp_async_bulk_s2g",
|
||||
n_attrs=2,
|
||||
helper_name=lambda *a: _bulk_s2g_parts(*a)[0],
|
||||
c_signature=(
|
||||
"(void* dst_ptr, void* src_ptr, unsigned int num_bytes, "
|
||||
"unsigned int byte_mask, unsigned long long cache_policy)"
|
||||
),
|
||||
body=lambda *a: _bulk_s2g_parts(*a)[1],
|
||||
)
|
||||
@@ -0,0 +1,809 @@
|
||||
# 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=line-too-long
|
||||
"""CUDA header generator for codegen.
|
||||
|
||||
The header generator is used to generate the header for the CUDA code.
|
||||
It's controlled by the predefined tags.
|
||||
The tags are used to identify the utility functions/classes necessary for the codegen.
|
||||
"""
|
||||
|
||||
import tvm_ffi
|
||||
|
||||
TAGS = {
|
||||
"cuda",
|
||||
"cuda/barrier",
|
||||
"cooperative_groups",
|
||||
"fp16",
|
||||
"bf16",
|
||||
"fp8",
|
||||
"fp6",
|
||||
"fp4",
|
||||
"int8",
|
||||
"math_constants",
|
||||
"mma",
|
||||
"warp_shuffle",
|
||||
"cast_smem_ptr_to_int",
|
||||
"get_tmem_addr",
|
||||
"gmma_descriptor",
|
||||
"smem_descriptor",
|
||||
"instr_descriptor",
|
||||
"instr_descriptor_block_scaled",
|
||||
"get_time_stamp",
|
||||
"nvshmem",
|
||||
"elect_one_sync",
|
||||
}
|
||||
|
||||
|
||||
@tvm_ffi.register_global_func("tirx.intrinsics.cuda.header_generator")
|
||||
def header_generator(tags):
|
||||
"""Generate the header for the CUDA code."""
|
||||
for tag in tags:
|
||||
if tag not in TAGS:
|
||||
raise ValueError(f"Invalid tag: {tag}")
|
||||
|
||||
header = ""
|
||||
if "nvshmem" in tags:
|
||||
header += R"""
|
||||
#include <nvshmem.h>
|
||||
#include <nvshmemx.h>
|
||||
"""
|
||||
|
||||
if "cuda/barrier" in tags or "cooperative_groups" in tags:
|
||||
header += (
|
||||
R"""
|
||||
#include <cuda/barrier>
|
||||
#include <cooperative_groups.h>
|
||||
"""
|
||||
+ "\n"
|
||||
)
|
||||
|
||||
# NVRTC has no host C++ stdlib and no <cuda.h>. Branch on __CUDACC_RTC__ so
|
||||
# the same emitted source compiles under both nvcc (offline) and NVRTC
|
||||
# (runtime) without any post-processing in tvm.support.nvcc.
|
||||
header += """
|
||||
#ifdef __CUDACC_RTC__
|
||||
#include <cuda/std/cstdint>
|
||||
using cuda::std::uint8_t;
|
||||
using cuda::std::uint16_t;
|
||||
using cuda::std::uint32_t;
|
||||
using cuda::std::uint64_t;
|
||||
using cuda::std::int8_t;
|
||||
using cuda::std::int16_t;
|
||||
using cuda::std::int32_t;
|
||||
using cuda::std::int64_t;
|
||||
|
||||
#include <cuda/std/type_traits>
|
||||
namespace std {
|
||||
using cuda::std::is_same;
|
||||
using cuda::std::is_same_v;
|
||||
using cuda::std::is_integral;
|
||||
using cuda::std::is_signed;
|
||||
using cuda::std::is_unsigned;
|
||||
using cuda::std::is_floating_point;
|
||||
using cuda::std::enable_if;
|
||||
using cuda::std::conditional;
|
||||
}
|
||||
|
||||
// NVRTC uses asm/volatile instead of __asm__/__volatile__ (gcc extension).
|
||||
#ifndef __asm__
|
||||
#define __asm__ asm
|
||||
#endif
|
||||
#ifndef __volatile__
|
||||
#define __volatile__ volatile
|
||||
#endif
|
||||
#else
|
||||
#include <cstdint>
|
||||
#include <type_traits>
|
||||
#include <cuda.h>
|
||||
#endif
|
||||
"""
|
||||
|
||||
if "fp16" in tags:
|
||||
header += R"""
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 530)
|
||||
#include <cuda_fp16.h>
|
||||
__device__ half max(half a, half b)
|
||||
{
|
||||
return __hgt(__half(a), __half(b)) ? a : b;
|
||||
}
|
||||
__device__ half min(half a, half b)
|
||||
{
|
||||
return __hlt(__half(a), __half(b)) ? a : b;
|
||||
}
|
||||
#endif // __CUDA_ARCH__ >= 530
|
||||
|
||||
// Pack two half values.
|
||||
static inline __device__ __host__ unsigned
|
||||
__pack_half2(const half x, const half y) {
|
||||
unsigned v0 = *((unsigned short *)&x);
|
||||
unsigned v1 = *((unsigned short *)&y);
|
||||
return (v1 << 16) | v0;
|
||||
}
|
||||
|
||||
#define CUDA_UNSUPPORTED_HALF_MATH_BINARY(HALF_MATH_NAME, FP32_MATH_NAME) \
|
||||
static inline __device__ __host__ half HALF_MATH_NAME(half x, half y) { \
|
||||
float tmp_x = __half2float(x); \
|
||||
float tmp_y = __half2float(y); \
|
||||
float result = FP32_MATH_NAME(tmp_x, tmp_y); \
|
||||
return __float2half(result); \
|
||||
}
|
||||
|
||||
#define CUDA_UNSUPPORTED_HALF_MATH_UNARY(HALF_MATH_NAME, FP32_MATH_NAME) \
|
||||
static inline __device__ __host__ half HALF_MATH_NAME(half x) { \
|
||||
float tmp_x = __half2float(x); \
|
||||
float result = FP32_MATH_NAME(tmp_x); \
|
||||
return __float2half(result); \
|
||||
}
|
||||
|
||||
// Some fp16 math functions are not supported in cuda_fp16.h,
|
||||
// so we define them here to make sure the generated CUDA code
|
||||
// is valid.
|
||||
#if defined(__CUDA_ARCH__)
|
||||
#if (__CUDA_ARCH__ >= 530)
|
||||
CUDA_UNSUPPORTED_HALF_MATH_BINARY(hpow, powf)
|
||||
#if ((__CUDACC_VER_MAJOR__ < 12) || ((__CUDACC_VER_MAJOR__ == 12) && (__CUDACC_VER_MINOR__ < 8)))
|
||||
CUDA_UNSUPPORTED_HALF_MATH_UNARY(htanh, tanhf)
|
||||
#endif
|
||||
CUDA_UNSUPPORTED_HALF_MATH_UNARY(htan, tanf)
|
||||
CUDA_UNSUPPORTED_HALF_MATH_UNARY(hatan, atanf)
|
||||
CUDA_UNSUPPORTED_HALF_MATH_UNARY(herf, erf)
|
||||
#else
|
||||
CUDA_UNSUPPORTED_HALF_MATH_UNARY(hexp, exp)
|
||||
#endif
|
||||
#endif
|
||||
|
||||
#undef CUDA_UNSUPPORTED_HALF_MATH_BINARY
|
||||
#undef CUDA_UNSUPPORTED_HALF_MATH_UNARY
|
||||
"""
|
||||
|
||||
if "bf16" in tags:
|
||||
header += R"""
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
|
||||
#include <cuda_bf16.h>
|
||||
__device__ nv_bfloat16 max(nv_bfloat16 a, nv_bfloat16 b)
|
||||
{
|
||||
return __hgt(a, b) ? a : b;
|
||||
}
|
||||
__device__ nv_bfloat16 min(nv_bfloat16 a, nv_bfloat16 b)
|
||||
{
|
||||
return __hlt(a, b) ? a : b;
|
||||
}
|
||||
#endif // __CUDA_ARCH__ >= 800
|
||||
// Pack two bfloat16 values.
|
||||
static inline __device__ __host__ unsigned
|
||||
__pack_nv_bfloat162(const nv_bfloat16 x, const nv_bfloat16 y) {
|
||||
unsigned v0 = *((unsigned short *)&x);
|
||||
unsigned v1 = *((unsigned short *)&y);
|
||||
return (v1 << 16) | v0;
|
||||
}
|
||||
|
||||
// Some bfp16 math functions are not supported in cuda_bfp16.h,
|
||||
// so we define them here to make sure the generated CUDA code
|
||||
// is valid.
|
||||
#define CUDA_UNSUPPORTED_HALF_MATH_BINARY(HALF_MATH_NAME, FP32_MATH_NAME) \
|
||||
static inline __device__ __host__ nv_bfloat16 HALF_MATH_NAME(nv_bfloat16 x, nv_bfloat16 y) { \
|
||||
float tmp_x = __bfloat162float(x); \
|
||||
float tmp_y = __bfloat162float(y); \
|
||||
float result = FP32_MATH_NAME(tmp_x, tmp_y); \
|
||||
return __float2bfloat16(result); \
|
||||
}
|
||||
|
||||
#define CUDA_UNSUPPORTED_HALF_MATH_UNARY(HALF_MATH_NAME, FP32_MATH_NAME) \
|
||||
static inline __device__ __host__ nv_bfloat16 HALF_MATH_NAME(nv_bfloat16 x) { \
|
||||
float tmp_x = __bfloat162float(x); \
|
||||
float result = FP32_MATH_NAME(tmp_x); \
|
||||
return __float2bfloat16(result); \
|
||||
}
|
||||
|
||||
CUDA_UNSUPPORTED_HALF_MATH_BINARY(hpow, powf)
|
||||
#if ((__CUDACC_VER_MAJOR__ < 12) || ((__CUDACC_VER_MAJOR__ == 12) && (__CUDACC_VER_MINOR__ < 8)))
|
||||
CUDA_UNSUPPORTED_HALF_MATH_UNARY(htanh, tanhf)
|
||||
#endif
|
||||
CUDA_UNSUPPORTED_HALF_MATH_UNARY(htan, tanf)
|
||||
CUDA_UNSUPPORTED_HALF_MATH_UNARY(hatan, atanf)
|
||||
CUDA_UNSUPPORTED_HALF_MATH_UNARY(herf, erf)
|
||||
|
||||
#undef CUDA_UNSUPPORTED_HALF_MATH_BINARY
|
||||
#undef CUDA_UNSUPPORTED_HALF_MATH_UNARY
|
||||
"""
|
||||
|
||||
if "fp8" in tags:
|
||||
header += R"""
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 890)
|
||||
#include <cuda_fp8.h>
|
||||
using fp8_e4_t = __nv_fp8_e4m3;
|
||||
using fp8_e4x2_t = __nv_fp8x2_e4m3;
|
||||
using fp8_e4x4_t = __nv_fp8x4_e4m3;
|
||||
struct fp8_e4x8_t {
|
||||
fp8_e4_t data[8];
|
||||
};
|
||||
struct fp8_e4x16_t {
|
||||
fp8_e4_t data[16];
|
||||
};
|
||||
using fp8_e5_t = __nv_fp8_e5m2;
|
||||
using fp8_e5x2_t = __nv_fp8x2_e5m2;
|
||||
using fp8_e5x4_t = __nv_fp8x4_e5m2;
|
||||
struct fp8_e5x8_t {
|
||||
fp8_e5_t data[8];
|
||||
};
|
||||
struct fp8_e5x16_t {
|
||||
fp8_e5_t data[16];
|
||||
};
|
||||
using fp8_e8_t = __nv_fp8_e8m0;
|
||||
using fp8_e8x2_t = __nv_fp8x2_e8m0;
|
||||
using fp8_e8x4_t = __nv_fp8x4_e8m0;
|
||||
struct fp8_e8x8_t {
|
||||
fp8_e8_t data[8];
|
||||
};
|
||||
struct fp8_e8x16_t {
|
||||
fp8_e8_t data[16];
|
||||
};
|
||||
#endif // __CUDA_ARCH__ >= 890
|
||||
"""
|
||||
|
||||
if "fp6" in tags:
|
||||
header += R"""
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
|
||||
#include <cuda_fp6.h>
|
||||
using fp6_e2_t = __nv_fp6_e2m3;
|
||||
using fp6_e2x2_t = __nv_fp6x2_e2m3;
|
||||
using fp6_e2x4_t = __nv_fp6x4_e2m3;
|
||||
struct fp6_e2x8_t {
|
||||
fp6_e2_t data[8];
|
||||
};
|
||||
struct fp6_e2x16_t {
|
||||
fp6_e2_t data[16];
|
||||
};
|
||||
using fp6_e3_t = __nv_fp6_e3m2;
|
||||
using fp6_e3x2_t = __nv_fp6x2_e3m2;
|
||||
using fp6_e3x4_t = __nv_fp6x4_e3m2;
|
||||
struct fp6_e3x8_t {
|
||||
fp6_e3_t data[8];
|
||||
};
|
||||
struct fp6_e3x16_t {
|
||||
fp6_e3_t data[16];
|
||||
};
|
||||
#endif // __CUDA_ARCH__ >= 1000
|
||||
"""
|
||||
|
||||
if "fp4" in tags:
|
||||
header += R"""
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
|
||||
#include <cuda_fp4.h>
|
||||
using fp4_e2_t = __nv_fp4_e2m1;
|
||||
using fp4_e2x2_t = __nv_fp4x2_e2m1;
|
||||
using fp4_e2x4_t = __nv_fp4x4_e2m1;
|
||||
struct fp4_e2x8_t {
|
||||
fp4_e2_t data[8];
|
||||
};
|
||||
struct fp4_e2x16_t {
|
||||
fp4_e2_t data[16];
|
||||
};
|
||||
#endif // __CUDA_ARCH__ >= 800
|
||||
"""
|
||||
|
||||
#########################################################
|
||||
# Vector type extensions
|
||||
#########################################################
|
||||
if "fp16" in tags or "bf16" in tags:
|
||||
header += R"""
|
||||
template <typename T, typename TVec2>
|
||||
struct __align__(8) half4_bfloat164 {
|
||||
T x, y, z, w;
|
||||
__host__ __device__ half4_bfloat164() : x(T(0)), y(T(0)), z(T(0)), w(T(0)) {}
|
||||
__host__ __device__ half4_bfloat164(T x, T y, T z, T w) : x(x), y(y), z(z), w(w) {}
|
||||
"""
|
||||
if "fp8" in tags:
|
||||
header += R"""
|
||||
__host__ __device__ explicit half4_bfloat164(const __nv_fp8x4_e4m3& fp8x4) {
|
||||
if constexpr (std::is_same_v<T, __half>) {
|
||||
__nv_fp8x2_e4m3 lo_part, hi_part;
|
||||
lo_part.__x = static_cast<__nv_fp8x2_storage_t>(fp8x4.__x & 0xFFFF);
|
||||
hi_part.__x = static_cast<__nv_fp8x2_storage_t>((fp8x4.__x >> 16) & 0xFFFF);
|
||||
TVec2 lo_half2 = static_cast<TVec2>(lo_part);
|
||||
TVec2 hi_half2 = static_cast<TVec2>(hi_part);
|
||||
x = reinterpret_cast<T*>(&lo_half2)[0];
|
||||
y = reinterpret_cast<T*>(&lo_half2)[1];
|
||||
z = reinterpret_cast<T*>(&hi_half2)[0];
|
||||
w = reinterpret_cast<T*>(&hi_half2)[1];
|
||||
} else {
|
||||
__nv_fp8_storage_t elem0_raw = static_cast<__nv_fp8_storage_t>(fp8x4.__x & 0xFF);
|
||||
__nv_fp8_storage_t elem1_raw = static_cast<__nv_fp8_storage_t>((fp8x4.__x >> 8) & 0xFF);
|
||||
__nv_fp8_storage_t elem2_raw = static_cast<__nv_fp8_storage_t>((fp8x4.__x >> 16) & 0xFF);
|
||||
__nv_fp8_storage_t elem3_raw = static_cast<__nv_fp8_storage_t>((fp8x4.__x >> 24) & 0xFF);
|
||||
__nv_fp8_e4m3 elem0, elem1, elem2, elem3;
|
||||
elem0.__x = elem0_raw;
|
||||
elem1.__x = elem1_raw;
|
||||
elem2.__x = elem2_raw;
|
||||
elem3.__x = elem3_raw;
|
||||
x = T(elem0);
|
||||
y = T(elem1);
|
||||
z = T(elem2);
|
||||
w = T(elem3);
|
||||
}
|
||||
}
|
||||
__host__ __device__ explicit operator __nv_fp8x4_e4m3() const {
|
||||
__nv_fp8x4_e4m3 result;
|
||||
TVec2 lo_half2 = *reinterpret_cast<const TVec2*>(&x);
|
||||
TVec2 hi_half2 = *reinterpret_cast<const TVec2*>(&z);
|
||||
__nv_fp8x2_e4m3 lo_part(lo_half2), hi_part(hi_half2);
|
||||
result.__x =
|
||||
(static_cast<uint32_t>(lo_part.__x) | (static_cast<uint32_t>(hi_part.__x) << 16));
|
||||
return result;
|
||||
}
|
||||
__host__ __device__ explicit half4_bfloat164(const __nv_fp8x4_e5m2& fp8x4) {
|
||||
__nv_fp8x2_e5m2 lo_part, hi_part;
|
||||
lo_part.__x = static_cast<__nv_fp8x2_storage_t>(fp8x4.__x & 0xFFFF);
|
||||
hi_part.__x = static_cast<__nv_fp8x2_storage_t>((fp8x4.__x >> 16) & 0xFFFF);
|
||||
TVec2 lo_half2 = static_cast<TVec2>(lo_part);
|
||||
TVec2 hi_half2 = static_cast<TVec2>(hi_part);
|
||||
x = reinterpret_cast<T*>(&lo_half2)[0];
|
||||
y = reinterpret_cast<T*>(&lo_half2)[1];
|
||||
z = reinterpret_cast<T*>(&hi_half2)[0];
|
||||
w = reinterpret_cast<T*>(&hi_half2)[1];
|
||||
}
|
||||
__host__ __device__ explicit operator __nv_fp8x4_e5m2() const {
|
||||
__nv_fp8x4_e5m2 result;
|
||||
TVec2 lo_half2 = *reinterpret_cast<const TVec2*>(&x);
|
||||
TVec2 hi_half2 = *reinterpret_cast<const TVec2*>(&z);
|
||||
__nv_fp8x2_e5m2 lo_part(lo_half2), hi_part(hi_half2);
|
||||
result.__x =
|
||||
(static_cast<uint32_t>(lo_part.__x) | (static_cast<uint32_t>(hi_part.__x) << 16));
|
||||
return result;
|
||||
}
|
||||
__host__ __device__ explicit half4_bfloat164(const __nv_fp8x4_e8m0& fp8x4) {
|
||||
__nv_fp8x2_e8m0 lo_part, hi_part;
|
||||
lo_part.__x = static_cast<__nv_fp8x2_storage_t>(fp8x4.__x & 0xFFFF);
|
||||
hi_part.__x = static_cast<__nv_fp8x2_storage_t>((fp8x4.__x >> 16) & 0xFFFF);
|
||||
TVec2 lo_half2 = static_cast<TVec2>(lo_part);
|
||||
TVec2 hi_half2 = static_cast<TVec2>(hi_part);
|
||||
x = reinterpret_cast<T*>(&lo_half2)[0];
|
||||
y = reinterpret_cast<T*>(&lo_half2)[1];
|
||||
z = reinterpret_cast<T*>(&hi_half2)[0];
|
||||
w = reinterpret_cast<T*>(&hi_half2)[1];
|
||||
}
|
||||
__host__ __device__ explicit operator __nv_fp8x4_e8m0() const {
|
||||
__nv_fp8x4_e8m0 result;
|
||||
TVec2 lo_half2 = *reinterpret_cast<const TVec2*>(&x);
|
||||
TVec2 hi_half2 = *reinterpret_cast<const TVec2*>(&z);
|
||||
__nv_fp8x2_e8m0 lo_part(lo_half2), hi_part(hi_half2);
|
||||
result.__x =
|
||||
(static_cast<uint32_t>(lo_part.__x) | (static_cast<uint32_t>(hi_part.__x) << 16));
|
||||
return result;
|
||||
}
|
||||
"""
|
||||
if "fp4" in tags:
|
||||
header += R"""
|
||||
__host__ __device__ explicit half4_bfloat164(const __nv_fp4x4_e2m1& fp4x4) {
|
||||
if constexpr (std::is_same_v<T, __half>) {
|
||||
__nv_fp4x2_storage_t lo_part = static_cast<__nv_fp4x2_storage_t>(fp4x4.__x & 0xFF);
|
||||
__nv_fp4x2_storage_t hi_part = static_cast<__nv_fp4x2_storage_t>((fp4x4.__x >> 8) & 0xFF);
|
||||
TVec2 lo_half2 = __half2(__nv_cvt_fp4x2_to_halfraw2(lo_part, __NV_E2M1));
|
||||
TVec2 hi_half2 = __half2(__nv_cvt_fp4x2_to_halfraw2(hi_part, __NV_E2M1));
|
||||
x = reinterpret_cast<T*>(&lo_half2)[0];
|
||||
y = reinterpret_cast<T*>(&lo_half2)[1];
|
||||
z = reinterpret_cast<T*>(&hi_half2)[0];
|
||||
w = reinterpret_cast<T*>(&hi_half2)[1];
|
||||
} else {
|
||||
__nv_fp4_e2m1 elem0, elem1, elem2, elem3;
|
||||
elem0.__x = static_cast<__nv_fp4_storage_t>(fp4x4.__x & 0xF);
|
||||
elem1.__x = static_cast<__nv_fp4_storage_t>((fp4x4.__x >> 4) & 0xF);
|
||||
elem2.__x = static_cast<__nv_fp4_storage_t>((fp4x4.__x >> 8) & 0xF);
|
||||
elem3.__x = static_cast<__nv_fp4_storage_t>((fp4x4.__x >> 12) & 0xF);
|
||||
x = T(elem0);
|
||||
y = T(elem1);
|
||||
z = T(elem2);
|
||||
w = T(elem3);
|
||||
}
|
||||
}
|
||||
__host__ __device__ explicit operator __nv_fp4x4_e2m1() const {
|
||||
TVec2 lo_half2 = *reinterpret_cast<const TVec2*>(&x);
|
||||
TVec2 hi_half2 = *reinterpret_cast<const TVec2*>(&z);
|
||||
return __nv_fp4x4_e2m1(lo_half2, hi_half2);
|
||||
}
|
||||
"""
|
||||
header += R"""
|
||||
};
|
||||
"""
|
||||
if "fp16" in tags:
|
||||
header += R"""
|
||||
using half4 = half4_bfloat164<__half, __half2>;
|
||||
__host__ __device__ half4 make_half4(__half x, __half y, __half z, __half w) {
|
||||
return half4(x, y, z, w);
|
||||
}
|
||||
"""
|
||||
if "bf16" in tags:
|
||||
header += R"""
|
||||
using nv_bfloat164 = half4_bfloat164<nv_bfloat16, nv_bfloat162>;
|
||||
__host__ __device__ nv_bfloat164 make_nv_bfloat164(nv_bfloat16 x, nv_bfloat16 y, nv_bfloat16 z, nv_bfloat16 w) {
|
||||
return nv_bfloat164(x, y, z, w);
|
||||
}
|
||||
__host__ __device__ nv_bfloat162 make_nv_bfloat162(nv_bfloat16 x, nv_bfloat16 y) {
|
||||
return nv_bfloat162(x, y);
|
||||
}
|
||||
""" # noqa: E501
|
||||
if "fp8" in tags:
|
||||
header += R"""
|
||||
__host__ __device__ nv_bfloat162 cast_to_nv_bfloat162(const __nv_fp8x2_e4m3& fp8x2) {
|
||||
__nv_fp8_e4m3 elem0, elem1;
|
||||
elem0.__x = static_cast<__nv_fp8_storage_t>(fp8x2.__x & 0xFF);
|
||||
elem1.__x = static_cast<__nv_fp8_storage_t>((fp8x2.__x >> 8) & 0xFF);
|
||||
nv_bfloat16 x = nv_bfloat16(elem0);
|
||||
nv_bfloat16 y = nv_bfloat16(elem1);
|
||||
return nv_bfloat162(x, y);
|
||||
}
|
||||
__host__ __device__ nv_bfloat162 cast_to_nv_bfloat162(const __nv_fp8x2_e5m2& fp8x2) {
|
||||
__nv_fp8_e5m2 elem0, elem1;
|
||||
elem0.__x = static_cast<__nv_fp8_storage_t>(fp8x2.__x & 0xFF);
|
||||
elem1.__x = static_cast<__nv_fp8_storage_t>((fp8x2.__x >> 8) & 0xFF);
|
||||
nv_bfloat16 x = nv_bfloat16(elem0);
|
||||
nv_bfloat16 y = nv_bfloat16(elem1);
|
||||
return nv_bfloat162(x, y);
|
||||
}
|
||||
__host__ __device__ nv_bfloat162 cast_to_nv_bfloat162(const __nv_fp8x2_e8m0& fp8x2) {
|
||||
__nv_fp8_e8m0 elem0, elem1;
|
||||
elem0.__x = static_cast<__nv_fp8_storage_t>(fp8x2.__x & 0xFF);
|
||||
elem1.__x = static_cast<__nv_fp8_storage_t>((fp8x2.__x >> 8) & 0xFF);
|
||||
nv_bfloat16 x = nv_bfloat16(elem0);
|
||||
nv_bfloat16 y = nv_bfloat16(elem1);
|
||||
return nv_bfloat162(x, y);
|
||||
}
|
||||
"""
|
||||
if "fp8" in tags:
|
||||
header += R"""
|
||||
__device__ __nv_fp8x2_e5m2 make___nv_fp8x2_e5m2(__nv_fp8_e5m2 x, __nv_fp8_e5m2 y) {
|
||||
__nv_fp8x2_e5m2 result;
|
||||
result.__x = (x.__x) | (y.__x << 8);
|
||||
return result;
|
||||
}
|
||||
__device__ __nv_fp8x4_e5m2 make___nv_fp8x4_e5m2(__nv_fp8_e5m2 a, __nv_fp8_e5m2 b, __nv_fp8_e5m2 c, __nv_fp8_e5m2 d) {
|
||||
__nv_fp8x4_e5m2 result;
|
||||
result.__x = (a.__x) | (b.__x << 8) | (c.__x << 16) | (d.__x << 24);
|
||||
return result;
|
||||
}
|
||||
__device__ __nv_fp8x2_e4m3 make___nv_fp8x2_e4m3(__nv_fp8_e4m3 x, __nv_fp8_e4m3 y) {
|
||||
__nv_fp8x2_e4m3 result;
|
||||
result.__x = (x.__x) | (y.__x << 8);
|
||||
return result;
|
||||
}
|
||||
__device__ __nv_fp8x4_e4m3 make___nv_fp8x4_e4m3(__nv_fp8_e4m3 a, __nv_fp8_e4m3 b, __nv_fp8_e4m3 c, __nv_fp8_e4m3 d) {
|
||||
__nv_fp8x4_e4m3 result;
|
||||
result.__x = (a.__x) | (b.__x << 8) | (c.__x << 16) | (d.__x << 24);
|
||||
return result;
|
||||
}
|
||||
__device__ __nv_fp8x2_e8m0 make___nv_fp8x2_e8m0(__nv_fp8_e8m0 x, __nv_fp8_e8m0 y) {
|
||||
__nv_fp8x2_e8m0 result;
|
||||
result.__x = (x.__x) | (y.__x << 8);
|
||||
return result;
|
||||
}
|
||||
__device__ __nv_fp8x4_e8m0 make___nv_fp8x4_e8m0(__nv_fp8_e8m0 a, __nv_fp8_e8m0 b, __nv_fp8_e8m0 c, __nv_fp8_e8m0 d) {
|
||||
__nv_fp8x4_e8m0 result;
|
||||
result.__x = (a.__x) | (b.__x << 8) | (c.__x << 16) | (d.__x << 24);
|
||||
return result;
|
||||
}
|
||||
""" # noqa: E501
|
||||
if "fp4" in tags:
|
||||
header += R"""
|
||||
__host__ __device__ nv_bfloat162 cast_to_nv_bfloat162(const __nv_fp4x2_e2m1& fp4x2) {
|
||||
__nv_fp4_e2m1 elem0, elem1;
|
||||
elem0.__x = static_cast<__nv_fp4_storage_t>(fp4x2.__x & 0xFF);
|
||||
elem1.__x = static_cast<__nv_fp4_storage_t>((fp4x2.__x >> 8) & 0xFF);
|
||||
nv_bfloat16 x = nv_bfloat16(elem0);
|
||||
nv_bfloat16 y = nv_bfloat16(elem1);
|
||||
return nv_bfloat162(x, y);
|
||||
}
|
||||
"""
|
||||
|
||||
if "int8" in tags:
|
||||
header += R"""
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 610)
|
||||
#include <sm_61_intrinsics.h>
|
||||
|
||||
#if defined(__CUDACC_RTC__)
|
||||
#define __SM_61_INTRINSICS_DECL__ __device__
|
||||
#else /* !__CUDACC_RTC__ */
|
||||
#define __SM_61_INTRINSICS_DECL__ static __device__ __inline__
|
||||
#endif /* __CUDACC_RTC__ */
|
||||
|
||||
#ifndef __CUDA_ARCH__
|
||||
#define __DEF_IF_HOST { }
|
||||
#else /* !__CUDA_ARCH__ */
|
||||
#define __DEF_IF_HOST ;
|
||||
#endif /* __CUDA_ARCH__ */
|
||||
|
||||
__SM_61_INTRINSICS_DECL__ int __dp4a(unsigned int srcA, int srcB, int c) __DEF_IF_HOST
|
||||
__SM_61_INTRINSICS_DECL__ int __dp4a(int srcA, unsigned int srcB, int c) __DEF_IF_HOST
|
||||
|
||||
#undef __DEF_IF_HOST
|
||||
|
||||
#if !defined(__CUDACC_RTC__) && defined(__CUDA_ARCH__)
|
||||
__SM_61_INTRINSICS_DECL__ int __dp4a(unsigned int srcA, int srcB, int c) {
|
||||
int ret;
|
||||
asm volatile ("dp4a.u32.s32 %0, %1, %2, %3;" : "=r"(ret) : "r"(srcA), "r"(srcB), "r"(c));
|
||||
return ret;
|
||||
}
|
||||
|
||||
__SM_61_INTRINSICS_DECL__ int __dp4a(int srcA, unsigned int srcB, int c) {
|
||||
int ret;
|
||||
asm volatile ("dp4a.s32.u32 %0, %1, %2, %3;" : "=r"(ret) : "r"(srcA), "r"(srcB), "r"(c));
|
||||
return ret;
|
||||
}
|
||||
#endif /* !__CUDACC_RTC__ && defined(__CUDA_ARCH__) */
|
||||
|
||||
#undef __SM_61_INTRINSICS_DECL__
|
||||
|
||||
#endif // __CUDA_ARCH__ >= 610
|
||||
"""
|
||||
if "math_constants" in tags:
|
||||
header += R"""
|
||||
#include <math_constants.h>
|
||||
"""
|
||||
if "mma" in tags:
|
||||
header += R"""
|
||||
#include <mma.h>
|
||||
"""
|
||||
|
||||
if "warp_shuffle" in tags:
|
||||
header += R"""
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 700)
|
||||
#define __shfl_sync(mask, var, lane, width) \
|
||||
__shfl((var), (lane), (width))
|
||||
|
||||
#define __shfl_down_sync(mask, var, offset, width) \
|
||||
__shfl_down((var), (offset), (width))
|
||||
|
||||
#define __shfl_up_sync(mask, var, offset, width) \
|
||||
__shfl_up((var), (offset), (width))
|
||||
#endif
|
||||
"""
|
||||
|
||||
if "cast_smem_ptr_to_int" in tags:
|
||||
header += R"""
|
||||
__forceinline__ __device__ unsigned int cast_smem_ptr_to_int(const void* const smem_ptr) {
|
||||
unsigned int smem_int;
|
||||
asm volatile ("{ .reg .u64 smem_int; cvta.to.shared.u64 smem_int, %1; cvt.u32.u64 %0, smem_int; }"
|
||||
: "=r"(smem_int) : "l"(smem_ptr));
|
||||
return smem_int;
|
||||
}
|
||||
"""
|
||||
header += R"""
|
||||
#if (((__CUDACC_VER_MAJOR__ == 11) && (__CUDACC_VER_MINOR__ >= 4)) || \
|
||||
(__CUDACC_VER_MAJOR__ > 11))
|
||||
#define TVM_ENABLE_L2_PREFETCH 1
|
||||
#else
|
||||
#define TVM_ENABLE_L2_PREFETCH 0
|
||||
#endif
|
||||
|
||||
#ifdef _WIN32
|
||||
using uint = unsigned int;
|
||||
using uchar = unsigned char;
|
||||
using ushort = unsigned short;
|
||||
using int64_t = long long;
|
||||
using uint64_t = unsigned long long;
|
||||
#else
|
||||
#define uint unsigned int
|
||||
#define uchar unsigned char
|
||||
#define ushort unsigned short
|
||||
#endif
|
||||
"""
|
||||
|
||||
if "get_tmem_addr" in tags:
|
||||
header += R"""
|
||||
__forceinline__ __device__ uint32_t get_tmem_addr(uint32_t idx, int row_offset, int col_offset) {
|
||||
int col_idx = idx & 0xFFFF;
|
||||
int row_idx = (idx >> 16) & 0xFFFF;
|
||||
col_idx += col_offset;
|
||||
row_idx += row_offset;
|
||||
col_idx = col_idx & 0xFFFF;
|
||||
row_idx = row_idx & 0xFFFF;
|
||||
|
||||
uint32_t new_idx = (row_idx << 16) | col_idx;
|
||||
return new_idx;
|
||||
}
|
||||
"""
|
||||
|
||||
if "get_time_stamp" in tags:
|
||||
header += R"""
|
||||
__forceinline__ __device__ uint32_t tvm_builtin_get_timestamp() {
|
||||
volatile uint32_t ret;
|
||||
asm volatile("mov.u32 %0, %globaltimer_lo;" : "=r"(ret));
|
||||
return ret;
|
||||
}
|
||||
"""
|
||||
|
||||
if "gmma_descriptor" in tags:
|
||||
header += R"""
|
||||
#ifndef HOST_DEVICE
|
||||
#define HOST_DEVICE __forceinline__ __host__ __device__
|
||||
#endif
|
||||
union GmmaDescriptor
|
||||
{
|
||||
HOST_DEVICE constexpr
|
||||
GmmaDescriptor() noexcept : desc_(0) {}
|
||||
HOST_DEVICE constexpr
|
||||
GmmaDescriptor(uint64_t desc) noexcept : desc_(desc) {}
|
||||
HOST_DEVICE constexpr
|
||||
GmmaDescriptor(GmmaDescriptor const& t) noexcept : desc_(t.desc_) {}
|
||||
HOST_DEVICE constexpr
|
||||
GmmaDescriptor(GmmaDescriptor && t) noexcept : desc_(t.desc_) {}
|
||||
|
||||
HOST_DEVICE constexpr
|
||||
GmmaDescriptor& operator=(GmmaDescriptor const& t) noexcept {
|
||||
desc_ = t.desc_;
|
||||
return *this;
|
||||
}
|
||||
|
||||
HOST_DEVICE constexpr
|
||||
GmmaDescriptor& operator=(GmmaDescriptor && t) noexcept {
|
||||
desc_ = t.desc_;
|
||||
return *this;
|
||||
}
|
||||
|
||||
uint64_t desc_;
|
||||
uint32_t reg32_[2];
|
||||
uint16_t reg16_[4];
|
||||
|
||||
// Bitfield implementation avoids the need for shifts in assignment
|
||||
struct {
|
||||
// start_address, bit [0,14), 4LSB not included
|
||||
uint16_t start_address_ : 14, : 2; // 14 bits [0,14), 2 bits unused
|
||||
// leading dimension byte offset, bit [16,30), 4LSB not included
|
||||
// For N: This is the stride from the first col to the second col of the 8x2 brick in INTERLEAVED
|
||||
// Unused for all SWIZZLE_* layouts (and assumed to be 1)
|
||||
// For T: This is the stride from the first 8 rows to the next 8 rows.
|
||||
uint16_t leading_byte_offset_ : 14, : 2; // 14 bits [0,14), 2 bits unused
|
||||
// stride dimension byte offset, bit [32,46), 4LSB not included
|
||||
// For N: This is the stride from the first 8 rows to the next 8 rows.
|
||||
// For T: This is the stride fro mthe first 8 cols to the next 8 cols.
|
||||
uint16_t stride_byte_offset_ : 14, : 2; // 14 bits [0,14), 2 bits unused
|
||||
// base_offset, bit [49,52)
|
||||
// Valid only for SWIZZLE_128B and SWIZZLE_64B
|
||||
uint8_t : 1, base_offset_ : 3, : 4; // 1 bit unused, 3 bits [1,4), 4 bits unused
|
||||
// layout type, bit [62,64)
|
||||
// SWIZZLE_NONE = 0, SWIZZLE_32B = 3, SWIZZLE_64B = 2, SWIZZLE_128B = 1
|
||||
uint8_t : 6, layout_type_ : 2; // 6 bits unused, 2 bits [6,8)
|
||||
} bitfield;
|
||||
|
||||
// Decay to a uint64_t
|
||||
HOST_DEVICE constexpr
|
||||
operator uint64_t() const noexcept { return desc_; }
|
||||
};
|
||||
""" # noqa: E501
|
||||
|
||||
if "smem_descriptor" in tags:
|
||||
header += R"""
|
||||
#ifndef HOST_DEVICE
|
||||
#define HOST_DEVICE __forceinline__ __host__ __device__
|
||||
#endif
|
||||
union SmemDescriptor
|
||||
{
|
||||
uint64_t desc_ = 0;
|
||||
// Bitfield implementation avoids the need for shifts in assignment
|
||||
struct {
|
||||
// start_address, bit [0,14), 4LSB not included
|
||||
uint16_t start_address_ : 14, : 2; // 14 bits [0,14), 2 bits unused
|
||||
// leading dimension byte offset, bit [16,30), 4LSB not included
|
||||
uint16_t leading_byte_offset_ : 14, : 2; // 14 bits [0,14), 2 bits unused
|
||||
// stride dimension byte offset, bit [32,46), 4LSB not included
|
||||
uint16_t stride_byte_offset_ : 14, version_ : 2; // 14 bits [0,14), 2 bits [14,16)
|
||||
// base_offset, bit [49,52). leading_byte_offset_mode, bit [52,53).
|
||||
uint8_t : 1, base_offset_ : 3, lbo_mode_ : 1, : 3; // 1 bit unused, 3 bits [1,4), 1 bit [4,5), 3 bits unused
|
||||
// layout type, bit [61,64), SWIZZLE_NONE matrix descriptor = 0, SWIZZLE_128B matrix descriptor = 2, SWIZZLE_64B descriptor = 4, SWIZZLE_32B descriptor = 6, SWIZZLE_128B_BASE32B = 1, N/A = 3, N/A = 5, N/A = 7
|
||||
uint8_t : 5, layout_type_ : 3; // 6 bits unused, 3 bits [5,8)
|
||||
};
|
||||
// Seperate the field, as we may only update one part of desc
|
||||
struct {
|
||||
uint32_t lo;
|
||||
uint32_t hi;
|
||||
};
|
||||
|
||||
// Decay to a uint64_t
|
||||
HOST_DEVICE constexpr
|
||||
operator uint64_t() const noexcept { return desc_; }
|
||||
};
|
||||
""" # noqa: E501
|
||||
|
||||
if "instr_descriptor" in tags:
|
||||
header += R"""
|
||||
#ifndef HOST_DEVICE
|
||||
#define HOST_DEVICE __forceinline__ __host__ __device__
|
||||
#endif
|
||||
union InstrDescriptor
|
||||
{
|
||||
uint32_t desc_;
|
||||
|
||||
struct {
|
||||
// Bitfield implementation avoids the need for shifts in assignment
|
||||
uint16_t sparse_id2_ : 2, // bit [ 0, 2) : Sparse meta data id2
|
||||
sparse_flag_ : 1, // bit [ 2, 3) : 0 = dense. 1 = sparse. 1 value valid only for F32F16/S8/MXF8F6F4
|
||||
saturate_ : 1, // bit [ 3, 4) : 0 = no saturate. 1 = saturate. 1 value valid only for S8
|
||||
c_format_ : 2, // bit [ 4, 6) : 0 = F16. 1 = F32, 2 = S32
|
||||
: 1, //
|
||||
a_format_ : 3, // bit [ 7,10) : MXF8F6F4Format:0 = E4M3, 1 = E5M2, 3 = E2M3, 4 = E3M2, 5 = E2M1. F32F16Format: 0 = F16, 1 = BF16, 2 = TF32. S8: 0 unsigned 8 bit, 1 signed 8 bit. Boolean MMA: 0 Boolean
|
||||
b_format_ : 3, // bit [10,13) : MXF8F6F4Format:0 = E4M3, 1 = E5M2, 3 = E2M3, 4 = E3M2, 5 = E2M1. F32F16Format: 0 = F16, 1 = BF16, 2 = TF32. S8: 0 unsigned 8 bit, 1 signed 8 bit. Boolean MMA: 0 Boolean
|
||||
a_negate_ : 1, // bit [13,14) : 0 = no negate. 1 = negate. 1 value valid only for F32F16Format and MXF8F6F4Format
|
||||
b_negate_ : 1, // bit [14,15) : 0 = no negate. 1 = negate. 1 value valid only for F32F16Format and MXF8F6F4Format
|
||||
a_major_ : 1; // bit [15,16) : 0 = K-major. 1 = MN-major. Major value of 1 is only valid for E4M3, E5M2, INT8 (signed and unsigned), F16, BF16 and TF32 source formats
|
||||
uint16_t b_major_ : 1, // bit [16,17) : 0 = K-major. 1 = MN-major. Major value of 1 is only valid for E4M3, E5M2, INT8 (signed and unsigned), F16, BF16 and TF32 source formats
|
||||
n_dim_ : 6, // bit [17,23) : 3 LSBs not included. Valid values range from 1 (N=8) to 32 (N=256). All values are not valid for all instruction formats
|
||||
: 1, //
|
||||
m_dim_ : 5, // bit [24,29) : 4 LSBs not included. Valid values are: 4 (M=64), 8 (M=128), 16 (M=256)
|
||||
: 1, //
|
||||
max_shift_ : 2; // bit [30,32) : Maximum shift for WS instruction. Encoded as follows: 0 = no shift, 1 = maximum shift of 8, 2 = maximum shift of 16, 3 = maximum shift of 32.
|
||||
};
|
||||
|
||||
// Decay to a uint32_t
|
||||
HOST_DEVICE constexpr explicit
|
||||
operator uint32_t() const noexcept { return desc_; }
|
||||
};
|
||||
""" # noqa: E501
|
||||
|
||||
if "instr_descriptor_block_scaled" in tags:
|
||||
header += R"""
|
||||
#ifndef HOST_DEVICE
|
||||
#define HOST_DEVICE __forceinline__ __host__ __device__
|
||||
#endif
|
||||
union InstrDescriptorBlockScaled
|
||||
{
|
||||
uint32_t desc_;
|
||||
|
||||
struct {
|
||||
// Bitfield implementation avoids the need for shifts in assignment
|
||||
uint16_t sparse_id2_ : 2, // bit [ 0, 2) : Sparse meta data id2
|
||||
sparse_flag_ : 1, // bit [ 2, 3) : 0 = dense. 1 = sparse. 1 value valid only for F32F16/S8/MXF8F6F4
|
||||
: 1, //
|
||||
b_sf_id_ : 2, // bit [ 4, 6) : Matrix B Scale Factor ID
|
||||
: 1, //
|
||||
a_format_ : 3, // bit [ 7, 9) : MXF8F6F4Format:0 = E4M3, 1 = E5M2, 3 = E2M3, 4 = E3M2, 5 = E2M1. F32F16Format: 0 = F16, 1 = BF16, 2 = TF32. S8: 0 unsigned 8 bit, 1 signed 8 bit. BMMA: 0 Boolean
|
||||
b_format_ : 3, // bit [10,12) : MXF8F6F4Format:0 = E4M3, 1 = E5M2, 3 = E2M3, 4 = E3M2, 5 = E2M1. F32F16Format: 0 = F16, 1 = BF16, 2 = TF32. S8: 0 unsigned 8 bit, 1 signed 8 bit. BMMA: 0 Boolean
|
||||
a_negate_ : 1, // bit [13,14) : 0 = no negate. 1 = negate. 1 value valid only for F32F16Format and MXF8F6F4Format
|
||||
b_negate_ : 1, // bit [14,15) : 0 = no negate. 1 = negate. 1 value valid only for F32F16Format and MXF8F6F4Format
|
||||
a_major_ : 1; // bit [15,16) : 0 = K-major. 1 = MN-major. Major value of 1 is only valid for E4M3, E5M2, INT8 (signed and unsigned), F16, BF16 and TF32 source formats
|
||||
uint16_t b_major_ : 1, // bit [16,17) : 0 = K-major. 1 = MN-major. Major value of 1 is only valid for E4M3, E5M2, INT8 (signed and unsigned), F16, BF16 and TF32 source formats
|
||||
n_dim_ : 6, // bit [17,23) : 3 LSBs not included. Valid values range from 1 (N=8) to 32 (N=256). All values are not valid for all instruction formats
|
||||
scale_format_ : 1, // bit [23,24) : 0=E4M3, 1=E8M0
|
||||
m_dim_ : 5, // bit [24,29) : 4 LSBs not included. Valid values are: 4 (M=64), 8 (M=128), 16 (M=256)
|
||||
a_sf_id_ : 2, // bit [29,31) : Matrix A Scale Factor ID
|
||||
: 1; //
|
||||
};
|
||||
|
||||
// Decay to a uint32_t
|
||||
HOST_DEVICE constexpr
|
||||
operator uint32_t() const noexcept { return desc_; }
|
||||
};
|
||||
""" # noqa: E501
|
||||
|
||||
if "elect_one_sync" in tags:
|
||||
header += R"""
|
||||
__forceinline__ __device__ uint32_t tvm_builtin_elect_one_sync() {{
|
||||
uint32_t pred = 0;
|
||||
uint32_t laneid = 0;
|
||||
asm volatile(
|
||||
"{\n"
|
||||
".reg .b32 %%rx;\n"
|
||||
".reg .pred %%px;\n"
|
||||
" elect.sync %%rx|%%px, %2;\n"
|
||||
"@%%px mov.s32 %1, 1;\n"
|
||||
" mov.s32 %0, %%rx;\n"
|
||||
"}\n"
|
||||
: "+r"(laneid), "+r"(pred)
|
||||
: "r"(0xFFFFFFFF));
|
||||
return pred;
|
||||
}}
|
||||
"""
|
||||
return header
|
||||
@@ -0,0 +1,501 @@
|
||||
# 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
|
||||
"""Math intrinsics.
|
||||
|
||||
PTX side:
|
||||
* ``add{.rnd}{.ftz}.f32x2`` / ``sub`` / ``mul`` / ``fma`` — packed f32x2.
|
||||
* ``ex2.approx.ftz.f32`` / ``rcp.approx.ftz.f32`` — special functions.
|
||||
* ``max.f32`` / ``min.f32`` — 3-operand reduction form.
|
||||
|
||||
CUDA side:
|
||||
* warp / CTA reductions (templated butterfly shuffle-XOR).
|
||||
"""
|
||||
|
||||
from tvm.backend.cuda.op import cuda_func_call
|
||||
|
||||
from ._schema import device_intrinsic
|
||||
from .registry import register_codegen
|
||||
from .utils import parse_str, validate_power_of_two_range
|
||||
|
||||
# =============================================================================
|
||||
# Packed f32x2 arithmetic — `add{.rnd}{.ftz}.f32x2 d, a, b ;` and friends.
|
||||
# Inputs are packed into a `.b64` register (low half = elem 0, high half =
|
||||
# elem 1); the body packs/unpacks via ``make_float2`` + ``reinterpret_cast``.
|
||||
# =============================================================================
|
||||
|
||||
# PTX add/sub/mul/fma over (f32 | f32x2 | f64), DPS form.
|
||||
# add{.rnd}{.ftz}{.sat}.f32 [d], a, b
|
||||
# add{.rnd}{.ftz}.f32x2 [d], a, b (a,b are packed-as-u64)
|
||||
# add{.rnd}.f64 [d], a, b
|
||||
# (sub / mul same shape; fma adds a `c` operand)
|
||||
# Inputs a/b/c are register operands (scalar fp32 / packed u64 / scalar fp64).
|
||||
# Result is written through `d` (a pointer).
|
||||
_PACKED_ROUNDING = ("rz", "rn", "rm", "rp")
|
||||
|
||||
|
||||
# Per-dtype operand types and asm constraints.
|
||||
# - c_in: C type of input register operand (matches PTX register type)
|
||||
# - out_cast: pointer cast applied at d_addr (callers may pass float*/double*/...)
|
||||
# - in_cstr / out_cstr: GCC asm constraint letter
|
||||
_DTYPE_INFO = {
|
||||
"f32": {"c_in": "float", "out_cast": "float*", "in_cstr": "f", "out_cstr": "f"},
|
||||
"f32x2": {
|
||||
"c_in": "unsigned long long",
|
||||
"out_cast": "uint64_t*",
|
||||
"in_cstr": "l",
|
||||
"out_cstr": "l",
|
||||
},
|
||||
"f64": {"c_in": "double", "out_cast": "double*", "in_cstr": "d", "out_cstr": "d"},
|
||||
}
|
||||
|
||||
|
||||
def _ptx_arith_modifier_string(dtype, rounding, ftz, sat):
|
||||
"""Build the `.rnd.ftz.sat` modifier substring + name suffix."""
|
||||
rnd = parse_str(rounding)
|
||||
assert rnd in _PACKED_ROUNDING, f"invalid rounding {rnd!r}, expected one of {_PACKED_ROUNDING}"
|
||||
ftz_b = bool(int(ftz)) if hasattr(ftz, "value") else bool(ftz)
|
||||
sat_b = bool(int(sat)) if hasattr(sat, "value") else bool(sat)
|
||||
if dtype == "f64" and (ftz_b or sat_b):
|
||||
raise ValueError("PTX <op>.f64 does not accept .ftz or .sat")
|
||||
if dtype == "f32x2" and sat_b:
|
||||
raise ValueError("PTX <op>.f32x2 does not accept .sat")
|
||||
mod = f".{rnd}"
|
||||
if ftz_b:
|
||||
mod += ".ftz"
|
||||
if sat_b:
|
||||
mod += ".sat"
|
||||
name_suffix = f"_{rnd}"
|
||||
if ftz_b:
|
||||
name_suffix += "_ftz"
|
||||
if sat_b:
|
||||
name_suffix += "_sat"
|
||||
return mod, name_suffix
|
||||
|
||||
|
||||
def _ptx_binary_arith_parts(op, dtype):
|
||||
"""Return (name_fn, sig, body_fn) for ptx_{op}_{dtype} binary form."""
|
||||
info = _DTYPE_INFO[dtype]
|
||||
# Destination is ``void*`` so callers can pass any element-type pointer
|
||||
# (float* / double* / uint64_t*); body reinterpret-casts to the right type.
|
||||
sig = f"(void* d, {info['c_in']} a, {info['c_in']} b)"
|
||||
|
||||
def _name(d, a, b, rounding, ftz, sat):
|
||||
_, suf = _ptx_arith_modifier_string(dtype, rounding, ftz, sat)
|
||||
return f"tvm_builtin_ptx_{op}_{dtype}{suf}"
|
||||
|
||||
out_c = info["out_cstr"]
|
||||
in_c = info["in_cstr"]
|
||||
out_cast = info["out_cast"]
|
||||
|
||||
def _body(d, a, b, rounding, ftz, sat):
|
||||
mod, _ = _ptx_arith_modifier_string(dtype, rounding, ftz, sat)
|
||||
return (
|
||||
f' asm volatile("{op}{mod}.{dtype} %0, %1, %2;"\n'
|
||||
f' : "={out_c}"(*reinterpret_cast<{out_cast}>(d))\n'
|
||||
f' : "{in_c}"(a), "{in_c}"(b));'
|
||||
)
|
||||
|
||||
return _name, sig, _body
|
||||
|
||||
|
||||
def _ptx_fma_parts(dtype):
|
||||
"""Return (name_fn, sig, body_fn) for ptx_fma_{dtype}."""
|
||||
info = _DTYPE_INFO[dtype]
|
||||
sig = f"(void* d, {info['c_in']} a, {info['c_in']} b, {info['c_in']} c)"
|
||||
|
||||
def _name(d, a, b, c, rounding, ftz, sat):
|
||||
_, suf = _ptx_arith_modifier_string(dtype, rounding, ftz, sat)
|
||||
return f"tvm_builtin_ptx_fma_{dtype}{suf}"
|
||||
|
||||
out_c = info["out_cstr"]
|
||||
in_c = info["in_cstr"]
|
||||
out_cast = info["out_cast"]
|
||||
|
||||
def _body(d, a, b, c, rounding, ftz, sat):
|
||||
mod, _ = _ptx_arith_modifier_string(dtype, rounding, ftz, sat)
|
||||
return (
|
||||
f' asm volatile("fma{mod}.{dtype} %0, %1, %2, %3;"\n'
|
||||
f' : "={out_c}"(*reinterpret_cast<{out_cast}>(d))\n'
|
||||
f' : "{in_c}"(a), "{in_c}"(b), "{in_c}"(c));'
|
||||
)
|
||||
|
||||
return _name, sig, _body
|
||||
|
||||
|
||||
# Register 12 ops: {add, sub, mul, fma} x {f32, f32x2, f64}.
|
||||
for _dtype in ("f32", "f32x2", "f64"):
|
||||
for _op in ("add", "sub", "mul"):
|
||||
_name_fn, _sig, _body_fn = _ptx_binary_arith_parts(_op, _dtype)
|
||||
device_intrinsic(
|
||||
f"ptx_{_op}_{_dtype}",
|
||||
n_attrs=3, # rounding, ftz, sat
|
||||
helper_name=_name_fn,
|
||||
c_signature=_sig,
|
||||
body=_body_fn,
|
||||
)
|
||||
_name_fn, _sig, _body_fn = _ptx_fma_parts(_dtype)
|
||||
device_intrinsic(
|
||||
f"ptx_fma_{_dtype}",
|
||||
n_attrs=3,
|
||||
helper_name=_name_fn,
|
||||
c_signature=_sig,
|
||||
body=_body_fn,
|
||||
)
|
||||
del _dtype, _op, _name_fn, _sig, _body_fn
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# ex2.approx.ftz.f32 / rcp.approx.ftz.f32 — 1 form each.
|
||||
# =============================================================================
|
||||
device_intrinsic(
|
||||
"ptx_exp2",
|
||||
c_signature="(float x)",
|
||||
return_type="float",
|
||||
body=(
|
||||
" float result;\n"
|
||||
' asm volatile("ex2.approx.ftz.f32 %0, %1;" : "=f"(result) : "f"(x));\n'
|
||||
" return result;"
|
||||
),
|
||||
)
|
||||
device_intrinsic(
|
||||
"ptx_rcp",
|
||||
c_signature="(float x)",
|
||||
return_type="float",
|
||||
body=(
|
||||
" float result;\n"
|
||||
' asm volatile("rcp.approx.ftz.f32 %0, %1;" : "=f"(result) : "f"(x));\n'
|
||||
" return result;"
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# 3-operand max.f32 / min.f32 — the f32, 3-operand form-table entry of the
|
||||
# redux/reduction-style fp32 max/min ops.
|
||||
# =============================================================================
|
||||
_ABC_SIG = "(float a, float b, float c)"
|
||||
device_intrinsic(
|
||||
"ptx_reduce3_max_f32",
|
||||
c_signature=_ABC_SIG,
|
||||
return_type="float",
|
||||
body=(
|
||||
" float result;\n"
|
||||
' asm volatile("max.f32 %0, %1, %2, %3;"\n'
|
||||
' : "=f"(result) : "f"(a), "f"(b), "f"(c));\n'
|
||||
" return result;"
|
||||
),
|
||||
)
|
||||
device_intrinsic(
|
||||
"ptx_reduce3_min_f32",
|
||||
c_signature=_ABC_SIG,
|
||||
return_type="float",
|
||||
body=(
|
||||
" float result;\n"
|
||||
' asm volatile("min.f32 %0, %1, %2, %3;"\n'
|
||||
' : "=f"(result) : "f"(a), "f"(b), "f"(c));\n'
|
||||
" return result;"
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
_BINARY_F32_SIG = "(float a, float b)"
|
||||
|
||||
|
||||
def _ptx_max_f32_body(a, b, ftz, nan):
|
||||
ftz_b = bool(int(ftz)) if hasattr(ftz, "value") else bool(ftz)
|
||||
nan_b = bool(int(nan)) if hasattr(nan, "value") else bool(nan)
|
||||
ftz_suffix = ".ftz" if ftz_b else ""
|
||||
nan_suffix = ".NaN" if nan_b else ""
|
||||
return (
|
||||
" float result;\n"
|
||||
f' asm volatile("max{ftz_suffix}{nan_suffix}.f32 %0, %1, %2;"\n'
|
||||
' : "=f"(result) : "f"(a), "f"(b));\n'
|
||||
" return result;"
|
||||
)
|
||||
|
||||
|
||||
def _ptx_max_f32_name(a, b, ftz, nan):
|
||||
ftz_b = bool(int(ftz)) if hasattr(ftz, "value") else bool(ftz)
|
||||
nan_b = bool(int(nan)) if hasattr(nan, "value") else bool(nan)
|
||||
suffix = ""
|
||||
if ftz_b:
|
||||
suffix += "_ftz"
|
||||
if nan_b:
|
||||
suffix += "_nan"
|
||||
return f"tvm_builtin_ptx_max_f32{suffix}"
|
||||
|
||||
|
||||
device_intrinsic(
|
||||
"ptx_max_f32",
|
||||
n_attrs=2,
|
||||
helper_name=_ptx_max_f32_name,
|
||||
c_signature=_BINARY_F32_SIG,
|
||||
return_type="float",
|
||||
body=_ptx_max_f32_body,
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# CUDA-side warp / CTA reductions (templated butterfly shuffle-XOR).
|
||||
# Emitted directly via ``cuda_func_call`` — the helper signature uses a
|
||||
# single template parameter ``T`` for both arg and return, which doesn't
|
||||
# match the operand-driven C signature pattern.
|
||||
# =============================================================================
|
||||
|
||||
# (accumulation expression, identity value for cross-warp padding)
|
||||
_OP_TABLE = {
|
||||
"sum": ("val += shuffled;", "T(0)"),
|
||||
"max": ("val = max(val, shuffled);", "-INFINITY"),
|
||||
"min": ("val = min(val, shuffled);", "INFINITY"),
|
||||
}
|
||||
|
||||
|
||||
def _validate_op(op_str, context):
|
||||
if op_str not in _OP_TABLE:
|
||||
raise ValueError(f"Unsupported {context} op '{op_str}', expected one of {list(_OP_TABLE)}")
|
||||
return _OP_TABLE[op_str]
|
||||
|
||||
|
||||
def _warp_reduce_source(func_name, width_int, step_expr):
|
||||
return (
|
||||
f"\ntemplate <typename T>\n"
|
||||
f"__forceinline__ __device__ T {func_name}(T val) {{\n"
|
||||
f" #pragma unroll\n"
|
||||
f" for (int mask = {width_int} >> 1; mask > 0; mask >>= 1) {{\n"
|
||||
" T shuffled = __shfl_xor_sync(0xFFFFFFFF, val, mask);\n"
|
||||
f" {step_expr}\n"
|
||||
" }\n"
|
||||
" return val;\n"
|
||||
"}\n"
|
||||
)
|
||||
|
||||
|
||||
@register_codegen("cuda_warp_reduce")
|
||||
def codegen_cuda_warp_reduce(value, op, width):
|
||||
op_str = parse_str(op)
|
||||
width_int = validate_power_of_two_range(width, 2, 32, "warp_reduce width")
|
||||
step_expr, _ = _validate_op(op_str, "warp_reduce")
|
||||
|
||||
func_name = f"tvm_builtin_cuda_warp_reduce_{op_str}_{width_int}"
|
||||
source_code = _warp_reduce_source(func_name, width_int, step_expr)
|
||||
return cuda_func_call(func_name, value, source_code=source_code, return_type=value.ty)
|
||||
|
||||
|
||||
@register_codegen("cuda_cta_reduce")
|
||||
def codegen_cuda_cta_reduce(value, op, num_warps, scratch):
|
||||
op_str = parse_str(op)
|
||||
nw = validate_power_of_two_range(num_warps, 1, 32, "cta_reduce num_warps")
|
||||
step_expr, identity = _validate_op(op_str, "cta_reduce")
|
||||
|
||||
warp_reduce_name = f"tvm_builtin_cuda_warp_reduce_{op_str}_32"
|
||||
func_name = f"tvm_builtin_cuda_cta_reduce_{op_str}_{nw}"
|
||||
|
||||
cta_body = (
|
||||
f"{_warp_reduce_source(warp_reduce_name, 32, step_expr)}"
|
||||
"template <typename T>\n"
|
||||
f"__forceinline__ __device__ T {func_name}(T val, void* scratch_raw) {{\n"
|
||||
" T* scratch = reinterpret_cast<T*>(scratch_raw);\n"
|
||||
f" val = {warp_reduce_name}(val);\n"
|
||||
" int tid = threadIdx.x + threadIdx.y * blockDim.x"
|
||||
" + threadIdx.z * blockDim.x * blockDim.y;\n"
|
||||
" int warp_id = tid / 32;\n"
|
||||
" int lane_id = tid % 32;\n"
|
||||
" if (lane_id == 0) scratch[warp_id] = val;\n"
|
||||
" __syncthreads();\n"
|
||||
" if (warp_id == 0) {\n"
|
||||
f" T partial = (lane_id < {nw}) ? scratch[lane_id] : {identity};\n"
|
||||
f" partial = {warp_reduce_name}(partial);\n"
|
||||
" if (lane_id == 0) scratch[0] = partial;\n"
|
||||
" }\n"
|
||||
" __syncthreads();\n"
|
||||
" return scratch[0];\n"
|
||||
"}\n"
|
||||
)
|
||||
return cuda_func_call(func_name, value, scratch, source_code=cta_body, return_type=value.ty)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Additional FP8/BF16 packing, integer, and activation helpers.
|
||||
# =============================================================================
|
||||
|
||||
# PTX integer bit-search form:
|
||||
# fns.b32 d, mask, base, offset;
|
||||
device_intrinsic(
|
||||
"ptx_fns_b32",
|
||||
helper_name="tvm_builtin_ptx_fns_b32",
|
||||
c_signature="(unsigned int mask, unsigned int base, int offset)",
|
||||
return_type="unsigned int",
|
||||
body=(
|
||||
" unsigned int ret;\n"
|
||||
' asm("fns.b32 %0, %1, %2, %3;" : "=r"(ret) : "r"(mask), "r"(base), "r"(offset));\n'
|
||||
" return ret;"
|
||||
),
|
||||
)
|
||||
|
||||
device_intrinsic(
|
||||
"cuda_ffs_u32",
|
||||
helper_name="tvm_builtin_ffs_u32",
|
||||
c_signature="(unsigned int value)",
|
||||
return_type="int",
|
||||
body=" return __ffs(value);",
|
||||
)
|
||||
|
||||
device_intrinsic(
|
||||
"ptx_add_rn_f32_bf16",
|
||||
helper_name="tvm_builtin_ptx_add_rn_f32_bf16",
|
||||
c_signature="(float acc, unsigned short x)",
|
||||
return_type="float",
|
||||
body=(' asm("add.rn.f32.bf16 %0, %1, %0;" : "+f"(acc) : "h"(x));\n return acc;'),
|
||||
)
|
||||
|
||||
|
||||
device_intrinsic(
|
||||
"cuda_make_float2",
|
||||
helper_name="tvm_builtin_make_float2",
|
||||
c_signature="(float x, float y)",
|
||||
return_type="unsigned long long",
|
||||
body=(
|
||||
" float2 value = make_float2(x, y);\n"
|
||||
" return *reinterpret_cast<unsigned long long*>(&value);"
|
||||
),
|
||||
)
|
||||
|
||||
device_intrinsic(
|
||||
"cuda_float2_x",
|
||||
helper_name="tvm_builtin_float2_x",
|
||||
c_signature="(unsigned long long packed)",
|
||||
return_type="float",
|
||||
body=(" float2 value = *reinterpret_cast<float2*>(&packed);\n return value.x;"),
|
||||
)
|
||||
|
||||
device_intrinsic(
|
||||
"cuda_float2_y",
|
||||
helper_name="tvm_builtin_float2_y",
|
||||
c_signature="(unsigned long long packed)",
|
||||
return_type="float",
|
||||
body=(" float2 value = *reinterpret_cast<float2*>(&packed);\n return value.y;"),
|
||||
)
|
||||
|
||||
device_intrinsic(
|
||||
"cuda_fmul2_rn",
|
||||
helper_name="tvm_builtin_fmul2_rn",
|
||||
c_signature="(unsigned long long a, unsigned long long b)",
|
||||
return_type="unsigned long long",
|
||||
body=(
|
||||
" float2 lhs = *reinterpret_cast<float2*>(&a);\n"
|
||||
" float2 rhs = *reinterpret_cast<float2*>(&b);\n"
|
||||
" float2 result = __fmul2_rn(lhs, rhs);\n"
|
||||
" return *reinterpret_cast<unsigned long long*>(&result);"
|
||||
),
|
||||
)
|
||||
|
||||
device_intrinsic(
|
||||
"cuda_fadd2_rn",
|
||||
helper_name="tvm_builtin_fadd2_rn",
|
||||
c_signature="(unsigned long long a, unsigned long long b)",
|
||||
return_type="unsigned long long",
|
||||
body=(
|
||||
" float2 lhs = *reinterpret_cast<float2*>(&a);\n"
|
||||
" float2 rhs = *reinterpret_cast<float2*>(&b);\n"
|
||||
" float2 result = __fadd2_rn(lhs, rhs);\n"
|
||||
" return *reinterpret_cast<unsigned long long*>(&result);"
|
||||
),
|
||||
)
|
||||
|
||||
device_intrinsic(
|
||||
"cuda_float22bfloat162_rn",
|
||||
helper_name="tvm_builtin_float22bfloat162_rn",
|
||||
c_signature="(float x, float y)",
|
||||
return_type="unsigned int",
|
||||
body=(
|
||||
" __nv_bfloat162 value = __float22bfloat162_rn(make_float2(x, y));\n"
|
||||
" return *reinterpret_cast<unsigned int*>(&value);"
|
||||
),
|
||||
extra_deps=("bf16",),
|
||||
)
|
||||
|
||||
device_intrinsic(
|
||||
"cuda_float22bfloat162_rn_from_float2",
|
||||
helper_name="tvm_builtin_float22bfloat162_rn_from_float2",
|
||||
c_signature="(unsigned long long packed)",
|
||||
return_type="unsigned int",
|
||||
body=(
|
||||
" float2 value = *reinterpret_cast<float2*>(&packed);\n"
|
||||
" __nv_bfloat162 result = __float22bfloat162_rn(value);\n"
|
||||
" return *reinterpret_cast<unsigned int*>(&result);"
|
||||
),
|
||||
extra_deps=("bf16",),
|
||||
)
|
||||
|
||||
device_intrinsic(
|
||||
"cuda_bfloat1622float2",
|
||||
helper_name="tvm_builtin_bfloat1622float2",
|
||||
c_signature="(unsigned int packed)",
|
||||
return_type="unsigned long long",
|
||||
body=(
|
||||
" __nv_bfloat162 value;\n"
|
||||
" *reinterpret_cast<unsigned int*>(&value) = packed;\n"
|
||||
" float2 result = __bfloat1622float2(value);\n"
|
||||
" return *reinterpret_cast<unsigned long long*>(&result);"
|
||||
),
|
||||
extra_deps=("bf16",),
|
||||
)
|
||||
|
||||
device_intrinsic(
|
||||
"cuda_hmin2",
|
||||
helper_name="tvm_builtin_hmin2",
|
||||
c_signature="(unsigned int a, unsigned int b)",
|
||||
return_type="unsigned int",
|
||||
body=(
|
||||
" __nv_bfloat162 lhs;\n"
|
||||
" __nv_bfloat162 rhs;\n"
|
||||
" *reinterpret_cast<unsigned int*>(&lhs) = a;\n"
|
||||
" *reinterpret_cast<unsigned int*>(&rhs) = b;\n"
|
||||
" __nv_bfloat162 result = __hmin2(lhs, rhs);\n"
|
||||
" return *reinterpret_cast<unsigned int*>(&result);"
|
||||
),
|
||||
extra_deps=("bf16",),
|
||||
)
|
||||
|
||||
device_intrinsic(
|
||||
"cuda_hmax2",
|
||||
helper_name="tvm_builtin_hmax2",
|
||||
c_signature="(unsigned int a, unsigned int b)",
|
||||
return_type="unsigned int",
|
||||
body=(
|
||||
" __nv_bfloat162 lhs;\n"
|
||||
" __nv_bfloat162 rhs;\n"
|
||||
" *reinterpret_cast<unsigned int*>(&lhs) = a;\n"
|
||||
" *reinterpret_cast<unsigned int*>(&rhs) = b;\n"
|
||||
" __nv_bfloat162 result = __hmax2(lhs, rhs);\n"
|
||||
" return *reinterpret_cast<unsigned int*>(&result);"
|
||||
),
|
||||
extra_deps=("bf16",),
|
||||
)
|
||||
|
||||
device_intrinsic(
|
||||
"cuda_fp8x4_e4m3_from_float4",
|
||||
helper_name="tvm_builtin_fp8x4_e4m3_from_float4",
|
||||
c_signature="(float x, float y, float z, float w)",
|
||||
return_type="unsigned int",
|
||||
body=(
|
||||
" __nv_fp8x4_e4m3 result = __nv_fp8x4_e4m3(make_float4(x, y, z, w));\n"
|
||||
" return *reinterpret_cast<unsigned int*>(&result);"
|
||||
),
|
||||
extra_deps=("fp8",),
|
||||
)
|
||||
@@ -0,0 +1,955 @@
|
||||
# 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.
|
||||
# ruff: noqa: E501
|
||||
# pylint: disable=redefined-builtin, invalid-name, too-many-arguments
|
||||
"""Memory ops (load / store / copy / atomic / address conversion / type punning).
|
||||
|
||||
PTX side:
|
||||
* ``ld.acquire.scope{.ss}.type`` scalar load forms.
|
||||
* ``ld.volatile{.ss}.type`` scalar load forms.
|
||||
* Legacy ``ld.global.acquire.gpu`` / ``ld.global.cg`` result-argument helper.
|
||||
* ``mapa.u64`` — map a SMEM ptr to a peer CTA's SMEM in the cluster.
|
||||
|
||||
CUDA side:
|
||||
* ``__ldg`` (cache-as-read-only load).
|
||||
* Templated ``atomicAdd`` / ``atomicCAS``.
|
||||
* half↔float type-punned conversions (single, packed, batch-of-8).
|
||||
* ``__cvta_generic_to_shared`` and ``cluster_addr → shared u32`` casts.
|
||||
"""
|
||||
|
||||
from tvm import DataType
|
||||
from tvm.backend.cuda.op import cuda_func_call
|
||||
|
||||
from ._schema import device_intrinsic
|
||||
from .registry import CODEGEN_REGISTRY, register_codegen
|
||||
from .utils import parse_str
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# __ldg — templated read-only cached load; ``T`` resolved at call time from
|
||||
# the ``dtype`` argument. Hand-written because the helper signature uses a
|
||||
# template parameter for both arg and return.
|
||||
# =============================================================================
|
||||
@register_codegen("cuda_ldg")
|
||||
def codegen_cuda_ldg(addr, dtype):
|
||||
dtype = DataType(parse_str(dtype))
|
||||
func_name = "tvm_builtin_cuda_ldg"
|
||||
source_code = f"""
|
||||
template <typename T>
|
||||
__forceinline__ __device__ T {func_name}(T* src) {{
|
||||
return __ldg(src);
|
||||
}}
|
||||
"""
|
||||
return cuda_func_call(func_name, addr, source_code=source_code, return_type=dtype)
|
||||
|
||||
|
||||
# Shared PTX scalar type metadata (ld/st/red/atom).
|
||||
_PTX_SCALAR_TYPE_INFO = {
|
||||
"b8": ("unsigned int", "r", "uint32"),
|
||||
"u8": ("unsigned int", "r", "uint32"),
|
||||
"s8": ("int", "r", "int32"),
|
||||
"b16": ("unsigned short", "h", "uint16"),
|
||||
"u16": ("unsigned short", "h", "uint16"),
|
||||
"s16": ("short", "h", "int16"),
|
||||
"b32": ("unsigned int", "r", "uint32"),
|
||||
"u32": ("unsigned int", "r", "uint32"),
|
||||
"s32": ("int", "r", "int32"),
|
||||
"b64": ("unsigned long long", "l", "uint64"),
|
||||
"u64": ("unsigned long long", "l", "uint64"),
|
||||
"s64": ("long long", "l", "int64"),
|
||||
"f32": ("float", "f", "float32"),
|
||||
"f64": ("double", "d", "float64"),
|
||||
}
|
||||
_PTX_LD_TYPE_RETURNS = {
|
||||
"b32": {"uint32": "unsigned int", "int32": "int"},
|
||||
"b64": {"uint64": "unsigned long long", "int64": "long long"},
|
||||
}
|
||||
_PTX_VEC_STORE_TYPE = {
|
||||
16: "uint4",
|
||||
8: "uint2",
|
||||
4: "unsigned int",
|
||||
2: "unsigned short",
|
||||
1: "unsigned char",
|
||||
}
|
||||
|
||||
|
||||
def _safe_attr(value):
|
||||
return parse_str(value).replace("::", "_").replace(".", "_")
|
||||
|
||||
|
||||
def _dot(value):
|
||||
value = parse_str(value)
|
||||
return f".{value}" if value else ""
|
||||
|
||||
|
||||
def _cache_suffix(cache):
|
||||
return ".L2::cache_hint" if cache else ""
|
||||
|
||||
|
||||
def _type_info(ptx_type):
|
||||
ptx_type = parse_str(ptx_type)
|
||||
if ptx_type not in _PTX_SCALAR_TYPE_INFO:
|
||||
raise ValueError(
|
||||
f"Unsupported PTX scalar type {ptx_type!r}; expected {sorted(_PTX_SCALAR_TYPE_INFO)}"
|
||||
)
|
||||
return (ptx_type, *_PTX_SCALAR_TYPE_INFO[ptx_type])
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# PTX ld forms (ISA table entries registered via ``ptx_ld`` and siblings):
|
||||
# ld{.weak}{.ss}{.cop}{.level::cache_hint}{.level::prefetch_size}{.vec}.type d, [a]{, cache-policy};
|
||||
# ld{.weak}{.ss}{.level1::eviction_priority}{.level2::eviction_priority}{.level::cache_hint}{.level::prefetch_size}{.vec}.type d, [a]{, cache-policy};
|
||||
# ld.volatile{.ss}{.level::prefetch_size}{.vec}.type d, [a];
|
||||
# ld.relaxed.scope{.ss}{.level1::eviction_priority}{.level2::eviction_priority}{.level::cache_hint}{.level::prefetch_size}{.vec}.type d, [a]{, cache-policy};
|
||||
# ld.acquire.scope{.ss}{.level1::eviction_priority}{.level2::eviction_priority}{.level::cache_hint}{.level::prefetch_size}{.vec}.type d, [a]{, cache-policy};
|
||||
# ld.mmio.sem.sys{.global}.type d, [a];
|
||||
# =============================================================================
|
||||
_PTX_LD_SCOPES = {"cta", "cluster", "gpu", "sys"}
|
||||
_PTX_LD_SPACES = {"global", "shared", "shared::cta", "shared::cluster", "local"}
|
||||
_PTX_LD_VOLATILE_SPACES = _PTX_LD_SPACES | {"const"}
|
||||
_PTX_LD_WEAK_SPACES = _PTX_LD_SPACES | {"const", "param::entry", "param::func"}
|
||||
_PTX_LD_COPS = {"", "ca", "cg", "cs", "lu", "cv"}
|
||||
_PTX_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"}
|
||||
|
||||
|
||||
def _bool_attr(value):
|
||||
return bool(int(value)) if hasattr(value, "value") else bool(value)
|
||||
|
||||
|
||||
def _parse_ld_attrs(return_dtype, ptx_type, scope=None, space="global"):
|
||||
return_dtype = parse_str(return_dtype)
|
||||
ptx_type = parse_str(ptx_type)
|
||||
scope = None if scope is None else parse_str(scope)
|
||||
space = parse_str(space)
|
||||
ptx_type, _ptx, constraint, default_tvm = _type_info(ptx_type)
|
||||
if ptx_type in _PTX_LD_TYPE_RETURNS:
|
||||
returns = _PTX_LD_TYPE_RETURNS[ptx_type]
|
||||
if return_dtype not in returns:
|
||||
raise ValueError(
|
||||
f"PTX ld type {ptx_type!r} cannot return TVM dtype {return_dtype!r}; "
|
||||
f"expected one of {sorted(returns)}"
|
||||
)
|
||||
c_type = returns[return_dtype]
|
||||
else:
|
||||
if return_dtype != default_tvm:
|
||||
raise ValueError(
|
||||
f"PTX ld type {ptx_type!r} cannot return TVM dtype {return_dtype!r}; "
|
||||
f"expected {default_tvm!r}"
|
||||
)
|
||||
c_type = _ptx
|
||||
if scope is not None and scope not in _PTX_LD_SCOPES:
|
||||
raise ValueError(
|
||||
f"Unsupported PTX ld scope {scope!r}; expected one of {sorted(_PTX_LD_SCOPES)}"
|
||||
)
|
||||
return return_dtype, ptx_type, scope, space, c_type, constraint
|
||||
|
||||
|
||||
def _validate_ld_space(space: str, allowed: set[str]) -> None:
|
||||
if space not in allowed:
|
||||
raise ValueError(
|
||||
f"Unsupported PTX ld state space {space!r}; expected one of {sorted(allowed)}"
|
||||
)
|
||||
|
||||
|
||||
def _ptx_level_suffix(has_cache, l1_evict, l2_evict, prefetch_size):
|
||||
suffix = _cache_suffix("cache" if has_cache else "")
|
||||
l1_evict = parse_str(l1_evict)
|
||||
l2_evict = parse_str(l2_evict)
|
||||
prefetch_size = parse_str(prefetch_size)
|
||||
if l1_evict:
|
||||
suffix += f".{l1_evict}"
|
||||
if l2_evict:
|
||||
suffix += f".{l2_evict}"
|
||||
if prefetch_size:
|
||||
suffix += f".{prefetch_size}"
|
||||
return suffix
|
||||
|
||||
|
||||
def _ptx_shared_addr(space, ptr_name="address"):
|
||||
if parse_str(space).startswith("shared"):
|
||||
return (
|
||||
f" unsigned int addr = (unsigned int)__cvta_generic_to_shared({ptr_name});\n",
|
||||
'"r"(addr)',
|
||||
)
|
||||
return "", f'"l"({ptr_name})'
|
||||
|
||||
|
||||
def _ptx_ld_vec_store(num_bytes, vec_len, ptx_type):
|
||||
if ptx_type == "u8" and vec_len == 1:
|
||||
return " *reinterpret_cast<unsigned char*>(dst_ptr) = static_cast<unsigned char>(r0);"
|
||||
store_type = _PTX_VEC_STORE_TYPE[num_bytes]
|
||||
if vec_len > 1:
|
||||
return (
|
||||
f" *reinterpret_cast<{store_type}*>(dst_ptr) = "
|
||||
+ "{"
|
||||
+ ", ".join(f"r{i}" for i in range(vec_len))
|
||||
+ "};"
|
||||
)
|
||||
return f" *reinterpret_cast<{store_type}*>(dst_ptr) = r0;"
|
||||
|
||||
|
||||
def _ptx_ld_form_parts(form, attr_args):
|
||||
if form == "weak":
|
||||
(
|
||||
return_dtype,
|
||||
weak,
|
||||
space,
|
||||
cop,
|
||||
vec,
|
||||
ptx_type,
|
||||
has_cache_hint,
|
||||
to_dst,
|
||||
l1_evict,
|
||||
l2_evict,
|
||||
prefetch_size,
|
||||
) = attr_args
|
||||
sem, scope = "", ""
|
||||
elif form == "relaxed":
|
||||
(
|
||||
return_dtype,
|
||||
scope,
|
||||
space,
|
||||
vec,
|
||||
ptx_type,
|
||||
has_cache_hint,
|
||||
to_dst,
|
||||
l1_evict,
|
||||
l2_evict,
|
||||
prefetch_size,
|
||||
) = attr_args
|
||||
sem, weak, cop = "", False, ""
|
||||
elif form == "acquire":
|
||||
(
|
||||
return_dtype,
|
||||
scope,
|
||||
space,
|
||||
vec,
|
||||
ptx_type,
|
||||
has_cache_hint,
|
||||
to_dst,
|
||||
l1_evict,
|
||||
l2_evict,
|
||||
prefetch_size,
|
||||
) = attr_args
|
||||
sem, weak, cop = "", False, ""
|
||||
elif form == "volatile":
|
||||
return_dtype, space, vec, ptx_type, to_dst, prefetch_size = attr_args
|
||||
sem, scope, weak, cop = "", "", False, ""
|
||||
has_cache_hint, l1_evict, l2_evict = False, "", ""
|
||||
elif form == "mmio":
|
||||
return_dtype, sem, scope, space, ptx_type, to_dst = attr_args
|
||||
weak, cop, vec = False, "", ""
|
||||
has_cache_hint, l1_evict, l2_evict, prefetch_size = False, "", "", ""
|
||||
else:
|
||||
raise ValueError(f"unknown ld form {form!r}")
|
||||
|
||||
return_dtype, ptx_type, scope, space, c_type, constraint = _parse_ld_attrs(
|
||||
return_dtype, ptx_type, scope if form in ("relaxed", "acquire") else None, space
|
||||
)
|
||||
sem = parse_str(sem)
|
||||
scope = parse_str(scope)
|
||||
space = parse_str(space)
|
||||
cop = parse_str(cop)
|
||||
vec = parse_str(vec)
|
||||
l1_evict = parse_str(l1_evict)
|
||||
l2_evict = parse_str(l2_evict)
|
||||
prefetch_size = parse_str(prefetch_size)
|
||||
weak = _bool_attr(weak)
|
||||
has_cache = _bool_attr(has_cache_hint)
|
||||
to_dst = _bool_attr(to_dst)
|
||||
if cop and cop not in _PTX_LD_COPS:
|
||||
raise ValueError(f"Unsupported PTX ld cache operation {cop!r}")
|
||||
if vec and vec not in _PTX_VEC:
|
||||
raise ValueError(f"Unsupported PTX ld vector modifier {vec!r}")
|
||||
if l1_evict and l1_evict not in _PTX_L1_EVICT:
|
||||
raise ValueError(f"Unsupported PTX ld L1 eviction {l1_evict!r}")
|
||||
if l2_evict and l2_evict not in _PTX_L2_EVICT:
|
||||
raise ValueError(f"Unsupported PTX ld L2 eviction {l2_evict!r}")
|
||||
if prefetch_size and prefetch_size not in _PTX_PREFETCH:
|
||||
raise ValueError(f"Unsupported PTX ld prefetch size {prefetch_size!r}")
|
||||
if form == "mmio":
|
||||
if sem not in ("acquire", "relaxed") or scope != "sys" or space != "global":
|
||||
raise ValueError("ld.mmio requires sem in {acquire, relaxed}, scope=sys, space=global")
|
||||
prefix = f"ld.mmio.{sem}.{scope}"
|
||||
elif form == "relaxed":
|
||||
if not scope:
|
||||
raise ValueError("ld.relaxed requires scope")
|
||||
_validate_ld_space(space, _PTX_LD_SPACES)
|
||||
prefix = f"ld.relaxed.{scope}{_dot(space)}"
|
||||
elif form == "acquire":
|
||||
if not scope:
|
||||
raise ValueError("ld.acquire requires scope")
|
||||
_validate_ld_space(space, _PTX_LD_SPACES)
|
||||
prefix = f"ld.acquire.{scope}{_dot(space)}"
|
||||
elif form == "volatile":
|
||||
_validate_ld_space(space, _PTX_LD_VOLATILE_SPACES)
|
||||
prefix = f"ld.volatile{_dot(space)}"
|
||||
else:
|
||||
_validate_ld_space(space, _PTX_LD_WEAK_SPACES)
|
||||
prefix = f"ld{'.weak' if weak else ''}{_dot(space)}{_dot(cop)}"
|
||||
level = _ptx_level_suffix(has_cache, l1_evict, l2_evict, prefetch_size)
|
||||
vec_len = int(vec[1:]) if vec else 1
|
||||
if vec and not to_dst:
|
||||
raise ValueError("vector ld requires to_dst")
|
||||
elem_bytes = (
|
||||
8
|
||||
if ptx_type.endswith("64")
|
||||
else 2
|
||||
if ptx_type in ("u16", "s16", "b16")
|
||||
else 1
|
||||
if ptx_type in ("u8", "s8", "b8")
|
||||
else 4
|
||||
)
|
||||
num_bytes = vec_len * elem_bytes if vec else elem_bytes
|
||||
name_parts = [
|
||||
"tvm_builtin_ptx_ld",
|
||||
form if form != "weak" else ("weak" if weak else "plain"),
|
||||
]
|
||||
if sem:
|
||||
name_parts.append(_safe_attr(sem))
|
||||
if scope:
|
||||
name_parts.append(_safe_attr(scope))
|
||||
name_parts.extend(
|
||||
[
|
||||
_safe_attr(space),
|
||||
_safe_attr(cop) if cop else "",
|
||||
_safe_attr(vec) if vec else "",
|
||||
ptx_type,
|
||||
return_dtype if not to_dst else "to_dst",
|
||||
]
|
||||
)
|
||||
if has_cache:
|
||||
name_parts.append("cache_hint")
|
||||
if l1_evict:
|
||||
name_parts.append(_safe_attr(l1_evict))
|
||||
if l2_evict:
|
||||
name_parts.append(_safe_attr(l2_evict))
|
||||
if prefetch_size:
|
||||
name_parts.append(_safe_attr(prefetch_size))
|
||||
name = "_".join(p for p in name_parts if p)
|
||||
cache_operand = ', "l"(cache_policy)' if has_cache else ""
|
||||
addr_decl, addr_operand = _ptx_shared_addr(space, "src_ptr" if to_dst else "address")
|
||||
if to_dst:
|
||||
reg_decls = "".join(f" {c_type} r{i};\n" for i in range(vec_len))
|
||||
if vec_len > 1:
|
||||
out_slot = "{" + ", ".join(f"%{i}" for i in range(vec_len)) + "}"
|
||||
out_constraints = ", ".join(f'"={constraint}"(r{i})' for i in range(vec_len))
|
||||
addr_idx = vec_len
|
||||
else:
|
||||
out_slot = "%0"
|
||||
out_constraints = f'"={constraint}"(r0)'
|
||||
addr_idx = 1
|
||||
cache_slot = f", %{addr_idx + 1}" if has_cache else ""
|
||||
instr = f"{prefix}{level}{_dot(vec)}.{ptx_type}"
|
||||
body = (
|
||||
f"{addr_decl}{reg_decls}"
|
||||
f' asm volatile("{instr} {out_slot}, [%{addr_idx}]{cache_slot};"\n'
|
||||
f" : {out_constraints}\n"
|
||||
f" : {addr_operand}{cache_operand});\n"
|
||||
f"{_ptx_ld_vec_store(num_bytes, vec_len, ptx_type)}"
|
||||
)
|
||||
return (
|
||||
name,
|
||||
"(void* dst_ptr, void* src_ptr, unsigned long long cache_policy)",
|
||||
"void",
|
||||
"",
|
||||
body,
|
||||
)
|
||||
cache_slot = ", %2" if has_cache else ""
|
||||
instr = f"{prefix}{level}{_dot(vec)}.{ptx_type}"
|
||||
body = (
|
||||
f" {c_type} ret;\n"
|
||||
f"{addr_decl}"
|
||||
f' asm volatile("{instr} %0, [%1]{cache_slot};"\n'
|
||||
f' : "={constraint}"(ret)\n'
|
||||
f" : {addr_operand}{cache_operand});\n"
|
||||
" return ret;"
|
||||
)
|
||||
sig = (
|
||||
"(void* address, unsigned long long cache_policy)" if form == "weak" else "(void* address)"
|
||||
)
|
||||
return name, sig, c_type, return_dtype, body
|
||||
|
||||
|
||||
def _register_ptx_ld(op_name, form, n_attrs):
|
||||
def _parts(*args):
|
||||
return _ptx_ld_form_parts(form, args[-n_attrs:])
|
||||
|
||||
device_intrinsic(
|
||||
op_name,
|
||||
n_attrs=n_attrs,
|
||||
helper_name=lambda *a, _p=_parts: _p(*a)[0],
|
||||
c_signature=lambda *a, _p=_parts: _p(*a)[1],
|
||||
return_type=lambda *a, _p=_parts: _p(*a)[2],
|
||||
tvm_return_type=lambda *a, _p=_parts: (
|
||||
None if _p(*a)[2] == "void" else parse_str(a[-n_attrs])
|
||||
),
|
||||
body=lambda *a, _p=_parts: _p(*a)[4],
|
||||
)
|
||||
|
||||
|
||||
_register_ptx_ld("ptx_ld", "weak", 11)
|
||||
_register_ptx_ld("ptx_ld_relaxed", "relaxed", 10)
|
||||
_register_ptx_ld("ptx_ld_acquire", "acquire", 10)
|
||||
_register_ptx_ld("ptx_ld_volatile", "volatile", 6)
|
||||
_register_ptx_ld("ptx_ld_mmio", "mmio", 6)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Legacy acquire-load lvalue API — compatibility wrapper over
|
||||
# ``ld.acquire.gpu.global`` / ``ld.global.cg`` forms, dispatched on dtype.
|
||||
# Wrapper picks .b32/.b64 + matching constraint by dtype.
|
||||
#
|
||||
# The body uses ``#if __CUDA_ARCH__ >= 700`` to select acquire on SM70+ and
|
||||
# fall back to .cg on older arches. This is two PTX form table entries
|
||||
# combined in one device helper for arch portability.
|
||||
# =============================================================================
|
||||
_LD_GLOBAL_ACQUIRE_DTYPES = {
|
||||
"uint32": ("uint32_t", "b32", "r"),
|
||||
"int32": ("int32_t", "b32", "r"),
|
||||
"uint64": ("uint64_t", "b64", "l"),
|
||||
"int64": ("int64_t", "b64", "l"),
|
||||
}
|
||||
|
||||
|
||||
def _ld_global_acquire_body(ptx_type: str, spec: str) -> str:
|
||||
return (
|
||||
" #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700\n"
|
||||
f' asm volatile ("ld.acquire.gpu.global.{ptx_type} %0, [%1];\\n"\n'
|
||||
f' : "={spec}"(res) : "l"(addr));\n'
|
||||
" #else\n"
|
||||
f' asm volatile ("ld.global.cg.{ptx_type} %0, [%1];\\n"\n'
|
||||
f' : "={spec}"(res) : "l"(addr));\n'
|
||||
" #endif"
|
||||
)
|
||||
|
||||
|
||||
for _dtype, (_c_type, _ptx_type, _spec) in _LD_GLOBAL_ACQUIRE_DTYPES.items():
|
||||
device_intrinsic(
|
||||
f"ptx_ld_global_acquire_{_dtype}",
|
||||
c_signature=f"({_c_type}& res, {_c_type}* addr)",
|
||||
body=_ld_global_acquire_body(_ptx_type, _spec),
|
||||
)
|
||||
del _dtype, _c_type, _ptx_type, _spec
|
||||
|
||||
|
||||
@register_codegen("ptx_ld_global_acquire")
|
||||
def codegen_ptx_ld_global_acquire(res, addr):
|
||||
"""Dispatch to the dtype-specific helper."""
|
||||
dtype = str(res.ty)
|
||||
if dtype not in _LD_GLOBAL_ACQUIRE_DTYPES:
|
||||
raise ValueError(f"Unsupported data type for ld.global.acquire: {dtype}")
|
||||
result = CODEGEN_REGISTRY[f"tirx.ptx_ld_global_acquire_{dtype}"]([res, addr])
|
||||
return result[0] if isinstance(result, tuple) else result
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Atomics — templated wrappers around CUDA's ``atomicAdd`` / ``atomicCAS``.
|
||||
# =============================================================================
|
||||
device_intrinsic(
|
||||
"cuda_atomic_add",
|
||||
helper_name="tvm_builtin_cuda_atomic_add",
|
||||
c_signature="(T* addr, T value)",
|
||||
body=" return atomicAdd(addr, value);",
|
||||
return_type="T",
|
||||
templated=True,
|
||||
tvm_return_type=lambda _addr, value: value.ty,
|
||||
)
|
||||
device_intrinsic(
|
||||
"cuda_atomic_cas",
|
||||
helper_name="tvm_builtin_cuda_atomic_cas",
|
||||
c_signature="(T* address, T compare, T val)",
|
||||
body=" return atomicCAS(address, compare, val);",
|
||||
return_type="T",
|
||||
templated=True,
|
||||
tvm_return_type=lambda _p, old, _n: old.ty,
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# half / bfloat16 ↔ float type-punned conversions.
|
||||
# =============================================================================
|
||||
device_intrinsic(
|
||||
"cuda_half2float",
|
||||
c_signature="(half src)",
|
||||
body=" return __half2float(src);",
|
||||
return_type="float",
|
||||
tvm_return_type="float32",
|
||||
)
|
||||
device_intrinsic(
|
||||
"cuda_bfloat162float",
|
||||
c_signature="(nv_bfloat16 src)",
|
||||
body=" return __bfloat162float(src);",
|
||||
return_type="float",
|
||||
tvm_return_type="float32",
|
||||
)
|
||||
device_intrinsic(
|
||||
"cuda_float22half2",
|
||||
c_signature="(void* dst, void* src)",
|
||||
body=(
|
||||
" half2* dst_p = (half2*) dst;\n"
|
||||
" float2* src_p = (float2*) src;\n"
|
||||
" *dst_p = __float22half2_rn(*src_p);"
|
||||
),
|
||||
)
|
||||
device_intrinsic(
|
||||
"cuda_half8tofloat8",
|
||||
c_signature="(void* src_addr, void* dst_addr)",
|
||||
body=(
|
||||
" half2* source = (half2*) src_addr;\n"
|
||||
" float2* dest = (float2*) dst_addr;\n"
|
||||
" for (int i = 0; i < 4; i++) {\n"
|
||||
" dest[i] = __half22float2(source[i]);\n"
|
||||
" }"
|
||||
),
|
||||
)
|
||||
device_intrinsic(
|
||||
"cuda_float8tohalf8",
|
||||
c_signature="(void* src_addr, void* dst_addr)",
|
||||
body=(
|
||||
" float2* source = (float2*) src_addr;\n"
|
||||
" half2* dest = (half2*) dst_addr;\n"
|
||||
" for (int i = 0; i < 4; i++) {\n"
|
||||
" dest[i] = __float22half2_rn(source[i]);\n"
|
||||
" }"
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Address-conversion helpers used by op-wrapper-side dispatch in tvm.tirx.op.
|
||||
# Each precomputes a value that the schema's specialized op then takes as a
|
||||
# typed scalar input (instead of doing the conversion inside the asm helper).
|
||||
# =============================================================================
|
||||
device_intrinsic(
|
||||
"cuda_cvta_generic_to_shared",
|
||||
c_signature="(void* p)",
|
||||
body=" return __cvta_generic_to_shared(p);",
|
||||
return_type="unsigned int",
|
||||
tvm_return_type="uint32",
|
||||
)
|
||||
|
||||
device_intrinsic(
|
||||
"cuda_smem_addr_from_uint64",
|
||||
c_signature="(uint64_t cluster_addr)",
|
||||
body=" return static_cast<unsigned int>(cluster_addr);",
|
||||
return_type="unsigned int",
|
||||
tvm_return_type="uint32",
|
||||
)
|
||||
|
||||
# =============================================================================
|
||||
# PTX mapa form:
|
||||
# mapa{.space}.type d, a, b;
|
||||
# .space = {.shared::cluster}; .type = {.u32, .u64}
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def _ptx_mapa_parts(_addr, _rank, space, ptx_type, return_dtype):
|
||||
space = parse_str(space)
|
||||
ptx_type = parse_str(ptx_type)
|
||||
return_dtype = parse_str(return_dtype)
|
||||
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}")
|
||||
c_type = "uint32_t" if ptx_type == "u32" else "uint64_t"
|
||||
constraint = "r" if ptx_type == "u32" else "l"
|
||||
name = f"tvm_builtin_ptx_mapa{('_' + _safe_attr(space)) if space else ''}_{ptx_type}"
|
||||
body = (
|
||||
f" {c_type} result;\n"
|
||||
f' asm volatile("mapa{_dot(space)}.{ptx_type} %0, %1, %2;"\n'
|
||||
f' : "={constraint}"(result) : "l"(addr), "r"(rank));\n'
|
||||
" return result;"
|
||||
)
|
||||
return name, c_type, return_dtype, body
|
||||
|
||||
|
||||
device_intrinsic(
|
||||
"ptx_mapa",
|
||||
n_attrs=3,
|
||||
helper_name=lambda *a: _ptx_mapa_parts(*a)[0],
|
||||
c_signature="(void* addr, uint32_t rank)",
|
||||
return_type=lambda *a: _ptx_mapa_parts(*a)[1],
|
||||
tvm_return_type=lambda *a: _ptx_mapa_parts(*a)[2],
|
||||
body=lambda *a: _ptx_mapa_parts(*a)[3],
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Generic PTX memory forms. Compatibility wrappers in ``tvm.tirx.op`` bind
|
||||
# concrete sem/scope/space/op/type parameters for existing call sites.
|
||||
# =============================================================================
|
||||
|
||||
|
||||
# PTX red scalar form:
|
||||
# red{.sem}{.scope}{.space}.op{.level::cache_hint}.type [a], b{, cache-policy};
|
||||
def _ptx_red_scalar_parts(*args):
|
||||
sem, scope, space, op, ptx_type, has_cache_hint = args[-6:]
|
||||
sem = parse_str(sem)
|
||||
scope = parse_str(scope)
|
||||
space = parse_str(space)
|
||||
op = parse_str(op)
|
||||
ptx_type, c_type, constraint, _tvm_dtype = _type_info(ptx_type)
|
||||
has_cache = (
|
||||
bool(int(has_cache_hint)) if hasattr(has_cache_hint, "value") else bool(has_cache_hint)
|
||||
)
|
||||
modifiers = f"{_dot(sem)}{_dot(scope)}{_dot(space)}"
|
||||
instr = f"red{modifiers}.{op}{_cache_suffix('cache' if has_cache else '')}.{ptx_type}"
|
||||
name = (
|
||||
"tvm_builtin_ptx_red_scalar"
|
||||
f"{_dot(sem).replace('.', '_')}{_dot(scope).replace('.', '_')}"
|
||||
f"_{_safe_attr(space)}_{op}_{ptx_type}{'_cache_hint' if has_cache else ''}"
|
||||
)
|
||||
cache_operand = ', "l"(cache_policy)' if has_cache else ""
|
||||
addr_decl = ""
|
||||
addr_operand = '"l"(address)'
|
||||
if space.startswith("shared"):
|
||||
addr_decl = " unsigned int addr = (unsigned int)__cvta_generic_to_shared(address);\n"
|
||||
addr_operand = '"r"(addr)'
|
||||
body = (
|
||||
f"{addr_decl}"
|
||||
f' asm volatile("{instr} [%0], %1{", %2" if has_cache else ""};"\n'
|
||||
" :\n"
|
||||
f' : {addr_operand}, "{constraint}"(value)'
|
||||
f"{cache_operand}\n"
|
||||
' : "memory");'
|
||||
)
|
||||
return name, f"(void* address, {c_type} value, unsigned long long cache_policy)", body
|
||||
|
||||
|
||||
device_intrinsic(
|
||||
"ptx_red_scalar",
|
||||
n_attrs=6,
|
||||
helper_name=lambda *a: _ptx_red_scalar_parts(*a)[0],
|
||||
c_signature=lambda *a: _ptx_red_scalar_parts(*a)[1],
|
||||
body=lambda *a: _ptx_red_scalar_parts(*a)[2],
|
||||
)
|
||||
|
||||
|
||||
# PTX atom scalar one-source-operand form:
|
||||
# atom{.sem}{.scope}{.space}.op{.level::cache_hint}.type d, [a], b{, cache-policy};
|
||||
def _ptx_atom_scalar_parts(*args):
|
||||
sem, scope, space, op, ptx_type, has_cache_hint = args[-6:]
|
||||
sem = parse_str(sem)
|
||||
scope = parse_str(scope)
|
||||
space = parse_str(space)
|
||||
op = parse_str(op)
|
||||
ptx_type, c_type, constraint, tvm_dtype = _type_info(ptx_type)
|
||||
has_cache = (
|
||||
bool(int(has_cache_hint)) if hasattr(has_cache_hint, "value") else bool(has_cache_hint)
|
||||
)
|
||||
modifiers = f"{_dot(sem)}{_dot(scope)}{_dot(space)}"
|
||||
instr = f"atom{modifiers}.{op}{_cache_suffix('cache' if has_cache else '')}.{ptx_type}"
|
||||
name = (
|
||||
"tvm_builtin_ptx_atom_scalar"
|
||||
f"{_dot(sem).replace('.', '_')}{_dot(scope).replace('.', '_')}"
|
||||
f"_{_safe_attr(space)}_{op}_{ptx_type}{'_cache_hint' if has_cache else ''}"
|
||||
)
|
||||
cache_operand = ', "l"(cache_policy)' if has_cache else ""
|
||||
addr_decl = ""
|
||||
addr_operand = '"l"(address)'
|
||||
if space.startswith("shared"):
|
||||
addr_decl = " unsigned int addr = (unsigned int)__cvta_generic_to_shared(address);\n"
|
||||
addr_operand = '"r"(addr)'
|
||||
body = (
|
||||
f"{addr_decl}"
|
||||
f" {c_type} ret;\n"
|
||||
f' asm volatile("{instr} %0, [%1], %2{", %3" if has_cache else ""};"\n'
|
||||
f' : "={constraint}"(ret)\n'
|
||||
f' : {addr_operand}, "{constraint}"(value)'
|
||||
f"{cache_operand}\n"
|
||||
' : "memory");\n'
|
||||
" return ret;"
|
||||
)
|
||||
return (
|
||||
name,
|
||||
f"(void* address, {c_type} value, unsigned long long cache_policy)",
|
||||
c_type,
|
||||
tvm_dtype,
|
||||
body,
|
||||
)
|
||||
|
||||
|
||||
device_intrinsic(
|
||||
"ptx_atom_scalar",
|
||||
n_attrs=6,
|
||||
helper_name=lambda *a: _ptx_atom_scalar_parts(*a)[0],
|
||||
c_signature=lambda *a: _ptx_atom_scalar_parts(*a)[1],
|
||||
return_type=lambda *a: _ptx_atom_scalar_parts(*a)[2],
|
||||
tvm_return_type=lambda *a: _ptx_atom_scalar_parts(*a)[3],
|
||||
body=lambda *a: _ptx_atom_scalar_parts(*a)[4],
|
||||
)
|
||||
|
||||
|
||||
# PTX prefetch tensormap form:
|
||||
# prefetch{.tensormap_space}.tensormap [a];
|
||||
def _prefetch_tensormap_parts(_tensor_map, tensormap_space):
|
||||
space = parse_str(tensormap_space)
|
||||
instr = f"prefetch{_dot(space)}.tensormap"
|
||||
name = f"tvm_builtin_ptx_prefetch{('_' + _safe_attr(space)) if space else ''}_tensormap"
|
||||
body = (
|
||||
f' asm volatile("{instr} [%0];"\n'
|
||||
" :\n"
|
||||
' : "l"(tensor_map_addr)\n'
|
||||
' : "memory");'
|
||||
)
|
||||
return name, body
|
||||
|
||||
|
||||
device_intrinsic(
|
||||
"ptx_prefetch_tensormap",
|
||||
n_attrs=1,
|
||||
helper_name=lambda *a: _prefetch_tensormap_parts(*a)[0],
|
||||
c_signature="(unsigned long long tensor_map_addr)",
|
||||
body=lambda *a: _prefetch_tensormap_parts(*a)[1],
|
||||
)
|
||||
|
||||
|
||||
# PTX st forms (ISA table entries registered via ``ptx_st`` and siblings):
|
||||
# st{.weak}{.ss}{.cop}{.level::cache_hint}{.vec}.type [a], b{, cache-policy};
|
||||
# st{.weak}{.ss}{.level1::eviction_priority}{.level2::eviction_priority}{.level::cache_hint}{.vec}.type [a], b{, cache-policy};
|
||||
# st.volatile{.ss}{.vec}.type [a], b;
|
||||
# st.relaxed.scope{.ss}{.level1::eviction_priority}{.level2::eviction_priority}{.level::cache_hint}{.vec}.type [a], b{, cache-policy};
|
||||
# st.release.scope{.ss}{.level1::eviction_priority}{.level2::eviction_priority}{.level::cache_hint}{.vec}.type [a], b{, cache-policy};
|
||||
# st.mmio.sem.sys{.global}.type [a], b;
|
||||
_PTX_ST_COPS = {"", "wb", "cg", "cs", "wt"}
|
||||
_PTX_ST_SPACES = {"global", "shared", "shared::cta", "shared::cluster", "local", "param::func"}
|
||||
|
||||
|
||||
def _ptx_st_load_src(num_bytes, vec_len, ptx_type, c_type):
|
||||
if ptx_type == "u8" and vec_len == 1:
|
||||
return " unsigned int r0 = *reinterpret_cast<unsigned char*>(src_ptr);\n"
|
||||
store_type = _PTX_VEC_STORE_TYPE[num_bytes]
|
||||
if vec_len > 1:
|
||||
return f" {store_type} src_ = *reinterpret_cast<{store_type}*>(src_ptr);\n" + "".join(
|
||||
f" {c_type} r{i} = src_.{c};\n" for i, c in enumerate("xyzw"[:vec_len])
|
||||
)
|
||||
return f" {c_type} r0 = *reinterpret_cast<{c_type}*>(src_ptr);\n"
|
||||
|
||||
|
||||
def _ptx_st_form_parts(form, attr_args, from_src):
|
||||
if form == "weak":
|
||||
weak, space, cop, vec, ptx_type, has_cache_hint, l1_evict, l2_evict = attr_args
|
||||
sem, scope = "", ""
|
||||
elif form == "relaxed":
|
||||
scope, space, vec, ptx_type, has_cache_hint, l1_evict, l2_evict = attr_args
|
||||
sem, weak, cop = "relaxed", False, ""
|
||||
elif form == "release":
|
||||
scope, space, vec, ptx_type, has_cache_hint, l1_evict, l2_evict = attr_args
|
||||
sem, weak, cop = "release", False, ""
|
||||
elif form == "volatile":
|
||||
space, vec, ptx_type = attr_args
|
||||
sem, scope, weak, cop = "", "", False, ""
|
||||
has_cache_hint, l1_evict, l2_evict = False, "", ""
|
||||
elif form == "mmio":
|
||||
sem, scope, space, ptx_type = attr_args
|
||||
weak, cop, vec = False, "", ""
|
||||
has_cache_hint, l1_evict, l2_evict = False, "", ""
|
||||
else:
|
||||
raise ValueError(f"unknown st form {form!r}")
|
||||
|
||||
sem = parse_str(sem)
|
||||
scope = parse_str(scope)
|
||||
space = parse_str(space)
|
||||
cop = parse_str(cop)
|
||||
vec = parse_str(vec)
|
||||
l1_evict = parse_str(l1_evict)
|
||||
l2_evict = parse_str(l2_evict)
|
||||
weak = _bool_attr(weak)
|
||||
has_cache = _bool_attr(has_cache_hint)
|
||||
ptx_type, c_type, constraint, _tvm_dtype = _type_info(ptx_type)
|
||||
if cop and cop not in _PTX_ST_COPS:
|
||||
raise ValueError(f"Unsupported PTX st cache operation {cop!r}")
|
||||
if vec and vec not in _PTX_VEC:
|
||||
raise ValueError(f"Unsupported PTX st vector modifier {vec!r}")
|
||||
if space not in _PTX_ST_SPACES and not (form == "mmio" and space == "global"):
|
||||
raise ValueError(f"Unsupported PTX st state space {space!r}")
|
||||
vec_len = int(vec[1:]) if vec else 1
|
||||
elem_bytes = (
|
||||
8
|
||||
if ptx_type.endswith("64")
|
||||
else 2
|
||||
if ptx_type in ("u16", "s16", "b16")
|
||||
else 1
|
||||
if ptx_type in ("u8", "s8", "b8")
|
||||
else 4
|
||||
)
|
||||
num_bytes = vec_len * elem_bytes if vec else elem_bytes
|
||||
use_cache_policy = form in ("weak", "relaxed", "release")
|
||||
if form == "mmio":
|
||||
if sem not in ("acquire", "relaxed", "release") or scope != "sys" or space != "global":
|
||||
raise ValueError("st.mmio requires sem, scope=sys, space=global")
|
||||
prefix = f"st.mmio.{sem}.{scope}"
|
||||
elif form == "relaxed":
|
||||
if not scope:
|
||||
raise ValueError("st.relaxed requires scope")
|
||||
prefix = f"st.relaxed.{scope}{_dot(space)}"
|
||||
elif form == "release":
|
||||
if not scope:
|
||||
raise ValueError("st.release requires scope")
|
||||
prefix = f"st.release.{scope}{_dot(space)}"
|
||||
elif form == "volatile":
|
||||
prefix = f"st.volatile{_dot(space)}"
|
||||
else:
|
||||
prefix = f"st{'.weak' if weak else ''}{_dot(space)}{_dot(cop)}"
|
||||
level = _ptx_level_suffix(has_cache, l1_evict, l2_evict, "")
|
||||
instr = f"{prefix}{level}{_dot(vec)}.{ptx_type}"
|
||||
name_parts = ["tvm_builtin_ptx_st", form if form != "weak" else ("weak" if weak else "plain")]
|
||||
if sem:
|
||||
name_parts.append(_safe_attr(sem))
|
||||
if scope:
|
||||
name_parts.append(_safe_attr(scope))
|
||||
name_parts.extend(
|
||||
[
|
||||
_safe_attr(space),
|
||||
_safe_attr(cop) if cop else "",
|
||||
_safe_attr(vec) if vec else "",
|
||||
ptx_type,
|
||||
"from_src" if from_src else "values",
|
||||
]
|
||||
)
|
||||
if has_cache:
|
||||
name_parts.append("cache_hint")
|
||||
name = "_".join(p for p in name_parts if p)
|
||||
values = f"{{{', '.join(f'%{i + 1}' for i in range(vec_len))}}}" if vec_len > 1 else "%1"
|
||||
value_constraints = "".join(f', "{constraint}"(value{i})' for i in range(vec_len))
|
||||
cache_slot = f", %{vec_len + 1}" if has_cache else ""
|
||||
cache_operand = ', "l"(cache_policy)' if has_cache else ""
|
||||
addr_decl, addr_operand = _ptx_shared_addr(space, "address")
|
||||
if from_src:
|
||||
load_regs = _ptx_st_load_src(num_bytes, vec_len, ptx_type, c_type)
|
||||
if vec_len > 1:
|
||||
in_constraints = ", ".join(f'"{constraint}"(r{i})' for i in range(vec_len))
|
||||
value_args = f", {in_constraints}"
|
||||
else:
|
||||
value_args = f', "{constraint}"(r0)'
|
||||
body = (
|
||||
f"{addr_decl}{load_regs}"
|
||||
f' asm volatile("{instr} [%0], {values}{cache_slot};"\n'
|
||||
" :\n"
|
||||
f" : {addr_operand}{value_args}{cache_operand}\n"
|
||||
' : "memory");'
|
||||
)
|
||||
if use_cache_policy:
|
||||
sig = "(void* address, void* src_ptr, unsigned long long cache_policy)"
|
||||
else:
|
||||
sig = "(void* address, void* src_ptr)"
|
||||
else:
|
||||
body = (
|
||||
f"{addr_decl}"
|
||||
f' asm volatile("{instr} [%0], {values}{cache_slot};"\n'
|
||||
" :\n"
|
||||
f" : {addr_operand}{value_constraints}{cache_operand}\n"
|
||||
' : "memory");'
|
||||
)
|
||||
if form == "mmio":
|
||||
sig = f"(void* address, {c_type} value0)"
|
||||
elif use_cache_policy:
|
||||
value_params = ", ".join(f"{c_type} value{i}" for i in range(vec_len))
|
||||
sig = f"(void* address, {value_params}, unsigned long long cache_policy)"
|
||||
else:
|
||||
value_params = ", ".join(f"{c_type} value{i}" for i in range(vec_len))
|
||||
sig = f"(void* address, {value_params})"
|
||||
return name, sig, body
|
||||
|
||||
|
||||
def _register_ptx_st(op_name, form, n_attrs, *, with_cache_policy=True):
|
||||
def codegen(*args):
|
||||
from_src = _bool_attr(args[-1])
|
||||
st_attrs = args[-n_attrs:-1]
|
||||
parts = _ptx_st_form_parts(form, st_attrs, from_src)
|
||||
forward = args[:-(n_attrs)]
|
||||
name, sig, body_str = parts
|
||||
source_code = f"\n__forceinline__ __device__ void {name}{sig} {{\n{body_str}\n}}\n"
|
||||
return cuda_func_call(name, *forward, source_code=source_code)
|
||||
|
||||
codegen.__name__ = f"codegen_{op_name}"
|
||||
register_codegen(op_name)(codegen)
|
||||
|
||||
|
||||
def _register_ptx_st_mmio(op_name, form, n_attrs):
|
||||
def codegen(*args):
|
||||
parts = _ptx_st_form_parts(form, args[-n_attrs:], False)
|
||||
forward = args[:-n_attrs]
|
||||
name, sig, body_str = parts
|
||||
source_code = f"\n__forceinline__ __device__ void {name}{sig} {{\n{body_str}\n}}\n"
|
||||
return cuda_func_call(name, *forward, source_code=source_code)
|
||||
|
||||
codegen.__name__ = f"codegen_{op_name}"
|
||||
register_codegen(op_name)(codegen)
|
||||
|
||||
|
||||
_register_ptx_st("ptx_st", "weak", 9)
|
||||
_register_ptx_st("ptx_st_relaxed", "relaxed", 8)
|
||||
_register_ptx_st("ptx_st_release", "release", 8)
|
||||
_register_ptx_st("ptx_st_volatile", "volatile", 4)
|
||||
_register_ptx_st_mmio("ptx_st_mmio", "mmio", 4)
|
||||
|
||||
|
||||
# PTX st.bulk form:
|
||||
# st.bulk{.weak}{.shared::cta} [a], size, initval;
|
||||
# ``initval`` is an immediate operand whose only legal value is 0.
|
||||
def _ptx_st_bulk_parts(_ptr, _num_bytes, weak, space):
|
||||
weak = bool(int(weak)) if hasattr(weak, "value") else bool(weak)
|
||||
space = parse_str(space)
|
||||
instr = f"st.bulk{'.weak' if weak else ''}{_dot(space)}"
|
||||
name = f"tvm_builtin_ptx_st_bulk{'_weak' if weak else ''}{('_' + _safe_attr(space)) if space else ''}"
|
||||
addr_arg = (
|
||||
'"r"((unsigned int)__cvta_generic_to_shared(ptr))' if space == "shared::cta" else '"l"(ptr)'
|
||||
)
|
||||
body = (
|
||||
f' asm volatile("{instr} [%0], %1, 0;"\n'
|
||||
" :\n"
|
||||
f" : {addr_arg}, "
|
||||
'"l"(static_cast<uint64_t>(num_bytes))\n'
|
||||
' : "memory");'
|
||||
)
|
||||
return name, body
|
||||
|
||||
|
||||
device_intrinsic(
|
||||
"ptx_st_bulk",
|
||||
n_attrs=2,
|
||||
helper_name=lambda *a: _ptx_st_bulk_parts(*a)[0],
|
||||
c_signature="(void* ptr, unsigned int num_bytes)",
|
||||
body=lambda *a: _ptx_st_bulk_parts(*a)[1],
|
||||
)
|
||||
|
||||
device_intrinsic(
|
||||
"cuda_uint_as_float",
|
||||
helper_name="tvm_builtin_uint_as_float",
|
||||
c_signature="(unsigned int bits)",
|
||||
return_type="float",
|
||||
body=" return __uint_as_float(bits);",
|
||||
)
|
||||
device_intrinsic(
|
||||
"cuda_float_as_uint",
|
||||
helper_name="tvm_builtin_float_as_uint",
|
||||
c_signature="(float x)",
|
||||
return_type="unsigned int",
|
||||
body=" return __float_as_uint(x);",
|
||||
)
|
||||
@@ -0,0 +1,253 @@
|
||||
# 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.
|
||||
# ruff: noqa: E501
|
||||
# pylint: disable=redefined-builtin, invalid-name
|
||||
"""Miscellaneous device helpers.
|
||||
|
||||
Catch-all for ops that don't fit the (sync / mma / cp_async / memory / math /
|
||||
nvshmem) feature buckets:
|
||||
|
||||
* PTX register-allocation control: ``setmaxnreg`` / ``mov`` from special reg.
|
||||
* Per-thread queries / scheduling hints: ``thread_rank`` / ``nano_sleep``.
|
||||
* Profiler timer hooks (``timer_init/start/end/finalize``).
|
||||
* Debug helpers: ``printf`` / ``trap`` on assert failure.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
|
||||
import tvm
|
||||
from tvm.backend.cuda.op import cuda_func_call
|
||||
|
||||
from ._schema import device_intrinsic
|
||||
from .registry import CODEGEN_REGISTRY, register_codegen
|
||||
from .utils import parse_str
|
||||
|
||||
# =============================================================================
|
||||
# setmaxnreg.{inc,dec}.sync.aligned.u32 — 1 PTX form (.action picks inc/dec).
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def _ptx_setmaxnreg(inc, nreg):
|
||||
inc = bool(int(inc)) if hasattr(inc, "value") else bool(inc)
|
||||
nreg = int(nreg)
|
||||
action = "inc" if inc else "dec"
|
||||
return (
|
||||
f"tvm_builtin_ptx_setmaxnreg_{action}_{nreg}",
|
||||
f' asm volatile("setmaxnreg.{action}.sync.aligned.u32 {nreg};");',
|
||||
)
|
||||
|
||||
|
||||
device_intrinsic(
|
||||
"ptx_setmaxnreg",
|
||||
n_attrs=2,
|
||||
helper_name=lambda inc, nreg: _ptx_setmaxnreg(inc, nreg)[0],
|
||||
body=lambda inc, nreg: _ptx_setmaxnreg(inc, nreg)[1],
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# mov.u32/u64 from special register — 1 PTX form (Form 2 of mov.type d, sreg).
|
||||
# Each (bits, reg) emits a distinct helper because the special reg name is
|
||||
# baked into the PTX text.
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def _ptx_fetch_register_body(bits):
|
||||
spec = "l" if bits == 64 else "r"
|
||||
|
||||
def _body(reg):
|
||||
reg = parse_str(reg)
|
||||
return (
|
||||
f" uint{bits}_t x;\n"
|
||||
f' asm volatile("mov.u{bits} %0, %{reg};" : "={spec}"(x));\n'
|
||||
f" return (int{bits}_t)x;"
|
||||
)
|
||||
|
||||
return _body
|
||||
|
||||
|
||||
for _bits in (32, 64):
|
||||
device_intrinsic(
|
||||
f"ptx_fetch_register_{_bits}",
|
||||
n_attrs=1,
|
||||
helper_name=(
|
||||
lambda *a, bits=_bits: (
|
||||
f"tvm_builtin_ptx_fetch_register_"
|
||||
f"{parse_str(a[-1]).replace('::', '_').replace('.', '_')}"
|
||||
)
|
||||
),
|
||||
return_type=f"int{_bits}_t",
|
||||
body=_ptx_fetch_register_body(_bits),
|
||||
)
|
||||
del _bits
|
||||
|
||||
|
||||
@register_codegen("ptx_fetch_register")
|
||||
def codegen_ptx_fetch_register(bits, reg):
|
||||
bits = int(bits)
|
||||
reg = parse_str(reg)
|
||||
if bits not in (32, 64):
|
||||
raise ValueError(f"Only support 32/64 bits for ptx_fetch_register, but got {bits}.")
|
||||
result = CODEGEN_REGISTRY[f"tirx.ptx_fetch_register_{bits}"]([reg])
|
||||
return result[0] if isinstance(result, tuple) else result
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Per-thread queries / scheduling hints.
|
||||
# =============================================================================
|
||||
device_intrinsic(
|
||||
"cuda_thread_rank",
|
||||
body=(
|
||||
" namespace cg = cooperative_groups;\n return cg::this_thread_block().thread_rank();"
|
||||
),
|
||||
return_type="int",
|
||||
tvm_return_type="int32",
|
||||
extra_deps=("cooperative_groups",),
|
||||
)
|
||||
device_intrinsic("cuda_nano_sleep", c_signature="(uint64_t time)", body=" __nanosleep(time);")
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Profiler timer hooks.
|
||||
# =============================================================================
|
||||
_COMMON_PARAMS = (
|
||||
"uint64_t* profiler_buffer, uint64_t* profiler_tag, "
|
||||
"uint32_t* profiler_write_offset, int profiler_write_stride, bool leader_cond"
|
||||
)
|
||||
_EVENT_PARAMS = f"int event_type, {_COMMON_PARAMS}"
|
||||
|
||||
|
||||
def _write_event(event_bits: str) -> str:
|
||||
return (
|
||||
"profiler_buffer[profiler_write_offset[0]] = "
|
||||
"((uint64_t)tvm_builtin_get_timestamp() << 32) | "
|
||||
f"(profiler_tag[0] | {event_bits});\n"
|
||||
" profiler_write_offset[0] += profiler_write_stride;"
|
||||
)
|
||||
|
||||
|
||||
device_intrinsic(
|
||||
"timer_init_cuda",
|
||||
c_signature=(
|
||||
"(uint64_t* profiler_buffer, uint64_t* profiler_tag, "
|
||||
"uint32_t* profiler_write_offset, int num_groups, int group_id)"
|
||||
),
|
||||
body=(
|
||||
" const uint32_t NBLOCKS = (uint32_t)(gridDim.x * gridDim.y * gridDim.z);\n"
|
||||
" const uint32_t BLOCK_IDX = (uint32_t)("
|
||||
"(blockIdx.z * gridDim.y + blockIdx.y) * gridDim.x + blockIdx.x);\n"
|
||||
" const uint32_t NGROUPS = num_groups;\n"
|
||||
" const uint32_t GROUP_ID = group_id;\n"
|
||||
" const uint32_t BLOCK_GROUP_IDX = BLOCK_IDX * NGROUPS + GROUP_ID;\n"
|
||||
" if ((blockIdx.x == 0) && (blockIdx.y == 0) && "
|
||||
"(blockIdx.z == 0) && (threadIdx.x == 0)) {\n"
|
||||
" profiler_buffer[0] = ((uint64_t)NGROUPS << 32) | NBLOCKS;\n"
|
||||
" }\n"
|
||||
" profiler_write_offset[0] = 1 + BLOCK_GROUP_IDX;\n"
|
||||
" profiler_tag[0] = (uint64_t)BLOCK_GROUP_IDX << 12;"
|
||||
),
|
||||
)
|
||||
|
||||
device_intrinsic(
|
||||
"timer_start_cuda",
|
||||
c_signature=f"({_EVENT_PARAMS})",
|
||||
body=(
|
||||
f" if (leader_cond) {{\n {_write_event('(uint32_t)event_type << 2 | 0x0')}\n }}\n"
|
||||
" __threadfence_block();"
|
||||
),
|
||||
extra_deps=("get_time_stamp",),
|
||||
)
|
||||
|
||||
device_intrinsic(
|
||||
"timer_end_cuda",
|
||||
c_signature=f"({_EVENT_PARAMS})",
|
||||
body=(
|
||||
" __threadfence_block();\n"
|
||||
f" if (leader_cond) {{\n {_write_event('(uint32_t)event_type << 2 | 0x1')}\n }}"
|
||||
),
|
||||
extra_deps=("get_time_stamp",),
|
||||
)
|
||||
|
||||
device_intrinsic(
|
||||
"timer_finalize_cuda",
|
||||
c_signature=f"({_COMMON_PARAMS})",
|
||||
body=(
|
||||
f" __threadfence_block();\n if (leader_cond) {{\n {_write_event('0x3')}\n }}"
|
||||
),
|
||||
extra_deps=("get_time_stamp",),
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Debug helpers — ``printf`` (variadic templated) and ``trap`` on assert.
|
||||
# =============================================================================
|
||||
device_intrinsic(
|
||||
"cuda_trap_when_assert_failed",
|
||||
c_signature="(bool cond)",
|
||||
body=' do {\n if (not (cond))\n asm("trap;");\n } while (0);',
|
||||
)
|
||||
|
||||
|
||||
@register_codegen("cuda_printf")
|
||||
def codegen_cuda_printf(fmt, *args):
|
||||
if isinstance(fmt, tvm.tirx.StringImm):
|
||||
fmt = fmt.value
|
||||
if not isinstance(fmt, str):
|
||||
raise ValueError("T.cuda.printf format must be a string literal")
|
||||
fmt_literal = json.dumps(fmt)
|
||||
arg_dtypes = [str(arg.ty) for arg in args]
|
||||
signature = "|".join([fmt, *arg_dtypes])
|
||||
digest = hashlib.sha1(signature.encode("utf-8")).hexdigest()
|
||||
func_name = f"tvm_builtin_cuda_printf_{len(args)}_{digest}"
|
||||
|
||||
def c_type(dtype: str) -> str:
|
||||
if dtype == "float32":
|
||||
return "float"
|
||||
if dtype == "float64":
|
||||
return "double"
|
||||
if dtype in {"int8", "int16", "int32"}:
|
||||
return "int"
|
||||
if dtype == "int64":
|
||||
return "long long"
|
||||
if dtype in {"uint8", "uint16", "uint32"}:
|
||||
return "unsigned int"
|
||||
if dtype == "uint64":
|
||||
return "unsigned long long"
|
||||
if dtype == "bool":
|
||||
return "int"
|
||||
if dtype == "handle":
|
||||
return "void*"
|
||||
raise ValueError(f"Unsupported T.cuda.printf argument dtype: {dtype}")
|
||||
|
||||
params = ", ".join(f"{c_type(dtype)} arg{i}" for i, dtype in enumerate(arg_dtypes))
|
||||
call_args = ", ".join(f"arg{i}" for i in range(len(args)))
|
||||
comma_call_args = f", {call_args}" if call_args else ""
|
||||
source_code = f"""
|
||||
__noinline__ __device__ void {func_name}({params}) {{
|
||||
printf({fmt_literal}{comma_call_args});
|
||||
}}
|
||||
"""
|
||||
return cuda_func_call(func_name, *args, source_code=source_code)
|
||||
|
||||
|
||||
device_intrinsic(
|
||||
"cuda_clock64",
|
||||
helper_name="tvm_builtin_clock64",
|
||||
return_type="unsigned long long",
|
||||
body=" return clock64();",
|
||||
)
|
||||
@@ -0,0 +1,482 @@
|
||||
# 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, too-many-arguments, too-many-locals, too-many-positional-arguments
|
||||
"""PTX MMA / ldmatrix / stmatrix intrinsics.
|
||||
|
||||
mma.sync.aligned has 7 form_kinds per the PTX docs (f16 / tf32 / bf16 / fp64
|
||||
/ int8 / fp8 / subbyte). Each form_kind is one ``device_intrinsic`` registration;
|
||||
the (shape, layouts, dtypes) modifier slots are attrs. Body computes the per-
|
||||
fragment register counts at codegen time from M*N*bits/threads/frag_size and
|
||||
hand-builds the asm constraint list.
|
||||
|
||||
ldmatrix / stmatrix each have a single PTX form (the .m8n8 .b16/.b8 variant
|
||||
that TIRx uses); ``num`` and ``trans`` are modifier attrs.
|
||||
"""
|
||||
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
|
||||
from tvm import DataType
|
||||
|
||||
from ._schema import device_intrinsic
|
||||
from .registry import CODEGEN_REGISTRY, register_codegen
|
||||
from .types import PTXDataType
|
||||
from .utils import parse_str
|
||||
|
||||
|
||||
@dataclass
|
||||
class FragAttrs:
|
||||
reg_type: str # asm constraint letter (r / f / d)
|
||||
size: int # bit width per register slot (32 or 64)
|
||||
ptr_type: str # C type for the cast
|
||||
|
||||
|
||||
_FRAG_ATTRS_MAP = {
|
||||
PTXDataType.BIT1: FragAttrs("r", 32, "uint32_t"),
|
||||
PTXDataType.INT4: FragAttrs("r", 32, "uint32_t"),
|
||||
PTXDataType.UINT4: FragAttrs("r", 32, "uint32_t"),
|
||||
PTXDataType.INT8: FragAttrs("r", 32, "uint32_t"),
|
||||
PTXDataType.UINT8: FragAttrs("r", 32, "uint32_t"),
|
||||
PTXDataType.FLOAT8_E4M3FN: FragAttrs("r", 32, "uint32_t"),
|
||||
PTXDataType.FLOAT8_E5M2: FragAttrs("r", 32, "uint32_t"),
|
||||
PTXDataType.BIT16: FragAttrs("r", 32, "uint32_t"),
|
||||
PTXDataType.FLOAT16: FragAttrs("r", 32, "uint32_t"),
|
||||
PTXDataType.BFLOAT16: FragAttrs("r", 32, "uint32_t"),
|
||||
PTXDataType.TENSOR_FLOAT32: FragAttrs("r", 32, "uint32_t"),
|
||||
PTXDataType.INT32: FragAttrs("r", 32, "int32_t"),
|
||||
PTXDataType.FLOAT32: FragAttrs("f", 32, "float"),
|
||||
PTXDataType.FLOAT64: FragAttrs("d", 64, "double"),
|
||||
}
|
||||
|
||||
|
||||
def _parse_mma_shape(shape_str):
|
||||
match = re.search(r"m(\d+)n(\d+)k(\d+)", shape_str)
|
||||
if not match:
|
||||
raise ValueError(f"Cannot parse MMA shape: {shape_str!r}")
|
||||
return tuple(map(int, match.groups()))
|
||||
|
||||
|
||||
def _classify_mma_form(d_type, a_type, b_type):
|
||||
"""Map (d, a, b) dtype triple to one of the 7 PTX form_kind tags."""
|
||||
fp16 = {"float16", "fp16"}
|
||||
tf32 = {"tensor_float32", "tf32"}
|
||||
bf16 = {"bfloat16", "bf16"}
|
||||
fp64 = {"float64", "fp64"}
|
||||
int_a = {"int8", "uint8", "s8", "u8"}
|
||||
fp8 = {"e4m3", "e5m2", "float8_e4m3fn", "float8_e4m3fnuz", "float8_e5m2"}
|
||||
subbyte = {"int4", "uint4", "bit1", "s4", "u4", "b1", "int1", "uint1"}
|
||||
if a_type in fp16 and b_type in fp16:
|
||||
return "f16"
|
||||
if a_type in tf32 and b_type in tf32:
|
||||
return "tf32"
|
||||
if a_type in bf16 and b_type in bf16:
|
||||
return "bf16"
|
||||
if a_type in fp64 and b_type in fp64:
|
||||
return "fp64"
|
||||
if a_type in int_a and b_type in int_a:
|
||||
return "int8"
|
||||
if a_type in fp8 and b_type in fp8:
|
||||
return "fp8"
|
||||
if a_type in subbyte and b_type in subbyte:
|
||||
return "subbyte"
|
||||
raise ValueError(
|
||||
f"Unknown ptx.mma form for d_type={d_type!r}, a_type={a_type!r}, b_type={b_type!r}"
|
||||
)
|
||||
|
||||
|
||||
def _frag(dtype_str):
|
||||
return _FRAG_ATTRS_MAP[PTXDataType.from_string(dtype_str)]
|
||||
|
||||
|
||||
def _mma_threads(shape, a_type):
|
||||
"""Special case: m8n8k4 with f16 a/b uses 8 threads per fragment."""
|
||||
m, n, k = _parse_mma_shape(shape)
|
||||
if m == 8 and n == 8 and k == 4 and a_type == "float16":
|
||||
return 8
|
||||
return 32
|
||||
|
||||
|
||||
# PTX dtype abbreviation -> element bit width. Used by _frag_count so that
|
||||
# callers passing the PTX abbreviation (e.g. "fp32") don't blow up in
|
||||
# ``DataType("fp32")``.
|
||||
_PTX_BITS = {
|
||||
"fp16": 16,
|
||||
"fp32": 32,
|
||||
"fp64": 64,
|
||||
"bf16": 16,
|
||||
"tf32": 32, # tensor-float32 packs 19 significant bits into a 32-bit slot
|
||||
"s8": 8,
|
||||
"u8": 8,
|
||||
"s32": 32,
|
||||
"s4": 4,
|
||||
"u4": 4,
|
||||
"b1": 1,
|
||||
"b16": 16,
|
||||
"e4m3": 8,
|
||||
"e5m2": 8,
|
||||
}
|
||||
|
||||
|
||||
def _frag_count(dtype, dim_a, dim_b, threads):
|
||||
if dtype in _PTX_BITS:
|
||||
bits = _PTX_BITS[dtype]
|
||||
else:
|
||||
bits = DataType(dtype).bits
|
||||
size = _frag(dtype).size
|
||||
return dim_a * dim_b * bits // threads // size
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Shared helpers for the 7 mma form_kinds.
|
||||
# Args layout for each form:
|
||||
# (d_ptr_in, a_ptr_in, b_ptr_in [, c_ptr_in], shape, a_layout, b_layout,
|
||||
# d_type, a_type, b_type, c_type, no_c_ptr [, saturate or bit_op])
|
||||
# n_attrs = 8 for f16/tf32/bf16/fp64/fp8 (last 8 = shape, layouts, 4 dtypes, no_c_ptr)
|
||||
# n_attrs = 9 for int8 (+ saturate) and subbyte (+ bit_op)
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def _mma_form_parts(args, *, has_saturate=False, has_bit_op=False):
|
||||
"""Compute (helper_name, c_signature, body) for one mma form invocation.
|
||||
|
||||
``args`` is the full positional arg tuple as received by codegen.
|
||||
The trailing ``n_attrs`` (8 or 9) entries are attrs.
|
||||
"""
|
||||
n_extra = (1 if has_saturate else 0) + (1 if has_bit_op else 0)
|
||||
n_attrs = 8 + n_extra
|
||||
# Split off attr args from the tail (operand args are ahead).
|
||||
attrs = args[-n_attrs:]
|
||||
shape = parse_str(attrs[0])
|
||||
a_layout = parse_str(attrs[1])
|
||||
b_layout = parse_str(attrs[2])
|
||||
d_type = parse_str(attrs[3])
|
||||
a_type = parse_str(attrs[4])
|
||||
b_type = parse_str(attrs[5])
|
||||
c_type = parse_str(attrs[6])
|
||||
no_c_ptr_raw = attrs[7]
|
||||
no_c_ptr = bool(int(no_c_ptr_raw)) if hasattr(no_c_ptr_raw, "value") else bool(no_c_ptr_raw)
|
||||
saturate = False
|
||||
bit_op = ""
|
||||
if has_saturate:
|
||||
s = attrs[8]
|
||||
saturate = bool(int(s)) if hasattr(s, "value") else bool(s)
|
||||
if has_bit_op:
|
||||
bit_op = parse_str(attrs[8])
|
||||
|
||||
# Fragment counts (same derivation as the contiguous form).
|
||||
m, n, k = _parse_mma_shape(shape)
|
||||
threads = _mma_threads(shape, a_type)
|
||||
d_cnt = _frag_count(d_type, m, n, threads)
|
||||
a_cnt = _frag_count(a_type, m, k, threads)
|
||||
b_cnt = _frag_count(b_type, k, n, threads)
|
||||
c_cnt = _frag_count(c_type, m, n, threads)
|
||||
|
||||
# C signature: one void* per register, ordered D regs, A regs, B regs,
|
||||
# then C regs (only when the accumulator is used).
|
||||
sig_parts = (
|
||||
[f"void* d_ptr{i}" for i in range(d_cnt)]
|
||||
+ [f"void* a_ptr{i}" for i in range(a_cnt)]
|
||||
+ [f"void* b_ptr{i}" for i in range(b_cnt)]
|
||||
)
|
||||
if not no_c_ptr:
|
||||
sig_parts += [f"void* c_ptr{i}" for i in range(c_cnt)]
|
||||
sig = "(" + ", ".join(sig_parts) + ")"
|
||||
|
||||
def _safe(s):
|
||||
return s.replace("::", "_").replace(".", "_")
|
||||
|
||||
name = (
|
||||
f"ptx_mma_{shape}_{a_layout}_{b_layout}"
|
||||
f"_{_safe(d_type)}_{_safe(a_type)}_{_safe(b_type)}_{_safe(c_type)}"
|
||||
f"{'_no_c_ptr' if no_c_ptr else ''}"
|
||||
f"{'_saturate' if saturate else ''}"
|
||||
)
|
||||
|
||||
d_frag = _frag(d_type)
|
||||
a_frag = _frag(a_type)
|
||||
b_frag = _frag(b_type)
|
||||
c_frag = _frag(c_type)
|
||||
|
||||
saturate_inst = ".satfinite" if saturate else ""
|
||||
# PTX b1 mma requires a `.popc` suffix after the bit op (e.g. `.xor.popc`).
|
||||
bit_op_inst = f".{bit_op}.popc" if bit_op else ""
|
||||
|
||||
d_type_inst = PTXDataType.from_string(d_type).to_string()
|
||||
c_type_inst = PTXDataType.from_string(c_type).to_string()
|
||||
a_type_inst = PTXDataType.from_string(a_type).to_string()
|
||||
b_type_inst = PTXDataType.from_string(b_type).to_string()
|
||||
|
||||
def _slot_arr(start, cnt):
|
||||
return "{" + ", ".join(f"%{start + i}" for i in range(cnt)) + "}"
|
||||
|
||||
args_template = (
|
||||
f"{_slot_arr(0, d_cnt)}, {_slot_arr(d_cnt, a_cnt)}, "
|
||||
f"{_slot_arr(d_cnt + a_cnt, b_cnt)}, {_slot_arr(d_cnt + a_cnt + b_cnt, c_cnt)}"
|
||||
)
|
||||
|
||||
# Each register binds to its OWN pointer via *(T*)X_ptrN (scatter).
|
||||
d_outs = ", ".join(
|
||||
f'"=r"(*({d_frag.ptr_type}*)d_ptr{i})'
|
||||
if d_frag.reg_type == "r"
|
||||
else f'"={d_frag.reg_type}"(*({d_frag.ptr_type}*)d_ptr{i})'
|
||||
for i in range(d_cnt)
|
||||
)
|
||||
a_inputs = ", ".join(
|
||||
f'"{a_frag.reg_type}"(*({a_frag.ptr_type}*)a_ptr{i})' for i in range(a_cnt)
|
||||
)
|
||||
b_inputs = ", ".join(
|
||||
f'"{b_frag.reg_type}"(*({b_frag.ptr_type}*)b_ptr{i})' for i in range(b_cnt)
|
||||
)
|
||||
if no_c_ptr:
|
||||
c_value = "0.f" if c_frag.reg_type == "f" else "0"
|
||||
c_inputs = ", ".join(f'"{c_frag.reg_type}"({c_value})' for _ in range(c_cnt))
|
||||
else:
|
||||
c_inputs = ", ".join(
|
||||
f'"{c_frag.reg_type}"(*({c_frag.ptr_type}*)c_ptr{i})' for i in range(c_cnt)
|
||||
)
|
||||
|
||||
body = (
|
||||
" asm volatile(\n"
|
||||
f' "mma.sync.aligned.{shape}.{a_layout}.{b_layout}{saturate_inst}'
|
||||
f'{d_type_inst}{a_type_inst}{b_type_inst}{c_type_inst}{bit_op_inst} "\n'
|
||||
f' "{args_template};\\n"\n'
|
||||
f" : {d_outs}\n"
|
||||
f" : {a_inputs}, {b_inputs}, {c_inputs}\n"
|
||||
" );"
|
||||
)
|
||||
return name, sig, body
|
||||
|
||||
|
||||
def _register_mma_form(form_kind, *, has_saturate=False, has_bit_op=False):
|
||||
n_attrs = 8 + (1 if has_saturate else 0) + (1 if has_bit_op else 0)
|
||||
|
||||
def _parts(*args, hs=has_saturate, hb=has_bit_op):
|
||||
return _mma_form_parts(args, has_saturate=hs, has_bit_op=hb)
|
||||
|
||||
device_intrinsic(
|
||||
f"_ptx_mma_{form_kind}",
|
||||
n_attrs=n_attrs,
|
||||
helper_name=lambda *a: _parts(*a)[0],
|
||||
c_signature=lambda *a: _parts(*a)[1],
|
||||
body=lambda *a: _parts(*a)[2],
|
||||
)
|
||||
|
||||
|
||||
# Form 1 — f16. Form 2 — tf32. Form 3 — bf16. Form 4 — fp64. Form 6 — fp8.
|
||||
# All share the same 8-attr layout (no saturate / bit_op).
|
||||
for _kind in ("f16", "tf32", "bf16", "fp64", "fp8"):
|
||||
_register_mma_form(_kind)
|
||||
del _kind
|
||||
|
||||
# Form 5 — int8 (+ saturate).
|
||||
_register_mma_form("int8", has_saturate=True)
|
||||
|
||||
# Form 7 — subbyte (+ bit_op for b1).
|
||||
_register_mma_form("subbyte", has_bit_op=True)
|
||||
|
||||
|
||||
@register_codegen("ptx_mma")
|
||||
def codegen_ptx_mma(
|
||||
shape,
|
||||
a_layout,
|
||||
b_layout,
|
||||
d_type,
|
||||
a_type,
|
||||
b_type,
|
||||
c_type,
|
||||
d_cnt,
|
||||
a_cnt,
|
||||
b_cnt,
|
||||
c_cnt,
|
||||
no_c_ptr,
|
||||
*rest,
|
||||
):
|
||||
"""Classify (d, a, b) dtype triple to one of 7 form_kinds and forward.
|
||||
|
||||
``rest`` = flattened per-register pointers (d_cnt + a_cnt + b_cnt + c_cnt of
|
||||
them) followed by ``saturate`` and optionally ``bit_op``.
|
||||
"""
|
||||
shape = parse_str(shape)
|
||||
a_layout = parse_str(a_layout)
|
||||
b_layout = parse_str(b_layout)
|
||||
d_type = parse_str(d_type)
|
||||
a_type = parse_str(a_type)
|
||||
b_type = parse_str(b_type)
|
||||
c_type = parse_str(c_type)
|
||||
d_cnt = int(d_cnt)
|
||||
a_cnt = int(a_cnt)
|
||||
b_cnt = int(b_cnt)
|
||||
c_cnt = int(c_cnt)
|
||||
no_c_ptr = bool(int(no_c_ptr)) if hasattr(no_c_ptr, "value") else bool(no_c_ptr)
|
||||
|
||||
n_ptrs = d_cnt + a_cnt + b_cnt + (0 if no_c_ptr else c_cnt)
|
||||
ptrs = list(rest[:n_ptrs])
|
||||
trailing = list(rest[n_ptrs:])
|
||||
saturate = bool(trailing[0]) if trailing else False
|
||||
bit_op_v = ""
|
||||
if len(trailing) >= 2:
|
||||
bo = trailing[1]
|
||||
bit_op_v = parse_str(bo) if isinstance(bo, str) else (bo if bo is not None else "")
|
||||
|
||||
kind = _classify_mma_form(d_type, a_type, b_type)
|
||||
|
||||
# op_args are the flattened per-register pointers (already in PTX order:
|
||||
# D regs, A regs, B regs, then C regs unless no_c_ptr).
|
||||
op_args = ptrs
|
||||
|
||||
attr_args = [shape, a_layout, b_layout, d_type, a_type, b_type, c_type, no_c_ptr]
|
||||
if kind == "int8":
|
||||
attr_args.append(saturate)
|
||||
elif kind == "subbyte":
|
||||
attr_args.append(bit_op_v)
|
||||
|
||||
result = CODEGEN_REGISTRY[f"tirx._ptx_mma_{kind}"](op_args + attr_args)
|
||||
return result[0] if isinstance(result, tuple) else result
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# ldmatrix / stmatrix — m8n8 fragment load/store. PTX docs lists 3 ldmatrix
|
||||
# forms (m8n8 + m8n16 + m16n16); TIRx uses only the m8n8 form. 1
|
||||
# device_intrinsic each. ``num`` (.x1/.x2/.x4) and ``trans`` are modifier
|
||||
# attrs; the asm body loops over per-register constraints based on
|
||||
# (num, dtype).
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def _ldmatrix_parts(*args):
|
||||
# args = (smem_ptr, dst0, dst1, ..., dst{N-1}, num, dtype, trans)
|
||||
# The last 3 entries are the codegen attrs (n_attrs=3).
|
||||
num = int(args[-3])
|
||||
dtype = parse_str(args[-2])
|
||||
trans_b = bool(int(args[-1])) if hasattr(args[-1], "value") else bool(args[-1])
|
||||
if num not in (1, 2, 4):
|
||||
raise ValueError(f"ldmatrix .num must be one of {{1, 2, 4}}, got {num}")
|
||||
if dtype not in ("b16", "b8"):
|
||||
raise ValueError(f"ldmatrix dtype must be 'b16' or 'b8', got {dtype!r}")
|
||||
n_regs = num if dtype == "b16" else num // 2
|
||||
trans_inst = ".trans" if trans_b else ""
|
||||
slot_list = "{" + ", ".join(f"%{i}" for i in range(n_regs)) + "}"
|
||||
reg_decls = ", ".join(f"r{i}" for i in range(n_regs))
|
||||
out_constraints = ", ".join(f'"=r"(r{i})' for i in range(n_regs))
|
||||
dst_assigns = "\n".join(f" *(uint32_t*)dst{i} = r{i};" for i in range(n_regs))
|
||||
name = f"ptx_ldmatrix_{num}_{dtype.replace('::', '_').replace('.', '_')}_{1 if trans_b else 0}"
|
||||
sig = "(void* smem_ptr, " + ", ".join(f"void* dst{i}" for i in range(n_regs)) + ")"
|
||||
body = (
|
||||
f" uint32_t {reg_decls};\n"
|
||||
" unsigned int addr = __cvta_generic_to_shared(smem_ptr);\n"
|
||||
" asm volatile(\n"
|
||||
f' "ldmatrix.sync.aligned.m8n8.x{num}{trans_inst}.shared.{dtype} '
|
||||
f'{slot_list}, [%{n_regs}];"\n'
|
||||
f" : {out_constraints}\n"
|
||||
f' : "r"(addr));\n'
|
||||
f"{dst_assigns}"
|
||||
)
|
||||
return name, sig, body
|
||||
|
||||
|
||||
device_intrinsic(
|
||||
"_ptx_ldmatrix_impl",
|
||||
n_attrs=3,
|
||||
c_signature=lambda *a: _ldmatrix_parts(*a)[1],
|
||||
helper_name=lambda *a: _ldmatrix_parts(*a)[0],
|
||||
body=lambda *a: _ldmatrix_parts(*a)[2],
|
||||
)
|
||||
|
||||
|
||||
@register_codegen("ptx_ldmatrix")
|
||||
def codegen_ptx_ldmatrix(trans, num, dtype, smem_ptr, *dst_handles):
|
||||
trans = bool(trans)
|
||||
num = int(num)
|
||||
dtype = parse_str(dtype)
|
||||
if dtype.startswith("."):
|
||||
dtype = dtype[1:]
|
||||
n_regs = num if dtype == "b16" else num // 2
|
||||
if len(dst_handles) != n_regs:
|
||||
raise ValueError(
|
||||
f"ldmatrix .x{num}.{dtype} codegen expects {n_regs} dst handles, got {len(dst_handles)}"
|
||||
)
|
||||
result = CODEGEN_REGISTRY["tirx._ptx_ldmatrix_impl"](
|
||||
[smem_ptr, *dst_handles, num, dtype, trans]
|
||||
)
|
||||
return result[0] if isinstance(result, tuple) else result
|
||||
|
||||
|
||||
def _stmatrix_parts(*args):
|
||||
# args = (smem_ptr, src0, src1, ..., src{N-1}, trans, num, dtype, shape, space)
|
||||
# The last 5 entries are codegen attrs (n_attrs=5).
|
||||
trans_arg, num_arg, dtype_arg, shape_arg, space_arg = args[-5:]
|
||||
n_regs = int(num_arg)
|
||||
trans_b = bool(int(trans_arg)) if hasattr(trans_arg, "value") else bool(trans_arg)
|
||||
dtype = parse_str(dtype_arg)
|
||||
shape = parse_str(shape_arg)
|
||||
space = parse_str(space_arg)
|
||||
if dtype.startswith("."):
|
||||
dtype = dtype[1:]
|
||||
if n_regs not in (1, 2, 4):
|
||||
raise ValueError(f"stmatrix .num must be one of {{1, 2, 4}}, got {n_regs}")
|
||||
if dtype not in ("b16", "b8"):
|
||||
raise ValueError(f"stmatrix .type must be b16 or b8, got {dtype!r}")
|
||||
if shape not in ("m8n8", "m16n8"):
|
||||
raise ValueError(f"stmatrix .shape must be m8n8 or m16n8, got {shape!r}")
|
||||
if space not in ("shared", "shared::cta"):
|
||||
raise ValueError(f"stmatrix state space must be shared or shared::cta, got {space!r}")
|
||||
if shape == "m16n8" and not trans_b:
|
||||
raise ValueError("stmatrix .m16n8 requires .trans")
|
||||
trans_inst = ".trans" if trans_b else ""
|
||||
slot_list = "{" + ", ".join(f"%{i}" for i in range(n_regs)) + "}"
|
||||
src_loads = "\n".join(f" uint32_t r{i} = *(uint32_t*)src{i};" for i in range(n_regs))
|
||||
in_constraints = ", ".join(f'"r"(r{i})' for i in range(n_regs))
|
||||
name = f"ptx_stmatrix_{shape}_{n_regs}_{1 if trans_b else 0}_{space.replace('::', '_')}_{dtype}"
|
||||
sig = "(void* smem_ptr, " + ", ".join(f"void* src{i}" for i in range(n_regs)) + ")"
|
||||
body = (
|
||||
f"{src_loads}\n"
|
||||
" unsigned int addr = __cvta_generic_to_shared(smem_ptr);\n"
|
||||
" asm volatile(\n"
|
||||
f' "stmatrix.sync.aligned.{shape}.x{n_regs}{trans_inst}.{space}.{dtype} '
|
||||
f'[%{n_regs}], {slot_list};"\n'
|
||||
" :\n"
|
||||
f' : {in_constraints}, "r"(addr));'
|
||||
)
|
||||
return name, sig, body
|
||||
|
||||
|
||||
device_intrinsic(
|
||||
"_ptx_stmatrix_impl",
|
||||
n_attrs=5,
|
||||
c_signature=lambda *a: _stmatrix_parts(*a)[1],
|
||||
helper_name=lambda *a: _stmatrix_parts(*a)[0],
|
||||
body=lambda *a: _stmatrix_parts(*a)[2],
|
||||
)
|
||||
|
||||
|
||||
@register_codegen("ptx_stmatrix")
|
||||
def codegen_ptx_stmatrix(trans, num, dtype, shape, space, smem_ptr, *src_handles):
|
||||
trans = bool(trans)
|
||||
num = int(num)
|
||||
dtype_str = parse_str(dtype)
|
||||
if dtype_str.startswith("."):
|
||||
dtype_str = dtype_str[1:]
|
||||
n_regs = num
|
||||
if len(src_handles) != n_regs:
|
||||
raise ValueError(
|
||||
f"stmatrix .x{num}.{dtype_str} codegen expects {n_regs} src handles, "
|
||||
f"got {len(src_handles)}"
|
||||
)
|
||||
result = CODEGEN_REGISTRY["tirx._ptx_stmatrix_impl"](
|
||||
[smem_ptr, *src_handles, trans, num, dtype, shape, space]
|
||||
)
|
||||
return result[0] if isinstance(result, tuple) else result
|
||||
@@ -0,0 +1,161 @@
|
||||
# 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
|
||||
@@ -0,0 +1,77 @@
|
||||
# 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.
|
||||
"""Codegen registry for CUDA HW ops.
|
||||
|
||||
User-facing Python wrappers are hand-written in :mod:`tvm.tirx.op` so that
|
||||
editors / static analyzers (Cursor, Pyright) can see their signatures. This
|
||||
module only handles the backend codegen side.
|
||||
"""
|
||||
|
||||
import functools
|
||||
|
||||
import tvm_ffi
|
||||
|
||||
CODEGEN_REGISTRY = {}
|
||||
|
||||
|
||||
def _canonical_device_intrin_name(op_name: str) -> str:
|
||||
if not op_name.startswith("tirx."):
|
||||
return op_name
|
||||
basename = op_name[len("tirx.") :]
|
||||
if "." in basename:
|
||||
return op_name
|
||||
for prefix, namespace in (
|
||||
("cuda_", "cuda"),
|
||||
("ptx_", "ptx"),
|
||||
("nvshmem_", "nvshmem"),
|
||||
("nki_", "nki"),
|
||||
):
|
||||
if basename.startswith(prefix):
|
||||
return f"tirx.{namespace}.{basename[len(prefix) :]}"
|
||||
return op_name
|
||||
|
||||
|
||||
@tvm_ffi.register_global_func("tirx.intrinsics.cuda.get_codegen")
|
||||
def get_codegen(op):
|
||||
"""get the codegen function for a given op"""
|
||||
return CODEGEN_REGISTRY.get(op, None)
|
||||
|
||||
|
||||
def register_codegen(op, backend="cuda"):
|
||||
"""Register a codegen function for a given op.
|
||||
|
||||
The codegen function should return a ``cuda_func_call`` statement, and
|
||||
optionally a list of tags that the codegen function needs.
|
||||
"""
|
||||
|
||||
def decorator(func):
|
||||
full_op_name = "tirx." + op
|
||||
canonical_op_name = _canonical_device_intrin_name(full_op_name)
|
||||
op_names = {full_op_name, canonical_op_name}
|
||||
|
||||
@functools.wraps(func)
|
||||
def wrapper(arg_list):
|
||||
res = func(*arg_list) # pylint: disable=not-callable
|
||||
if isinstance(res, tuple):
|
||||
return res[0], res[1]
|
||||
return res, list()
|
||||
|
||||
for op_name in op_names:
|
||||
CODEGEN_REGISTRY[op_name] = wrapper
|
||||
return wrapper
|
||||
|
||||
return decorator
|
||||
@@ -0,0 +1,570 @@
|
||||
# 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
|
||||
"""Synchronization primitives.
|
||||
|
||||
PTX side:
|
||||
* ``bar.arrive`` / ``bar.sync`` — named-barrier alias of ``barrier.arrive/sync``
|
||||
* ``fence{.sem}.scope`` / ``fence.proxy.async`` / ``fence.mbarrier_init``
|
||||
* ``barrier.cluster.arrive`` / ``barrier.cluster.wait``
|
||||
* ``mbarrier.init`` / ``mbarrier.arrive[.expect_tx]`` (local + remote) / ``mbarrier.try_wait``
|
||||
* ``elect.sync`` — warp leader election
|
||||
* warp-vote ``__any_sync``
|
||||
|
||||
CUDA-side helpers:
|
||||
* ``__threadfence`` / ``__syncwarp`` / ``__syncthreads`` / ``__syncthreads_and|or``
|
||||
* cooperative-groups grid sync
|
||||
* cluster sync (open-coded ``barrier.cluster.arrive/wait`` pair)
|
||||
* warpgroup sync (``bar.sync``)
|
||||
"""
|
||||
|
||||
from tvm.tirx.operator.intrinsics._common import (
|
||||
CLUSTER_BARRIER_SEM,
|
||||
FENCE_PROXY_ASYNC_SPACE,
|
||||
FENCE_SCOPE,
|
||||
FENCE_SEM,
|
||||
)
|
||||
|
||||
from ._schema import device_intrinsic
|
||||
from .registry import CODEGEN_REGISTRY, register_codegen
|
||||
from .utils import parse_str
|
||||
|
||||
# =============================================================================
|
||||
# bar.arrive / bar.sync — alias of barrier.arrive/sync. 1 form each.
|
||||
# bar.sync a, b ;
|
||||
# bar.arrive a, b ;
|
||||
# =============================================================================
|
||||
device_intrinsic(
|
||||
"ptx_bar_arrive",
|
||||
c_signature="(int name_bar_id, int thread_count)",
|
||||
body=(
|
||||
' asm volatile("bar.arrive %0, %1;" : : "r"(name_bar_id), "r"(thread_count) : "memory");'
|
||||
),
|
||||
)
|
||||
device_intrinsic(
|
||||
"ptx_bar_sync",
|
||||
c_signature="(int name_bar_id, int thread_count)",
|
||||
body=(
|
||||
' asm volatile("bar.sync %0, %1;" : : "r"(name_bar_id), "r"(thread_count) : "memory");'
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# fence{.sem}.scope — 1 form (sem/scope are modifier values).
|
||||
# =============================================================================
|
||||
def _ptx_fence(sem, scope):
|
||||
sem, scope = parse_str(sem), parse_str(scope)
|
||||
assert sem in FENCE_SEM, f"invalid fence sem {sem!r}, expected one of {FENCE_SEM}"
|
||||
assert scope in FENCE_SCOPE, f"invalid fence scope {scope!r}, expected one of {FENCE_SCOPE}"
|
||||
return (
|
||||
f"tvm_builtin_ptx_fence_{sem}_{scope}",
|
||||
f' asm volatile("fence.{sem}.{scope};" ::: "memory");',
|
||||
)
|
||||
|
||||
|
||||
device_intrinsic(
|
||||
"ptx_fence",
|
||||
n_attrs=2,
|
||||
helper_name=lambda sem, scope: _ptx_fence(sem, scope)[0],
|
||||
body=lambda sem, scope: _ptx_fence(sem, scope)[1],
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# fence.proxy.async{.<space>} — 1 form, optional .space modifier.
|
||||
# =============================================================================
|
||||
def _ptx_fence_proxy_async(space):
|
||||
space = parse_str(space)
|
||||
assert space in FENCE_PROXY_ASYNC_SPACE, (
|
||||
f"invalid fence.proxy.async space {space!r}, expected one of {FENCE_PROXY_ASYNC_SPACE}"
|
||||
)
|
||||
suffix = f".{space}" if space else ""
|
||||
name_safe = "_" + space.replace("::", "_").replace(".", "_") if space else ""
|
||||
return (
|
||||
f"tvm_builtin_ptx_fence_proxy_async{name_safe}",
|
||||
f' asm volatile("fence.proxy.async{suffix};" ::: "memory");',
|
||||
)
|
||||
|
||||
|
||||
device_intrinsic(
|
||||
"ptx_fence_proxy_async",
|
||||
n_attrs=1,
|
||||
helper_name=lambda space: _ptx_fence_proxy_async(space)[0],
|
||||
body=lambda space: _ptx_fence_proxy_async(space)[1],
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# fence.mbarrier_init.release.cluster — 1 form, no operands.
|
||||
# =============================================================================
|
||||
device_intrinsic(
|
||||
"ptx_fence_mbarrier_init",
|
||||
body=' asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory");',
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# barrier.cluster.arrive{.sem}{.aligned} — 1 form.
|
||||
# =============================================================================
|
||||
def _ptx_barrier_cluster_arrive(sem, aligned):
|
||||
sem = parse_str(sem)
|
||||
aligned = bool(int(aligned)) if hasattr(aligned, "value") else bool(aligned)
|
||||
assert sem in CLUSTER_BARRIER_SEM, (
|
||||
f"invalid cluster.arrive sem {sem!r}, expected one of {CLUSTER_BARRIER_SEM}"
|
||||
)
|
||||
sem_suffix = f".{sem}" if sem else ""
|
||||
aligned_suffix = ".aligned" if aligned else ""
|
||||
name_sem = "_" + sem.replace("::", "_").replace(".", "_") if sem else ""
|
||||
name_aligned = "_aligned" if aligned else ""
|
||||
return (
|
||||
f"tvm_builtin_ptx_barrier_cluster_arrive{name_sem}{name_aligned}",
|
||||
f' asm volatile("barrier.cluster.arrive{sem_suffix}{aligned_suffix};" ::: "memory");',
|
||||
)
|
||||
|
||||
|
||||
device_intrinsic(
|
||||
"ptx_barrier_cluster_arrive",
|
||||
n_attrs=2,
|
||||
helper_name=lambda sem, aligned: _ptx_barrier_cluster_arrive(sem, aligned)[0],
|
||||
body=lambda sem, aligned: _ptx_barrier_cluster_arrive(sem, aligned)[1],
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# barrier.cluster.wait{.acquire}{.aligned} — 1 form.
|
||||
# =============================================================================
|
||||
def _ptx_barrier_cluster_wait(acquire, aligned):
|
||||
acquire = bool(int(acquire)) if hasattr(acquire, "value") else bool(acquire)
|
||||
aligned = bool(int(aligned)) if hasattr(aligned, "value") else bool(aligned)
|
||||
acq_suffix = ".acquire" if acquire else ""
|
||||
aligned_suffix = ".aligned" if aligned else ""
|
||||
return (
|
||||
f"tvm_builtin_ptx_barrier_cluster_wait"
|
||||
f"{'_acquire' if acquire else ''}{'_aligned' if aligned else ''}",
|
||||
f' asm volatile("barrier.cluster.wait{acq_suffix}{aligned_suffix};" ::: "memory");',
|
||||
)
|
||||
|
||||
|
||||
device_intrinsic(
|
||||
"ptx_barrier_cluster_wait",
|
||||
n_attrs=2,
|
||||
helper_name=lambda acquire, aligned: _ptx_barrier_cluster_wait(acquire, aligned)[0],
|
||||
body=lambda acquire, aligned: _ptx_barrier_cluster_wait(acquire, aligned)[1],
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# clusterlaunchcontrol.try_cancel / query_cancel — Blackwell Cluster Launch
|
||||
# Control (CLC) work-stealing, written from the PTX ISA spec (section
|
||||
# "clusterlaunchcontrol", PTX ISA 8.6). try_cancel async-requests cancelling the
|
||||
# next cluster's launch, writing a 16B response to smem + signalling mbar. query
|
||||
# decodes the response: on success it extracts the cancelled cluster's first
|
||||
# ctaid.x (via the get_first_ctaid::x form); a single uint32 is returned, with
|
||||
# 0xFFFFFFFF as the "no work stolen" sentinel (a device helper returns one scalar).
|
||||
# =============================================================================
|
||||
device_intrinsic(
|
||||
"ptx_clc_try_cancel",
|
||||
c_signature="(void* handle, void* mbar)",
|
||||
body=(
|
||||
" unsigned int addr = (unsigned int)__cvta_generic_to_shared(handle);\n"
|
||||
" unsigned int bar = (unsigned int)__cvta_generic_to_shared(mbar);\n"
|
||||
" asm volatile(\n"
|
||||
' "clusterlaunchcontrol.try_cancel.async.shared::cta.mbarrier::complete_tx::bytes"\n'
|
||||
' ".multicast::cluster::all.b128 [%0], [%1];\\n"\n'
|
||||
' :: "r"(addr), "r"(bar) : "memory");'
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
device_intrinsic(
|
||||
"ptx_clc_query_cancel",
|
||||
c_signature="(void* handle)",
|
||||
return_type="uint32_t",
|
||||
tvm_return_type="uint32",
|
||||
body=(
|
||||
" unsigned int addr = (unsigned int)__cvta_generic_to_shared(handle);\n"
|
||||
" unsigned int first_ctaid_x;\n"
|
||||
" asm volatile(\n"
|
||||
' "{\\n"\n'
|
||||
' ".reg .pred canceled;\\n"\n'
|
||||
' ".reg .b128 response;\\n"\n'
|
||||
' "ld.shared.b128 response, [%1];\\n"\n'
|
||||
' "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 canceled, response;\\n"\n'
|
||||
' "mov.u32 %0, 0xffffffff;\\n"\n'
|
||||
' "@canceled clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128"\n'
|
||||
' " %0, response;\\n"\n'
|
||||
' "}\\n"\n'
|
||||
' : "=r"(first_ctaid_x) : "r"(addr) : "memory");\n'
|
||||
' asm volatile("fence.proxy.async.shared::cta;\\n" ::: "memory");\n'
|
||||
" return first_ctaid_x;"
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# mbarrier.init.shared.b64 [addr], count ; — 1 form.
|
||||
# =============================================================================
|
||||
device_intrinsic(
|
||||
"ptx_mbarrier_init",
|
||||
c_signature="(void* barrier, int thread_count)",
|
||||
body=(
|
||||
" unsigned int barrier_addr = __cvta_generic_to_shared(barrier);\n"
|
||||
' asm volatile("mbarrier.init.shared.b64 [%0], %1;"'
|
||||
' : : "r"(barrier_addr), "r"(thread_count) : "memory");'
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# mbarrier.arrive — local + remote (cluster-mapped) forms. 2 PTX forms.
|
||||
# Form local: mbarrier.arrive.shared.b64 _, [bar];
|
||||
# Form remote: { setp+@p mapa.shared::cluster.u32 + @p mbarrier.arrive.shared::cluster.b64 }
|
||||
# Dispatcher picks by arg count (1 vs 3).
|
||||
# =============================================================================
|
||||
device_intrinsic(
|
||||
"_ptx_mbarrier_arrive_local",
|
||||
helper_name="tvm_builtin_ptx_mbarrier_arrive",
|
||||
c_signature="(void* barrier)",
|
||||
body=(
|
||||
" unsigned int barrier_addr = __cvta_generic_to_shared(barrier);\n"
|
||||
' asm volatile("mbarrier.arrive.shared.b64 _, [%0];"\n'
|
||||
' :: "r"(barrier_addr) : "memory");'
|
||||
),
|
||||
)
|
||||
device_intrinsic(
|
||||
"_ptx_mbarrier_arrive_remote",
|
||||
helper_name="tvm_builtin_ptx_mbarrier_arrive_remote",
|
||||
c_signature="(void* barrier, int cta_id, int pred)",
|
||||
body=(
|
||||
" unsigned int barrier_addr = __cvta_generic_to_shared(barrier);\n"
|
||||
" asm volatile(\n"
|
||||
' "{\\n"\n'
|
||||
' ".reg .pred p;\\n"\n'
|
||||
' ".reg .b32 remAddr32;\\n"\n'
|
||||
' "setp.ne.s32 p, %2, 0;\\n"\n'
|
||||
' "@p mapa.shared::cluster.u32 remAddr32, %0, %1;\\n"\n'
|
||||
' "@p mbarrier.arrive.shared::cluster.b64 _, [remAddr32];\\n"\n'
|
||||
' "}\\n"\n'
|
||||
' :: "r"(barrier_addr), "r"(cta_id), "r"(pred) : "memory");'
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
# Same cross-CTA arrive, but with an explicit arrival-count operand
|
||||
# (``..., [remAddr32], count``). Matches the ``tma::cluster::arrive`` spelling.
|
||||
device_intrinsic(
|
||||
"_ptx_mbarrier_arrive_remote_count",
|
||||
helper_name="tvm_builtin_ptx_mbarrier_arrive_remote_count",
|
||||
c_signature="(void* barrier, int cta_id, int pred, int count)",
|
||||
body=(
|
||||
" unsigned int barrier_addr = __cvta_generic_to_shared(barrier);\n"
|
||||
" asm volatile(\n"
|
||||
' "{\\n"\n'
|
||||
' ".reg .pred p;\\n"\n'
|
||||
' ".reg .b32 remAddr32;\\n"\n'
|
||||
' "setp.ne.s32 p, %2, 0;\\n"\n'
|
||||
' "@p mapa.shared::cluster.u32 remAddr32, %0, %1;\\n"\n'
|
||||
' "@p mbarrier.arrive.shared::cluster.b64 _, [remAddr32], %3;\\n"\n'
|
||||
' "}\\n"\n'
|
||||
' :: "r"(barrier_addr), "r"(cta_id), "r"(pred), "r"(count) : "memory");'
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@register_codegen("ptx_mbarrier_arrive")
|
||||
def _codegen_mbarrier_arrive(*args):
|
||||
"""Dispatch by arg count: 1 -> local, 3 -> remote, 4 -> remote+count."""
|
||||
if len(args) == 1:
|
||||
result = CODEGEN_REGISTRY["tirx._ptx_mbarrier_arrive_local"](list(args))
|
||||
elif len(args) == 3:
|
||||
result = CODEGEN_REGISTRY["tirx._ptx_mbarrier_arrive_remote"](list(args))
|
||||
elif len(args) == 4:
|
||||
result = CODEGEN_REGISTRY["tirx._ptx_mbarrier_arrive_remote_count"](list(args))
|
||||
else:
|
||||
raise ValueError(f"ptx_mbarrier_arrive expects 1, 3, or 4 args, got {len(args)}")
|
||||
return result[0] if isinstance(result, tuple) else result
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# mbarrier.arrive.expect_tx — local + remote (cluster-mapped) forms.
|
||||
# =============================================================================
|
||||
device_intrinsic(
|
||||
"_ptx_mbarrier_arrive_expect_tx_local",
|
||||
helper_name="tvm_builtin_ptx_mbarrier_arrive_expect_tx",
|
||||
c_signature="(void* barrier, int byte_count)",
|
||||
body=(
|
||||
" unsigned int barrier_addr = __cvta_generic_to_shared(barrier);\n"
|
||||
' asm volatile("mbarrier.arrive.expect_tx.shared.b64 _, [%0], %1;"\n'
|
||||
' :: "r"(barrier_addr), "r"(byte_count) : "memory");'
|
||||
),
|
||||
)
|
||||
device_intrinsic(
|
||||
"_ptx_mbarrier_arrive_expect_tx_remote",
|
||||
helper_name="tvm_builtin_ptx_mbarrier_arrive_expect_tx_remote",
|
||||
c_signature="(void* barrier, int cta_id, int pred, int byte_count)",
|
||||
body=(
|
||||
" unsigned int barrier_addr = __cvta_generic_to_shared(barrier);\n"
|
||||
" asm volatile(\n"
|
||||
' "{\\n"\n'
|
||||
' ".reg .pred p;\\n"\n'
|
||||
' ".reg .b32 remAddr32;\\n"\n'
|
||||
' "setp.ne.s32 p, %2, 0;\\n"\n'
|
||||
' "@p mapa.shared::cluster.u32 remAddr32, %0, %1;\\n"\n'
|
||||
' "@p mbarrier.arrive.expect_tx.shared::cluster.b64 _, [remAddr32], %3;\\n"\n'
|
||||
' "}\\n"\n'
|
||||
' :: "r"(barrier_addr), "r"(cta_id), "r"(pred), "r"(byte_count) : "memory");'
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@register_codegen("ptx_mbarrier_arrive_expect_tx")
|
||||
def _codegen_mbarrier_arrive_expect_tx(*args):
|
||||
"""Dispatch by arg count: 2 -> local, 4 -> remote. Remote arg order from
|
||||
the user is (bar, byte_count, cta_id, pred); reorder to match the helper
|
||||
signature (bar, cta_id, pred, byte_count)."""
|
||||
if len(args) == 2:
|
||||
result = CODEGEN_REGISTRY["tirx._ptx_mbarrier_arrive_expect_tx_local"](list(args))
|
||||
elif len(args) == 4:
|
||||
bar, byte_count, cta_id, pred = args
|
||||
result = CODEGEN_REGISTRY["tirx._ptx_mbarrier_arrive_expect_tx_remote"](
|
||||
[bar, cta_id, pred, byte_count]
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"ptx_mbarrier_arrive_expect_tx expects 2 or 4 args, got {len(args)}")
|
||||
return result[0] if isinstance(result, tuple) else result
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# mbarrier.try_wait.parity.shared::cta.b64 — 1 form. Body wraps the asm in a
|
||||
# label loop (TIRx convention; the magic ``ticks = 0x989680`` is the timeout
|
||||
# hint in ns).
|
||||
# =============================================================================
|
||||
device_intrinsic(
|
||||
"ptx_mbarrier_try_wait",
|
||||
c_signature="(void* barrier, int phase)",
|
||||
body=(
|
||||
" unsigned int barrier_addr_int = __cvta_generic_to_shared(barrier);\n"
|
||||
" unsigned int ticks = 0x989680;\n"
|
||||
" asm volatile(\n"
|
||||
' "{\\n"\n'
|
||||
' ".reg .pred P1;\\n"\n'
|
||||
' "LAB_WAIT:\\n"\n'
|
||||
' "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1, %2;\\n"\n'
|
||||
' "@P1 bra.uni DONE;\\n"\n'
|
||||
' "bra.uni LAB_WAIT;\\n"\n'
|
||||
' "DONE:\\n"\n'
|
||||
' "}\\n"\n'
|
||||
' :: "r"(barrier_addr_int), "r"(phase), "r"(ticks) : "memory");'
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
# mbarrier.try_wait.parity.acquire.cluster — cluster-scope acquire wait used for
|
||||
# cross-CTA barrier handshakes (e.g. the tmem-finished handoff).
|
||||
device_intrinsic(
|
||||
"ptx_mbarrier_try_wait_acquire_cluster",
|
||||
c_signature="(void* barrier, int phase)",
|
||||
body=(
|
||||
" unsigned int barrier_addr_int = __cvta_generic_to_shared(barrier);\n"
|
||||
" asm volatile(\n"
|
||||
' "{\\n"\n'
|
||||
' ".reg .pred P1;\\n"\n'
|
||||
' "LAB_WAIT_AC:\\n"\n'
|
||||
' "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 P1, [%0], %1;\\n"\n'
|
||||
' "@P1 bra.uni DONE_AC;\\n"\n'
|
||||
' "bra.uni LAB_WAIT_AC;\\n"\n'
|
||||
' "DONE_AC:\\n"\n'
|
||||
' "}\\n"\n'
|
||||
' :: "r"(barrier_addr_int), "r"(phase) : "memory");'
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# mbarrier.try_wait.parity — ONE-SHOT non-blocking variant. Returns true
|
||||
# if the requested parity has already been reached, false otherwise.
|
||||
# The TIRx-standard ``ptx_mbarrier_try_wait`` above wraps this in a
|
||||
# label loop that retries until success; this one-shot form is the
|
||||
# building block for bounded-retry debug waits (Nymph's
|
||||
# ``debug_bounded_wait`` lowering mode wraps it in a Python-counted
|
||||
# loop so the kernel cannot hang forever at a mis-protocoled wait).
|
||||
# =============================================================================
|
||||
device_intrinsic(
|
||||
"ptx_mbarrier_try_wait_once",
|
||||
c_signature="(void* barrier, int phase, int ticks)",
|
||||
return_type="uint32_t",
|
||||
body=(
|
||||
" unsigned int barrier_addr_int = __cvta_generic_to_shared(barrier);\n"
|
||||
" unsigned int ticks_u = (unsigned int)ticks;\n"
|
||||
" unsigned int result;\n"
|
||||
" asm volatile(\n"
|
||||
' "{\\n"\n'
|
||||
' ".reg .pred P1;\\n"\n'
|
||||
' "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2, %3;\\n"\n'
|
||||
' "selp.u32 %0, 1, 0, P1;\\n"\n'
|
||||
' "}\\n"\n'
|
||||
' : "=r"(result) : "r"(barrier_addr_int), "r"(phase), "r"(ticks_u) : "memory");\n'
|
||||
" return result;"
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# elect.sync — TIRx uses the CUDA builtin ``tvm_builtin_elect_one_sync()``
|
||||
# helper (declared in the CUDA header tags), not direct PTX.
|
||||
# =============================================================================
|
||||
device_intrinsic(
|
||||
"ptx_elect_sync",
|
||||
helper_name="tvm_builtin_elect_one_sync_op",
|
||||
return_type="uint32_t",
|
||||
body=" return tvm_builtin_elect_one_sync();",
|
||||
extra_deps=("elect_one_sync",),
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# __any_sync — warp-vote (pure CUDA helper).
|
||||
# =============================================================================
|
||||
device_intrinsic(
|
||||
"ptx_any_sync",
|
||||
c_signature="(unsigned mask, int pred)",
|
||||
body=" return __any_sync(mask, pred);",
|
||||
return_type="int",
|
||||
tvm_return_type="int32",
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# CUDA-side sync helpers (zero-arg void unless noted).
|
||||
# =============================================================================
|
||||
device_intrinsic("cuda_thread_fence", body=" __threadfence();")
|
||||
device_intrinsic("cuda_warp_sync", body=" __syncwarp();")
|
||||
device_intrinsic("cuda_cta_sync", body=" __syncthreads();")
|
||||
device_intrinsic(
|
||||
"cuda_grid_sync",
|
||||
body=" namespace cg = cooperative_groups;\n cg::this_grid().sync();",
|
||||
extra_deps=("cooperative_groups",),
|
||||
)
|
||||
device_intrinsic(
|
||||
"cuda_cluster_sync",
|
||||
body=(' asm("barrier.cluster.arrive.aligned;");\n asm("barrier.cluster.wait.aligned;");'),
|
||||
)
|
||||
device_intrinsic(
|
||||
"cuda_warpgroup_sync",
|
||||
c_signature="(int name_bar_id)",
|
||||
body=' asm volatile("bar.sync %0, 128;" : : "r"(name_bar_id));',
|
||||
)
|
||||
device_intrinsic(
|
||||
"cuda_syncthreads_and",
|
||||
c_signature="(int predicate)",
|
||||
body=" return __syncthreads_and(predicate);",
|
||||
return_type="int",
|
||||
tvm_return_type="int32",
|
||||
)
|
||||
device_intrinsic(
|
||||
"cuda_syncthreads_or",
|
||||
c_signature="(int predicate)",
|
||||
body=" return __syncthreads_or(predicate);",
|
||||
return_type="int",
|
||||
tvm_return_type="int32",
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Additional mbarrier, grid-sync, and warp collective helpers.
|
||||
# =============================================================================
|
||||
|
||||
|
||||
# PTX mbarrier parity wait form:
|
||||
# mbarrier.test_wait.parity{.sem.scope}{.shared{::cta}}.b64 waitComplete, [addr], phaseParity;
|
||||
def _mbarrier_test_wait_parity_parts(_barrier, _phase, sem, scope, space):
|
||||
sem = parse_str(sem)
|
||||
scope = parse_str(scope)
|
||||
space = parse_str(space)
|
||||
if sem and sem not in ("acquire", "relaxed"):
|
||||
raise ValueError(f"Unsupported mbarrier.test_wait.parity sem {sem!r}")
|
||||
if scope and scope not in ("cta", "cluster"):
|
||||
raise ValueError(f"Unsupported mbarrier.test_wait.parity scope {scope!r}")
|
||||
if space not in ("shared", "shared::cta"):
|
||||
raise ValueError(f"Unsupported mbarrier.test_wait.parity space {space!r}")
|
||||
sem_scope = f".{sem}.{scope}" if sem else ""
|
||||
name = (
|
||||
"tvm_builtin_ptx_mbarrier_test_wait_parity"
|
||||
f"{('_' + sem + '_' + scope) if sem else ''}_{space.replace('::', '_')}_b64"
|
||||
)
|
||||
body = (
|
||||
" unsigned int ready = 0;\n"
|
||||
" asm volatile(\n"
|
||||
' "{\\n\\t"\n'
|
||||
' ".reg .pred P1; \\n\\t"\n'
|
||||
f' "mbarrier.test_wait.parity{sem_scope}.{space}.b64 P1, [%1], %2; \\n\\t"\n'
|
||||
' "selp.b32 %0, 1, 0, P1; \\n\\t"\n'
|
||||
' "}" : "=r"(ready) : "r"((unsigned int)__cvta_generic_to_shared(barrier)), '
|
||||
'"r"(phase) : "memory");\n'
|
||||
" return ready;"
|
||||
)
|
||||
return name, body
|
||||
|
||||
|
||||
device_intrinsic(
|
||||
"ptx_mbarrier_test_wait_parity",
|
||||
n_attrs=3,
|
||||
helper_name=lambda *a: _mbarrier_test_wait_parity_parts(*a)[0],
|
||||
c_signature="(void* barrier, int phase)",
|
||||
return_type="unsigned int",
|
||||
tvm_return_type="uint32",
|
||||
body=lambda *a: _mbarrier_test_wait_parity_parts(*a)[1],
|
||||
)
|
||||
|
||||
device_intrinsic(
|
||||
"cuda_ballot_sync",
|
||||
helper_name="tvm_builtin_ballot_sync",
|
||||
c_signature="(unsigned int mask, int pred)",
|
||||
return_type="unsigned int",
|
||||
body=" return __ballot_sync(mask, pred);",
|
||||
)
|
||||
device_intrinsic(
|
||||
"cuda_reduce_add_sync_u32",
|
||||
helper_name="tvm_builtin_reduce_add_sync_u32",
|
||||
c_signature="(unsigned int mask, unsigned int value)",
|
||||
return_type="unsigned int",
|
||||
body=" return __reduce_add_sync(mask, value);",
|
||||
)
|
||||
device_intrinsic(
|
||||
"cuda_reduce_min_sync_u32",
|
||||
helper_name="tvm_builtin_reduce_min_sync_u32",
|
||||
c_signature="(unsigned int mask, unsigned int value)",
|
||||
return_type="unsigned int",
|
||||
body=" return __reduce_min_sync(mask, value);",
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# griddepcontrol.wait / griddepcontrol.launch_dependents (sm_90+)
|
||||
# Programmatic Dependent Launch (PDL) synchronization. Both carry memory
|
||||
# clobber to prevent CSE / cross-barrier reordering.
|
||||
# =============================================================================
|
||||
device_intrinsic(
|
||||
"ptx_griddepcontrol_wait",
|
||||
body=' asm volatile("griddepcontrol.wait;" ::: "memory");',
|
||||
)
|
||||
|
||||
device_intrinsic(
|
||||
"ptx_griddepcontrol_launch_dependents",
|
||||
body=' asm volatile("griddepcontrol.launch_dependents;" ::: "memory");',
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,71 @@
|
||||
# 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.
|
||||
"""PTX data types for CUDA codegen."""
|
||||
|
||||
import enum
|
||||
|
||||
import tvm_ffi
|
||||
|
||||
from_string_func = tvm_ffi.get_global_func("tirx.intrinsics.cuda.PTXDTypeFromString")
|
||||
to_string_func = tvm_ffi.get_global_func("tirx.intrinsics.cuda.PTXDTypeToString")
|
||||
|
||||
|
||||
class PTXDataType(enum.Enum):
|
||||
"""
|
||||
A Python equivalent of the provided C++ DataType enum class.
|
||||
|
||||
Inherits from IntEnum so that members behave both as enum members
|
||||
and as integers, mirroring the C++ behavior.
|
||||
|
||||
see also src/target/source/ptx.cc
|
||||
"""
|
||||
|
||||
INT4 = 0
|
||||
UINT4 = 1
|
||||
INT8 = 2
|
||||
UINT8 = 3
|
||||
INT16 = 4
|
||||
UINT16 = 5
|
||||
INT32 = 6
|
||||
UINT32 = 7
|
||||
INT64 = 8
|
||||
UINT64 = 9
|
||||
FLOAT4_E2M1FN = 10
|
||||
FLOAT6_E2M3FN = 11
|
||||
FLOAT6_E3M2FN = 12
|
||||
FLOAT8_E4M3FN = 13
|
||||
FLOAT8_E4M3FNUZ = 14
|
||||
FLOAT8_E5M2 = 15
|
||||
FLOAT8_E8M0FNU = 16
|
||||
FLOAT16 = 17
|
||||
BFLOAT16 = 18
|
||||
FLOAT16X2 = 19
|
||||
FLOAT32 = 20
|
||||
TENSOR_FLOAT32 = 21
|
||||
FLOAT64 = 22
|
||||
BIT1 = 23
|
||||
BIT8 = 24
|
||||
BIT16 = 25
|
||||
BIT32 = 26
|
||||
BIT64 = 27
|
||||
|
||||
@classmethod
|
||||
def from_string(cls, s_type: str) -> "PTXDataType":
|
||||
return PTXDataType(from_string_func(s_type))
|
||||
|
||||
def to_string(self) -> str:
|
||||
return to_string_func(self.value)
|
||||
@@ -0,0 +1,82 @@
|
||||
# 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.
|
||||
"""Common utility functions for CUDA op codegen."""
|
||||
|
||||
|
||||
def parse_str(arg) -> str:
|
||||
"""Parse TIR StringImm or Python str to a plain str.
|
||||
|
||||
TIR StringImm values stringify to quoted strings, e.g., ``'"float16"'``;
|
||||
Python strs do not. Idempotent — passing an already-parsed str returns it
|
||||
unchanged, so dispatchers that parse once before forwarding to inner
|
||||
codegens won't double-strip the value.
|
||||
"""
|
||||
s = str(arg)
|
||||
if len(s) >= 2 and s[0] == '"' and s[-1] == '"':
|
||||
return s[1:-1]
|
||||
return s
|
||||
|
||||
|
||||
def is_power_of_two(n: int) -> bool:
|
||||
"""Check if n is a power of two."""
|
||||
return n > 0 and (n & (n - 1)) == 0
|
||||
|
||||
|
||||
def validate_cta_group(cta_group, context: str = "") -> int:
|
||||
"""Validate that cta_group is 1 or 2 and return it as int.
|
||||
|
||||
Args:
|
||||
cta_group: The cta_group value (can be int or TIR IntImm)
|
||||
context: Optional context string for error message (e.g., "allocating Tensor Memory")
|
||||
|
||||
Returns:
|
||||
The validated cta_group as int
|
||||
|
||||
Raises:
|
||||
ValueError: If cta_group is not 1 or 2
|
||||
"""
|
||||
cta_group = int(cta_group)
|
||||
if cta_group not in [1, 2]:
|
||||
ctx = f" involved in {context}" if context else ""
|
||||
raise ValueError(
|
||||
f"The number of cta_group{ctx} is incorrect, expected 1 or 2, got {cta_group}"
|
||||
)
|
||||
return cta_group
|
||||
|
||||
|
||||
def validate_power_of_two_range(value, min_val: int, max_val: int, name: str) -> int:
|
||||
"""Validate that value is within range and is a power of two.
|
||||
|
||||
Args:
|
||||
value: The value to validate
|
||||
min_val: Minimum allowed value (inclusive)
|
||||
max_val: Maximum allowed value (inclusive)
|
||||
name: Name of the parameter for error messages
|
||||
|
||||
Returns:
|
||||
The validated value as int
|
||||
|
||||
Raises:
|
||||
ValueError: If value is out of range or not a power of two
|
||||
"""
|
||||
value = int(value)
|
||||
if not (min_val <= value <= max_val and is_power_of_two(value)):
|
||||
raise ValueError(
|
||||
f"The {name} is invalid, expect a value within range [{min_val}, {max_val}] "
|
||||
f"and be a power of 2, got {value}"
|
||||
)
|
||||
return value
|
||||
@@ -0,0 +1,403 @@
|
||||
# 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, too-many-arguments, too-many-locals, too-many-positional-arguments
|
||||
"""PTX WGMMA operations (Hopper warpgroup MMA).
|
||||
|
||||
One ``device_intrinsic`` registration per PTX form table entry. Bodies are
|
||||
hand-written ``asm volatile(...)`` strings. Variable-arity register vectors
|
||||
(``mma_async`` accumulators / A fragments) materialize via the same
|
||||
device_intrinsic with an attr-driven ``c_signature`` callable.
|
||||
"""
|
||||
|
||||
import tvm
|
||||
|
||||
from ._schema import device_intrinsic
|
||||
from .registry import CODEGEN_REGISTRY, register_codegen
|
||||
from .types import PTXDataType
|
||||
from .utils import parse_str
|
||||
|
||||
# =============================================================================
|
||||
# wgmma.fence / commit_group / wait_group — one PTX form each.
|
||||
# =============================================================================
|
||||
device_intrinsic(
|
||||
"ptx_wgmma_fence",
|
||||
helper_name="ptx_wgmma_fence",
|
||||
body=' asm volatile("wgmma.fence.sync.aligned;" ::: "memory");',
|
||||
)
|
||||
device_intrinsic(
|
||||
"ptx_wgmma_commit_group",
|
||||
helper_name="ptx_wgmma_commit_group",
|
||||
body=' asm volatile("wgmma.commit_group.sync.aligned;" ::: "memory");',
|
||||
)
|
||||
device_intrinsic(
|
||||
"ptx_wgmma_wait_group",
|
||||
n_attrs=1,
|
||||
helper_name=lambda n: f"ptx_wgmma_wait_group_{int(n)}",
|
||||
body=lambda n: f' asm volatile("wgmma.wait_group.sync.aligned {int(n)};" ::: "memory");',
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# wgmma_encode_matrix_descriptor — pure-C bitfield struct fill (no asm).
|
||||
# =============================================================================
|
||||
device_intrinsic(
|
||||
"ptx_wgmma_encode_matrix_descriptor",
|
||||
helper_name="ptx_wgmma_encode_matrix_descriptor",
|
||||
c_signature="(uint64_t* desc, void* addr, int ldo, int sdo, int swizzle)",
|
||||
body=(
|
||||
" GmmaDescriptor _desc{}; // value-init: reading uncovered pad bits is UB\n"
|
||||
"\n"
|
||||
" switch (swizzle) {\n"
|
||||
" case 0: _desc.bitfield.layout_type_ = uint8_t(0); break; // No swizzle\n"
|
||||
" case 1: _desc.bitfield.layout_type_ = uint8_t(3); break; // 32B swizzle\n"
|
||||
" case 2: _desc.bitfield.layout_type_ = uint8_t(2); break; // 64B swizzle\n"
|
||||
" case 3: _desc.bitfield.layout_type_ = uint8_t(1); break; // 128B swizzle\n"
|
||||
" }\n"
|
||||
"\n"
|
||||
" uint32_t start_address = __cvta_generic_to_shared(addr);\n"
|
||||
" _desc.bitfield.start_address_ = static_cast<uint16_t>(start_address >> 4);\n"
|
||||
"\n"
|
||||
" constexpr uint8_t base_offset = 0;\n"
|
||||
" _desc.bitfield.base_offset_ = base_offset;\n"
|
||||
"\n"
|
||||
" _desc.bitfield.stride_byte_offset_ = static_cast<uint32_t>(sdo);\n"
|
||||
" _desc.bitfield.leading_byte_offset_ = static_cast<uint32_t>(ldo);\n"
|
||||
"\n"
|
||||
" *desc = (uint64_t)_desc;"
|
||||
),
|
||||
extra_deps=("gmma_descriptor",),
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# wgmma_noop_barrier — empty asm with one inout register operand. Two
|
||||
# device_intrinsic calls, one per supported dtype; dispatcher picks the form
|
||||
# based on the operand's runtime dtype.
|
||||
# =============================================================================
|
||||
device_intrinsic(
|
||||
"ptx_wgmma_noop_barrier_uint32",
|
||||
helper_name="ptx_wgmma_fence_uint32_t",
|
||||
c_signature="(uint32_t reg)",
|
||||
body=' asm volatile("" : "+r"(reg) :: "memory");',
|
||||
)
|
||||
device_intrinsic(
|
||||
"ptx_wgmma_noop_barrier_float32",
|
||||
helper_name="ptx_wgmma_fence_float",
|
||||
c_signature="(float reg)",
|
||||
body=' asm volatile("" : "+f"(reg) :: "memory");',
|
||||
)
|
||||
|
||||
|
||||
@register_codegen("ptx_wgmma_noop_barrier")
|
||||
def codegen_ptx_wgmma_noop_barrier(reg):
|
||||
dtype = str(reg.dtype)
|
||||
dtype_enum = PTXDataType.from_string(dtype)
|
||||
if dtype_enum == PTXDataType.UINT32:
|
||||
op_name = "tirx.ptx_wgmma_noop_barrier_uint32"
|
||||
elif dtype_enum == PTXDataType.FLOAT32:
|
||||
op_name = "tirx.ptx_wgmma_noop_barrier_float32"
|
||||
else:
|
||||
raise ValueError(f"Only support uint32/float32 for wgmma_fence, but got {dtype}.")
|
||||
result = CODEGEN_REGISTRY[op_name]([reg])
|
||||
return result[0] if isinstance(result, tuple) else result
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# wgmma.mma_async ss / rs — 2 PTX form table entries. Accumulator count and
|
||||
# A-register count vary with (M, N, K, in_dtype) but are fully determined by
|
||||
# attrs at codegen time.
|
||||
#
|
||||
# Args layout for ss form (forwarded operand args first, then 9 attr args):
|
||||
# *p_acc[0..num_accums-1], p_descA, p_descB, p_scaleD,
|
||||
# M, N, K, in_dtype, out_dtype, transA, transB, scaleA, scaleB
|
||||
#
|
||||
# Args layout for rs form:
|
||||
# *p_acc[0..num_accums-1], *p_A[0..num_A_regs-1], p_descB, p_scaleD,
|
||||
# M, N, K, in_dtype, out_dtype, transA, transB, scaleA, scaleB
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def _coerce_wgmma_attrs(attrs):
|
||||
"""Decode the trailing 9 attrs (M, N, K, in_dtype, out_dtype, transA,
|
||||
transB, scaleA, scaleB) into native Python types."""
|
||||
M, N, K = int(attrs[0]), int(attrs[1]), int(attrs[2])
|
||||
in_dtype = parse_str(attrs[3])
|
||||
out_dtype = parse_str(attrs[4])
|
||||
transA = bool(int(attrs[5])) if hasattr(attrs[5], "value") else bool(attrs[5])
|
||||
transB = bool(int(attrs[6])) if hasattr(attrs[6], "value") else bool(attrs[6])
|
||||
scaleA = bool(int(float(attrs[7])))
|
||||
scaleB = bool(int(float(attrs[8])))
|
||||
if out_dtype != "float32":
|
||||
raise ValueError("WGMMA codegen only supports float32 as output dtype.")
|
||||
allow_transpose = in_dtype in {"float16", "bfloat16"}
|
||||
if not allow_transpose and (transA or transB):
|
||||
raise ValueError("Transpose is only supported for .f16/.bf16 types in WGMMA.")
|
||||
return M, N, K, in_dtype, out_dtype, transA, transB, scaleA, scaleB, allow_transpose
|
||||
|
||||
|
||||
def _safe(s):
|
||||
return s.replace("::", "_").replace(".", "_")
|
||||
|
||||
|
||||
def _wgmma_helper_name(prefix, M, N, K, in_dtype, out_dtype, transA, transB, scaleA, scaleB):
|
||||
return (
|
||||
f"{prefix}_{M}x{N}x{K}_{_safe(in_dtype)}_{_safe(out_dtype)}"
|
||||
f"_{1 if scaleA else 0}_{1 if scaleB else 0}"
|
||||
f"_{1 if transA else 0}_{1 if transB else 0}"
|
||||
)
|
||||
|
||||
|
||||
def _wgmma_in_bits(in_dtype):
|
||||
return tvm.runtime.DataType(in_dtype).bits
|
||||
|
||||
|
||||
def _wgmma_ss_parts(*args):
|
||||
M, N, K, in_dtype, out_dtype, transA, transB, scaleA, scaleB, allow_transpose = (
|
||||
_coerce_wgmma_attrs(args[-9:])
|
||||
)
|
||||
num_accums = M * N // 128
|
||||
|
||||
name = _wgmma_helper_name(
|
||||
"ptx_wgmma_mma_async_ss", M, N, K, in_dtype, out_dtype, transA, transB, scaleA, scaleB
|
||||
)
|
||||
sig = (
|
||||
"("
|
||||
+ ", ".join(
|
||||
[f"float& p_acc{i}" for i in range(num_accums)]
|
||||
+ ["uint64_t p_descA", "uint64_t p_descB", "int p_scaleD"]
|
||||
)
|
||||
+ ")"
|
||||
)
|
||||
descA_idx = num_accums
|
||||
descB_idx = num_accums + 1
|
||||
scaleD_idx = num_accums + 2
|
||||
scaleA_idx = num_accums + 3
|
||||
scaleB_idx = num_accums + 4
|
||||
transA_idx = num_accums + 5
|
||||
transB_idx = num_accums + 6
|
||||
accum_r_list = ", ".join(f"%{i}" for i in range(num_accums))
|
||||
accum_constraints = ", ".join(f'"+f"(p_acc{i})' for i in range(num_accums))
|
||||
itype = PTXDataType.from_string(in_dtype)
|
||||
otype = PTXDataType.from_string(out_dtype)
|
||||
if allow_transpose:
|
||||
transpose_r_code = f", %{transA_idx}, %{transB_idx}"
|
||||
transpose_constraints = f', "n"({1 if transA else 0}), "n"({1 if transB else 0})'
|
||||
else:
|
||||
transpose_r_code = ""
|
||||
transpose_constraints = ""
|
||||
instr = (
|
||||
f"wgmma.mma_async.sync.aligned.m{M}n{N}k{K}"
|
||||
f"{otype.to_string()}{itype.to_string()}{itype.to_string()}"
|
||||
)
|
||||
asm_inputs = (
|
||||
f'"l"(p_descA), "l"(p_descB), "r"(p_scaleD),'
|
||||
f' "n"({1 if scaleA else 0}), "n"({1 if scaleB else 0})'
|
||||
f"{transpose_constraints}"
|
||||
)
|
||||
body = (
|
||||
" asm volatile(\n"
|
||||
' "{ \\n"\n'
|
||||
' ".reg .pred p;\\n"\n'
|
||||
f' "setp.ne.b32 p, %{scaleD_idx}, 0;\\n"\n'
|
||||
f' "{instr} "\n'
|
||||
f' "{{{accum_r_list}}},"\n'
|
||||
f' "%{descA_idx}, %{descB_idx},"\n'
|
||||
f' "p, %{scaleA_idx}, %{scaleB_idx}{transpose_r_code};\\n"\n'
|
||||
' "}\\n"\n'
|
||||
f" : {accum_constraints}\n"
|
||||
f" : {asm_inputs}\n"
|
||||
" );"
|
||||
)
|
||||
return name, sig, body
|
||||
|
||||
|
||||
device_intrinsic(
|
||||
"_ptx_wgmma_mma_async_ss_impl",
|
||||
n_attrs=9,
|
||||
helper_name=lambda *a: _wgmma_ss_parts(*a)[0],
|
||||
c_signature=lambda *a: _wgmma_ss_parts(*a)[1],
|
||||
body=lambda *a: _wgmma_ss_parts(*a)[2],
|
||||
)
|
||||
|
||||
|
||||
def _wgmma_rs_parts(*args):
|
||||
M, N, K, in_dtype, out_dtype, transA, transB, scaleA, scaleB, allow_transpose = (
|
||||
_coerce_wgmma_attrs(args[-9:])
|
||||
)
|
||||
num_accums = M * N // 128
|
||||
in_bits = _wgmma_in_bits(in_dtype)
|
||||
num_A_regs = M * K // 128 // (32 // in_bits)
|
||||
|
||||
name = _wgmma_helper_name(
|
||||
"ptx_wgmma_mma_async_rs", M, N, K, in_dtype, out_dtype, transA, transB, scaleA, scaleB
|
||||
)
|
||||
sig = (
|
||||
"("
|
||||
+ ", ".join(
|
||||
[f"float& p_acc{i}" for i in range(num_accums)]
|
||||
+ [f"uint32_t& p_A{i}" for i in range(num_A_regs)]
|
||||
+ ["uint64_t p_descB", "int p_scaleD"]
|
||||
)
|
||||
+ ")"
|
||||
)
|
||||
|
||||
accum_r_list = ", ".join(f"%{i}" for i in range(num_accums))
|
||||
A_reg_r_list = ", ".join(f"%{num_accums + i}" for i in range(num_A_regs))
|
||||
base_idx = num_accums + num_A_regs
|
||||
descB_idx = base_idx
|
||||
scaleD_idx = base_idx + 1
|
||||
scaleA_idx = base_idx + 2
|
||||
scaleB_idx = base_idx + 3
|
||||
transB_idx = base_idx + 4
|
||||
accum_constraints = ", ".join(f'"+f"(p_acc{i})' for i in range(num_accums))
|
||||
A_reg_constraints = ", ".join(f'"r"(p_A{i})' for i in range(num_A_regs))
|
||||
itype = PTXDataType.from_string(in_dtype)
|
||||
otype = PTXDataType.from_string(out_dtype)
|
||||
if allow_transpose:
|
||||
transpose_r_code = f", %{transB_idx}"
|
||||
transpose_constraints = f', "n"({1 if transB else 0})'
|
||||
else:
|
||||
transpose_r_code, transpose_constraints = "", ""
|
||||
instr = (
|
||||
f"wgmma.mma_async.sync.aligned.m{M}n{N}k{K}"
|
||||
f"{otype.to_string()}{itype.to_string()}{itype.to_string()}"
|
||||
)
|
||||
asm_inputs = (
|
||||
f'{A_reg_constraints}, "l"(p_descB), "r"(p_scaleD),'
|
||||
f' "n"({1 if scaleA else 0}), "n"({1 if scaleB else 0})'
|
||||
f"{transpose_constraints}"
|
||||
)
|
||||
body = (
|
||||
" asm volatile(\n"
|
||||
' "{ \\n"\n'
|
||||
' ".reg .pred p;\\n"\n'
|
||||
f' "setp.ne.b32 p, %{scaleD_idx}, 0;\\n"\n'
|
||||
f' "{instr} "\n'
|
||||
f' "{{{accum_r_list}}},"\n'
|
||||
f' "{{{A_reg_r_list}}}, %{descB_idx},"\n'
|
||||
f' "p, %{scaleA_idx}, %{scaleB_idx}{transpose_r_code};\\n"\n'
|
||||
' "}\\n"\n'
|
||||
f" : {accum_constraints}\n"
|
||||
f" : {asm_inputs}\n"
|
||||
" );"
|
||||
)
|
||||
return name, sig, body
|
||||
|
||||
|
||||
device_intrinsic(
|
||||
"_ptx_wgmma_mma_async_rs_impl",
|
||||
n_attrs=9,
|
||||
helper_name=lambda *a: _wgmma_rs_parts(*a)[0],
|
||||
c_signature=lambda *a: _wgmma_rs_parts(*a)[1],
|
||||
body=lambda *a: _wgmma_rs_parts(*a)[2],
|
||||
)
|
||||
|
||||
|
||||
# User-facing wrappers: just normalise types + reorder positional args to
|
||||
# put operands first, then attrs, matching the schema convention.
|
||||
|
||||
|
||||
def _wgmma_user_wrapper_ss(*args):
|
||||
M, N, K, in_dtype, out_dtype, transA, transB, scaleA, scaleB, scaleD, descA, descB, *accums = (
|
||||
args
|
||||
)
|
||||
M = int(M)
|
||||
N = int(N)
|
||||
K = int(K)
|
||||
in_dtype = parse_str(in_dtype)
|
||||
out_dtype = parse_str(out_dtype)
|
||||
transA = bool(transA)
|
||||
transB = bool(transB)
|
||||
scaleA = bool(int(float(scaleA)))
|
||||
scaleB = bool(int(float(scaleB)))
|
||||
expected = M * N // 128
|
||||
if len(accums) != expected:
|
||||
raise ValueError(
|
||||
"The number of arguments is incorrect. Expected "
|
||||
f"{12 + expected} total args (meaning {expected} accumulator args), "
|
||||
f"but got {len(accums)}."
|
||||
)
|
||||
return [
|
||||
*accums,
|
||||
descA,
|
||||
descB,
|
||||
scaleD,
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
transA,
|
||||
transB,
|
||||
scaleA,
|
||||
scaleB,
|
||||
]
|
||||
|
||||
|
||||
@register_codegen("ptx_wgmma_mma_async_ss")
|
||||
def codegen_ptx_wgmma_mma_async_ss(*args):
|
||||
forwarded = _wgmma_user_wrapper_ss(*args)
|
||||
result = CODEGEN_REGISTRY["tirx._ptx_wgmma_mma_async_ss_impl"](forwarded)
|
||||
return result[0] if isinstance(result, tuple) else result
|
||||
|
||||
|
||||
def _wgmma_user_wrapper_rs(*args):
|
||||
M, N, K, in_dtype, out_dtype, transA, transB, scaleA, scaleB, scaleD, descB, *reg_list = args
|
||||
M = int(M)
|
||||
N = int(N)
|
||||
K = int(K)
|
||||
in_dtype = parse_str(in_dtype)
|
||||
out_dtype = parse_str(out_dtype)
|
||||
transA = bool(transA)
|
||||
transB = bool(transB)
|
||||
scaleA = bool(int(float(scaleA)))
|
||||
scaleB = bool(int(float(scaleB)))
|
||||
if out_dtype != "float32":
|
||||
raise ValueError("This generator only supports float32 as the output dtype for WGMMA.")
|
||||
in_dtype_bits = tvm.runtime.DataType(in_dtype).bits
|
||||
if in_dtype_bits is None:
|
||||
raise ValueError(f"Bit width not defined for input dtype: {in_dtype}")
|
||||
expected_A_cnt = M * K // 128 // (32 // in_dtype_bits)
|
||||
expected_accm_cnt = M * N // 128
|
||||
if len(reg_list) != expected_A_cnt + expected_accm_cnt:
|
||||
raise ValueError(
|
||||
f"Incorrect number of A registers. Expected {expected_A_cnt}, got {len(reg_list)}"
|
||||
)
|
||||
A_regs = reg_list[:expected_A_cnt]
|
||||
accums = reg_list[expected_A_cnt:]
|
||||
return [
|
||||
*accums,
|
||||
*A_regs,
|
||||
descB,
|
||||
scaleD,
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
transA,
|
||||
transB,
|
||||
scaleA,
|
||||
scaleB,
|
||||
]
|
||||
|
||||
|
||||
@register_codegen("ptx_wgmma_mma_async_rs")
|
||||
def codegen_ptx_wgmma_mma_async_rs(*args):
|
||||
forwarded = _wgmma_user_wrapper_rs(*args)
|
||||
result = CODEGEN_REGISTRY["tirx._ptx_wgmma_mma_async_rs_impl"](forwarded)
|
||||
return result[0] if isinstance(result, tuple) else result
|
||||
Reference in New Issue
Block a user