tvm_tirx.backend.cuda.codegen.device_intrinsic

tvm_tirx.backend.cuda.codegen.device_intrinsic#

tvm_tirx.backend.cuda.codegen.device_intrinsic(op_name, *, helper_name=None, c_signature='()', body, n_attrs=0, return_type='void', tvm_return_type=None, templated=False, extra_deps=())#

Register a CUDA device-helper intrinsic.

Parameters:

  • op_name (str) – 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 (str | Callable | None) – Literal C function name, OR (*args) -> str to compute it from attr values. Defaults to f"tvm_builtin_{op_name}".

  • c_signature (str | Callable) – Literal C parameter list including outer parens ("(int x, int y)"), OR (*args) -> str to compute it from attr values. Defaults to "()".

  • body (str | Callable) – Literal C body string (already indented), OR (*args) -> str.

  • n_attrs (int) – 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 (str | Callable) – C return type. Default "void". Either a literal string or (*args) -> str when the helper return type depends on attrs.

  • tvm_return_type (str | Callable | None) – 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 (bool) – Prefix the helper with template <typename T>.

  • extra_deps (tuple) – 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.

Return type:

None