# 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();", )