chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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();",
|
||||
)
|
||||
Reference in New Issue
Block a user