wehub-resource-sync
254 lines
8.6 KiB
Python
254 lines
8.6 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
# 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();",
|
|
)
|