This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new ea0cfa320f [TIRx][CUDA] Preserve kernel launch calls and add CUDA host
source bundling (#20395)
ea0cfa320f is described below
commit ea0cfa320fcdea0567e2f4c19d363625c66630d1
Author: Tianqi Chen <[email protected]>
AuthorDate: Sun Sep 20 22:10:32 2026 -0400
[TIRx][CUDA] Preserve kernel launch calls and add CUDA host source bundling
(#20395)
Preserve kernel launch metadata through `tirx.call_ffi_kernel`. Ordinary
host targets lower these calls through packed functions, while
`cuda_host` emits direct CUDA launches with explicit diagnostics for
unsupported launch modes.
Add `tvm.backend.cuda.export_cuda_host(mod)` to concatenate retained
CUDA device source before its host wrappers into one translation unit.
Existing CUDA device code generation stays unchanged; source generation,
compilation and loading remain separate steps.
CUDA-host wrappers use tvm-ffi and CUDA libraries for streams, scoped
device selection, tensor-map encoding and error handling. Tensor-map
descriptors use aligned host storage and the device's by-value launch
ABI. Unsupported TVM runtime services are diagnosed during code
generation.
Tensor-map initialization uses `tirx.tensormap_encode_tiled` to retain
the final descriptor dtype and fixed options in call attributes, with
pointers and dimensions as runtime operands. CUDA-host codegen emits
typed arrays and a direct driver call; ordinary hosts keep the existing
packed encoder ABI.
---
include/tvm/tirx/attrs.h | 73 ++++
include/tvm/tirx/builtin.h | 18 +
python/tvm/backend/cuda/__init__.py | 3 +
python/tvm/backend/cuda/host.py | 78 ++++
.../backend/cuda/tile_primitive/copy_async/tma.py | 17 +-
python/tvm/tirx/__init__.py | 1 +
python/tvm/tirx/build.py | 2 +-
python/tvm/tirx/op.py | 84 +++++
python/tvm/tirx/script/builder/ir.py | 4 +
src/backend/cuda/codegen/codegen_cuda_host.cc | 401 +++++++++++++++++++++
src/backend/cuda/codegen/target_kind.cc | 2 +
src/tirx/op/builtin.cc | 32 ++
src/tirx/script/printer/expr.cc | 27 ++
src/tirx/transform/lower_tirx_dedup_tensormap.cc | 50 ++-
src/tirx/transform/lower_tvm_builtin.cc | 35 +-
src/tirx/transform/split_host_device.cc | 6 +-
tests/python/codegen/test_target_codegen_cuda.py | 41 +++
.../test_tir_transform_lower_tvm_builtin.py | 10 +-
.../test_tir_transform_split_host_device.py | 23 +-
.../tile_primitive/cuda/copy_async/test_tma.py | 32 +-
20 files changed, 868 insertions(+), 71 deletions(-)
diff --git a/include/tvm/tirx/attrs.h b/include/tvm/tirx/attrs.h
new file mode 100644
index 0000000000..1d173b3561
--- /dev/null
+++ b/include/tvm/tirx/attrs.h
@@ -0,0 +1,73 @@
+/*
+ * 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.
+ */
+#ifndef TVM_TIRX_ATTRS_H_
+#define TVM_TIRX_ATTRS_H_
+
+#include <tvm/ffi/container/array.h>
+#include <tvm/ir/attrs.h>
+
+namespace tvm {
+namespace tirx {
+
+/*! \brief Launch metadata for call_ffi_kernel. */
+struct CallFFIKernelAttr : public AttrsNode {
+ /*!
+ * \brief Ordered launch tags describing the suffix of the call arguments.
+ *
+ * The first call argument is the kernel symbol, followed by kernel operands
+ * and launch values. Flag-only tags consume no value; dynamic shared-memory
+ * bytes, when present, are last. Runtime expressions remain in Call.args.
+ */
+ ffi::Array<ffi::String> launch_params;
+
+ static void RegisterReflection() {
+ ffi::reflection::ObjectDef<CallFFIKernelAttr>().def_ro("launch_params",
+
&CallFFIKernelAttr::launch_params);
+ }
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.CallFFIKernelAttr",
CallFFIKernelAttr, AttrsNode);
+};
+
+/*! \brief Fixed encoding options for tensormap_encode_tiled. */
+struct TensorMapEncodeTiledAttr : public AttrsNode {
+ DLDataType descriptor_dtype;
+ int64_t rank;
+ int64_t interleave;
+ int64_t swizzle;
+ int64_t l2_promotion;
+ int64_t oob_fill;
+ int64_t force_cu_dtype;
+
+ static void RegisterReflection() {
+ ffi::reflection::ObjectDef<TensorMapEncodeTiledAttr>()
+ .def_ro("descriptor_dtype",
&TensorMapEncodeTiledAttr::descriptor_dtype)
+ .def_ro("rank", &TensorMapEncodeTiledAttr::rank)
+ .def_ro("interleave", &TensorMapEncodeTiledAttr::interleave)
+ .def_ro("swizzle", &TensorMapEncodeTiledAttr::swizzle)
+ .def_ro("l2_promotion", &TensorMapEncodeTiledAttr::l2_promotion)
+ .def_ro("oob_fill", &TensorMapEncodeTiledAttr::oob_fill)
+ .def_ro("force_cu_dtype", &TensorMapEncodeTiledAttr::force_cu_dtype);
+ }
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.TensorMapEncodeTiledAttr",
TensorMapEncodeTiledAttr,
+ AttrsNode);
+};
+
+} // namespace tirx
+} // namespace tvm
+
+#endif // TVM_TIRX_ATTRS_H_
diff --git a/include/tvm/tirx/builtin.h b/include/tvm/tirx/builtin.h
index 6b8939389c..4f66ab328a 100644
--- a/include/tvm/tirx/builtin.h
+++ b/include/tvm/tirx/builtin.h
@@ -349,6 +349,24 @@ TVM_DLL const Op& tvm_stack_make_array();
*/
TVM_DLL const Op& tvm_call_packed();
+/*!
+ * \brief Launch a kernel using the packed-function argument convention.
+ *
+ * Arguments are the kernel symbol, kernel operands, then launch values.
+ * CallFFIKernelAttr::launch_params describes the launch-value suffix.
+ * Host backends may consume this call directly or lower it to tvm_call_packed.
+ */
+TVM_DLL const Op& call_ffi_kernel();
+
+/*!
+ * \brief Encode a tiled tensor map at invocation time.
+ *
+ * TensorMapEncodeTiledAttr stores the descriptor dtype, rank and fixed
options.
+ * Arguments are descriptor and data pointers, global dimensions (rank), byte
+ * strides (rank - 1), box dimensions (rank), then element strides (rank).
+ */
+TVM_DLL const Op& tensormap_encode_tiled();
+
/*!
* \brief See pesudo code
*
diff --git a/python/tvm/backend/cuda/__init__.py
b/python/tvm/backend/cuda/__init__.py
index 89f0b9b71e..1b334c6f45 100644
--- a/python/tvm/backend/cuda/__init__.py
+++ b/python/tvm/backend/cuda/__init__.py
@@ -23,6 +23,8 @@ from tvm_ffi.libinfo import load_lib_ctypes
from tvm.base import _LOADED_LIBS
+from .host import export_cuda_host
+
_LAZY_SUBMODULES = {
"codegen",
"cpp",
@@ -112,6 +114,7 @@ def __getattr__(name: str):
__all__ = [
"codegen",
"cpp",
+ "export_cuda_host",
"iket",
"lang",
"op",
diff --git a/python/tvm/backend/cuda/host.py b/python/tvm/backend/cuda/host.py
new file mode 100644
index 0000000000..65e31dc16f
--- /dev/null
+++ b/python/tvm/backend/cuda/host.py
@@ -0,0 +1,78 @@
+# 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.
+"""CUDA host and device source assembly."""
+
+from tvm_ffi import Module
+
+
+def export_cuda_host(mod: Module) -> str:
+ """Return one CUDA C++ translation unit from a built CUDA-host module.
+
+ Build ``mod`` with a CUDA target whose host is ``"cuda_host"``. Its CUDA
+ device imports must retain their CUDA C++ source. Device definitions are
+ emitted before the host wrappers that launch them. This function only
+ reads the modules and returns source; compilation, loading and writing
+ files are left to the caller. Sources are concatenated without rewriting
+ declarations or deduplicating helpers, so the imports must be compatible
+ within one translation unit and provide the types needed by NVCC's host
pass.
+ Generated host wrappers use tvm-ffi and CUDA libraries. TVM runtime
workspace
+ allocation and other unsupported runtime services are rejected during
codegen.
+ Tensor-map encoding additionally requires linking the CUDA driver library.
+
+ Parameters
+ ----------
+ mod : tvm.runtime.Module
+ The built CUDA-host source module with its CUDA device imports.
+
+ Returns
+ -------
+ source : str
+ CUDA C++ source suitable for compilation with NVCC and tvm-ffi headers.
+
+ Raises
+ ------
+ ValueError
+ If the host or its imports are incompatible, or CUDA source is absent
+ (for example, after loading a device module saved as a binary).
+ """
+ if not isinstance(mod, Module):
+ raise TypeError("export_cuda_host expects a runtime Module")
+ if mod.kind != "c" or "cu" not in mod.get_write_formats():
+ raise ValueError("Expected a source module built with
host='cuda_host'")
+ host_source = mod.inspect_source()
+ if not host_source:
+ raise ValueError("CUDA-host module has no source")
+
+ sources = []
+ visited = set()
+
+ def collect(device_mod):
+ if device_mod in visited:
+ return
+ visited.add(device_mod)
+ if device_mod.kind != "cuda":
+ raise ValueError(f"Expected a CUDA device import, got
{device_mod.kind!r}")
+ source = device_mod.inspect_source("cuda")
+ if not source:
+ raise ValueError("CUDA device import has no CUDA C++ source to
bundle")
+ for imported in device_mod.imports:
+ collect(imported)
+ sources.append(source)
+
+ for imported in mod.imports:
+ collect(imported)
+ return "\n\n".join([*sources, host_source])
diff --git a/python/tvm/backend/cuda/tile_primitive/copy_async/tma.py
b/python/tvm/backend/cuda/tile_primitive/copy_async/tma.py
index 1f75023fb2..679681d4e5 100644
--- a/python/tvm/backend/cuda/tile_primitive/copy_async/tma.py
+++ b/python/tvm/backend/cuda/tile_primitive/copy_async/tma.py
@@ -2048,21 +2048,20 @@ def _get_or_encode_descriptor(spec: TensorMapSpec,
sctx: DispatchContext):
@T.prim_func(check_well_formed=False)
def create_tensor_map():
T.Bind(T.tvm_stack_alloca("tensormap", 1), var=tensor_map)
- T.call_packed(
- "runtime.cuTensorMapEncodeTiled",
+ T.tensormap_encode_tiled(
tensor_map,
- spec.descriptor_dtype,
- spec.rank,
spec.base,
*spec.global_dims,
*spec.global_strides,
*spec.box_dims,
*spec.element_strides,
- spec.interleave,
- spec.swizzle,
- spec.l2_promotion,
- spec.oob_fill,
- *([spec.force_cu_dtype] if spec.force_cu_dtype >= 0 else []),
+ descriptor_dtype=spec.descriptor_dtype,
+ rank=spec.rank,
+ interleave=spec.interleave,
+ swizzle=spec.swizzle,
+ l2_promotion=spec.l2_promotion,
+ oob_fill=spec.oob_fill,
+ force_cu_dtype=spec.force_cu_dtype,
)
T.tvm_kernel_replace_point()
# fmt: on
diff --git a/python/tvm/tirx/__init__.py b/python/tvm/tirx/__init__.py
index 500e938755..d0877c142a 100644
--- a/python/tvm/tirx/__init__.py
+++ b/python/tvm/tirx/__init__.py
@@ -61,6 +61,7 @@ from .function import PrimFunc, IndexMap
from .op import call_packed_lowered, call_cpacked_lowered,
register_intrin_lowering
from .op import call_packed, call_cpacked, call_intrin, call_pure_extern,
call_extern
+from .op import CallFFIKernelAttr, call_ffi_kernel, TensorMapEncodeTiledAttr,
tensormap_encode_tiled
from .op import call_llvm_intrin, call_llvm_pure_intrin, all, any, min_value,
max_value, trace
from .op import tvm_stack_alloca, tvm_stack_make_shape, tvm_stack_make_array
from .op import tvm_tuple, handle_add_byte_offset, tvm_struct_get,
tvm_struct_set
diff --git a/python/tvm/tirx/build.py b/python/tvm/tirx/build.py
index f15650a3ba..214da90ede 100644
--- a/python/tvm/tirx/build.py
+++ b/python/tvm/tirx/build.py
@@ -102,7 +102,7 @@ def split_host_device_mods(mod: IRModule) ->
tuple[IRModule, dict[Target, IRModu
def is_host_func(f):
target = f.attrs.get("target", tvm.target.Target("llvm"))
- return target.kind.name in ["llvm", "c"]
+ return target.kind.name in ["llvm", "c", "cuda_host"]
host_mod = tvm.tirx.transform.Filter(is_host_func)(mod)
device_mod = tvm.tirx.transform.Filter(lambda f: not is_host_func(f))(mod)
diff --git a/python/tvm/tirx/op.py b/python/tvm/tirx/op.py
index 291ec8720d..2369f9d5b5 100644
--- a/python/tvm/tirx/op.py
+++ b/python/tvm/tirx/op.py
@@ -255,6 +255,90 @@ def call_packed(*args, span=None):
return Call(Op.get("tirx.tvm_call_packed"), call_args, span=span,
ret_ty="int32")
+@tvm_ffi.register_object("tirx.CallFFIKernelAttr")
+class CallFFIKernelAttr(tvm.ir.Attrs):
+ """Ordered launch tags for an explicit FFI kernel call."""
+
+ launch_params: list[str]
+
+ def __init__(self, launch_params):
+ self.__init_handle_by_constructor__(_ffi_api.CallFFIKernelAttr,
launch_params)
+
+
+def call_ffi_kernel(*args, launch_params, ret_ty="int32", span=None):
+ """Call a kernel with its symbol, kernel operands, then launch values.
+
+ ``launch_params`` contains ordered tags for the launch-value suffix.
+ Flag-only tags consume no argument, and dynamic shared-memory bytes are
+ last when present. Host codegen may launch directly; other hosts use the
+ existing packed-function calling convention.
+ """
+ return Call(
+ "tirx.call_ffi_kernel",
+ args,
+ attrs=CallFFIKernelAttr(launch_params),
+ ret_ty=ret_ty,
+ span=span,
+ )
+
+
+@tvm_ffi.register_object("tirx.TensorMapEncodeTiledAttr")
+class TensorMapEncodeTiledAttr(tvm.ir.Attrs):
+ """Descriptor dtype and fixed options for tiled tensor-map encoding."""
+
+ def __init__(
+ self,
+ descriptor_dtype,
+ rank,
+ interleave=0,
+ swizzle=0,
+ l2_promotion=0,
+ oob_fill=0,
+ force_cu_dtype=-1,
+ ):
+ self.__init_handle_by_constructor__(
+ _ffi_api.TensorMapEncodeTiledAttr,
+ descriptor_dtype,
+ rank,
+ interleave,
+ swizzle,
+ l2_promotion,
+ oob_fill,
+ force_cu_dtype,
+ )
+
+
+def tensormap_encode_tiled(
+ *args,
+ descriptor_dtype,
+ rank,
+ interleave=0,
+ swizzle=0,
+ l2_promotion=0,
+ oob_fill=0,
+ force_cu_dtype=-1,
+ span=None,
+):
+ """Encode a tiled tensor map using runtime pointers and shape operands.
+
+ Arguments are the descriptor and data pointers, global dimensions (rank),
+ byte strides (rank - 1), box dimensions (rank), and element strides (rank).
+ The dtype describes the final descriptor units, including any promotion.
+ CUDA-host codegen encodes directly; other hosts use the runtime packed
call.
+ """
+ if not 1 <= rank <= 5 or len(args) != 4 * rank + 1:
+ raise ValueError("tensormap_encode_tiled requires rank 1..5 and 4 *
rank + 1 operands")
+ return Call(
+ "tirx.tensormap_encode_tiled",
+ args,
+ attrs=TensorMapEncodeTiledAttr(
+ descriptor_dtype, rank, interleave, swizzle, l2_promotion,
oob_fill, force_cu_dtype
+ ),
+ ret_ty="int32",
+ span=span,
+ )
+
+
def call_cpacked(*args, span=None):
"""Build expression by call an external packed function.
diff --git a/python/tvm/tirx/script/builder/ir.py
b/python/tvm/tirx/script/builder/ir.py
index 874e981a27..4364d5102d 100644
--- a/python/tvm/tirx/script/builder/ir.py
+++ b/python/tvm/tirx/script/builder/ir.py
@@ -3202,6 +3202,8 @@ tvm_stack_alloca = _op_wrapper(_tir_op.tvm_stack_alloca)
tvm_stack_make_shape = _op_wrapper(_tir_op.tvm_stack_make_shape)
tvm_stack_make_array = _op_wrapper(_tir_op.tvm_stack_make_array)
call_packed = _op_wrapper(_tir_op.call_packed)
+call_ffi_kernel = _op_wrapper(_tir_op.call_ffi_kernel)
+tensormap_encode_tiled = _op_wrapper(_tir_op.tensormap_encode_tiled)
call_cpacked = _op_wrapper(_tir_op.call_cpacked)
call_packed_lowered = _op_wrapper(_tir_op.call_packed_lowered)
call_cpacked_lowered = _op_wrapper(_tir_op.call_cpacked_lowered)
@@ -3520,6 +3522,8 @@ __all__ = [
"tvm_stack_make_shape",
"tvm_stack_make_array",
"call_packed",
+ "call_ffi_kernel",
+ "tensormap_encode_tiled",
"call_cpacked",
"call_packed_lowered",
"call_cpacked_lowered",
diff --git a/src/backend/cuda/codegen/codegen_cuda_host.cc
b/src/backend/cuda/codegen/codegen_cuda_host.cc
new file mode 100644
index 0000000000..5fb4aee7e5
--- /dev/null
+++ b/src/backend/cuda/codegen/codegen_cuda_host.cc
@@ -0,0 +1,401 @@
+/*
+ * 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.
+ */
+
+/*! \file codegen_cuda_host.cc
+ * \brief C host wrappers with direct CUDA kernel launches.
+ */
+#include <tvm/ffi/reflection/registry.h>
+#include <tvm/tirx/attrs.h>
+#include <tvm/tirx/type.h>
+
+#include <algorithm>
+#include <array>
+#include <limits>
+#include <string>
+#include <unordered_map>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+#include "../../../runtime/metadata.h"
+#include "../../../target/source/codegen_c_host.h"
+
+namespace tvm {
+namespace codegen {
+
+class CodeGenCUDAHost : public CodeGenCHost {
+ public:
+ void Init(Target target) {
+ CodeGenCHost::Init(false, true, true, target->str(), {});
+ // Retain C-host wrapper initialization, with an FFI/CUDA-only prelude.
+ decl_stream.str("");
+ decl_stream.clear();
+ decl_stream << "#include <tvm/ffi/c_api.h>\n"
+ << "#include <tvm/ffi/extra/c_env_api.h>\n"
+ << "#include <tvm/ffi/extra/cuda/device_guard.h>\n"
+ << "#include <tvm/ffi/function.h>\n"
+ << "#include <cuda.h>\n#include <cuda_runtime.h>\n#include
<math.h>\n"
+ << "#define TVM_DLL TVM_FFI_DLL_EXPORT\n";
+ InitGlobalContext();
+ check_error_ = name_supply_->FreshName("tvm_cuda_host_check");
+ decl_stream << "static int " << check_error_
+ << "(cudaError_t error, bool allow_unloading = false) {\n"
+ << " if (error == cudaSuccess || "
+ << "(allow_unloading && error == cudaErrorCudartUnloading))
return 0;\n"
+ << " const char* parts[] = {\"CUDA launch: \",
cudaGetErrorString(error)};\n"
+ << " TVMFFIErrorSetRaisedFromCStrParts(\"CUDAError\", parts,
2);\n"
+ << " cudaGetLastError();\n"
+ << " return -1;\n}\n";
+ }
+
+ using CodeGenCHost::PrintType;
+
+ void PrintType(const Type& type, std::ostream& os) override {
+ if (type.as<tirx::TensorMapTypeNode>()) {
+ os << "CUtensorMap";
+ } else {
+ CodeGenCHost::PrintType(type, os);
+ }
+ }
+
+ void AddFunction(const GlobalVar& gvar, const PrimFunc& func) override {
+ // A guard destructor may report a CUDA error on any return path. Keep it
+ // inside the FFI exception boundary along with device selection and setup.
+ InitFuncState(func);
+ PrintFunctionSignature(GetFunctionName(gvar), func, stream);
+ stream << " {\n TVM_FFI_SAFE_CALL_BEGIN();\n try {\n";
+ int scope = BeginScope();
+ PrintStmt(func->body);
+ EndScope(scope);
+ // CUDADeviceGuard reports errors by throwing. Clear the reported CUDA
+ // last-error state so a subsequent successful launch is not blamed for it.
+ stream << " } catch (...) {\n cudaGetLastError();\n throw;\n }\n"
+ << " TVM_FFI_SAFE_CALL_END();\n}\n\n";
+ if (auto symbol = func->GetAttr<ffi::String>(tvm::attr::kGlobalSymbol)) {
+ function_names_.push_back(symbol.value());
+ if (func->HasNonzeroAttr(tirx::attr::kIsEntryFunc) &&
!has_tvm_ffi_main_func_) {
+ function_names_.push_back(ffi::symbol::tvm_ffi_main);
+ PrintFuncPrefix(stream);
+ PrintType(func->ret_type, stream);
+ stream << " " << ffi::symbol::tvm_ffi_main
+ << "(void* self, void* args, int num_args, void* result) {\n"
+ << " return " << symbol.value() << "(self, args, num_args,
result);\n}\n";
+ }
+ }
+ }
+
+ ffi::Array<ffi::String> GetFunctionNames() { return function_names_; }
+
+ void Dispatch_(const CallNode* op, std::ostream& os) override {
+ if (op->op.same_as(tirx::builtin::tvm_stack_alloca()) &&
+ op->args[0].as_or_throw<StringImm>()->value == "tensormap") {
+ auto count = op->args[1].as_or_throw<IntImm>()->value.as<int64_t>();
+ TVM_FFI_CHECK(count.has_value() && *count > 0, ValueError)
+ << "cuda_host requires a positive constant tensor-map allocation
count";
+ std::string name = name_supply_->FreshName("cuda_tensormap");
+ PrintIndent();
+ stream << "alignas(64) CUtensorMap " << name << "[" << *count << "];\n";
+ os << name;
+ return;
+ }
+ if (op->op.same_as(tirx::builtin::tvm_call_packed_lowered())) {
+ const auto& name = op->args[0].as_or_throw<StringImm>()->value;
+ if (name == "runtime.cuTensorMapEncodeTiled" || name ==
"runtime.cuTensorMapInit") {
+ TVM_FFI_THROW(ValueError)
+ << "cuda_host requires tensormap_encode_tiled instead of a packed
tensor-map encoder";
+ }
+ if (name == "__tvm_set_device") {
+ std::string args =
+ "((TVMFFIAny*)" + PrintExpr(op->args[1]) + " + " +
PrintExpr(op->args[2]) + ")";
+ std::string guard = name_supply_->FreshName("cuda_device_guard");
+ PrintIndent();
+ stream << "if (" << args << "[0].v_int64 != kDLCUDA) {\n"
+ << " TVMFFIErrorSetRaisedFromCStr(\"ValueError\", "
+ << "\"cuda_host requires a CUDA device\");\n return -1;\n}\n";
+ PrintIndent();
+ stream << "tvm::ffi::CUDADeviceGuard " << guard << "(" << args <<
"[1].v_int64);\n";
+ os << "0";
+ return;
+ }
+ std::string function_name = name;
+ TVM_FFI_CHECK(function_name.rfind("runtime.", 0) != 0 &&
+ function_name.rfind("device_api.", 0) != 0 &&
+ function_name.rfind("__tvm_", 0) != 0,
+ ValueError)
+ << "cuda_host does not provide TVM runtime service: " << name;
+ }
+ if (auto call_op = op->op.as<Op>()) {
+ if (op_attr_global_symbol_.count(call_op.value())) {
+ const auto& symbol = op_attr_global_symbol_[call_op.value()];
+ TVM_FFI_CHECK(std::string(symbol).rfind("TVMBackend", 0) != 0,
ValueError)
+ << "cuda_host does not provide TVM runtime operation: " << symbol;
+ }
+ }
+ if (op->op.same_as(tirx::builtin::call_extern()) ||
+ op->op.same_as(tirx::builtin::call_pure_extern())) {
+ const auto& symbol = op->args[0].as_or_throw<StringImm>()->value;
+ TVM_FFI_CHECK(std::string(symbol).rfind("TVMBackend", 0) != 0,
ValueError)
+ << "cuda_host does not provide TVM runtime operation: " << symbol;
+ }
+ if (op->op.same_as(tirx::builtin::tensormap_encode_tiled())) {
+ PrintTensorMapEncode(op);
+ os << "0";
+ return;
+ }
+ if (!op->op.same_as(tirx::builtin::call_ffi_kernel())) {
+ CodeGenCHost::Dispatch_(op, os);
+ return;
+ }
+ const auto* attr = op->attrs.as<tirx::CallFFIKernelAttr>();
+ TVM_FFI_CHECK(attr, ValueError) << "cuda_host kernel calls require
CallFFIKernelAttr";
+ TVM_FFI_CHECK(!op->args.empty() && op->args[0].as<StringImmNode>(),
ValueError)
+ << "cuda_host kernel calls require a string kernel symbol";
+ const auto& symbol = op->args[0].as<StringImmNode>()->value;
+
+ std::array<int, 6> axes;
+ axes.fill(-1);
+ int shared_memory = -1;
+ size_t num_launch_values = 0;
+ std::unordered_set<std::string> seen;
+ for (size_t i = 0; i < attr->launch_params.size(); ++i) {
+ std::string tag = attr->launch_params[i];
+ TVM_FFI_CHECK(seen.insert(tag).second, ValueError)
+ << "cuda_host duplicate launch parameter: " << tag;
+ // These are flags, not values in the argument suffix. Reject
unsupported
+ // launch semantics before decoding any operands, rather than dropping
them.
+ if (tag == runtime::launch_param::kUseProgramaticDependentLaunch ||
+ tag == runtime::launch_param::kUseCooperativeLaunch ||
+ tag == runtime::launch_param::kUseRequiredBlockDimension) {
+ TVM_FFI_THROW(ValueError) << "cuda_host does not support launch flag:
" << tag;
+ } else if (tag == runtime::launch_param::kUseDynamicSharedMemoryTag) {
+ TVM_FFI_CHECK_EQ(i + 1, attr->launch_params.size(), ValueError)
+ << "cuda_host dynamic shared memory must be the last launch
parameter";
+ shared_memory = static_cast<int>(num_launch_values++);
+ } else {
+ int axis = -1;
+ for (int j = 0; j < 3; ++j) {
+ if (tag == std::string("blockIdx.") + "xyz"[j]) axis = j;
+ if (tag == std::string("threadIdx.") + "xyz"[j]) axis = j + 3;
+ }
+ TVM_FFI_CHECK_GE(axis, 0, ValueError)
+ << "cuda_host does not support launch parameter: " << tag;
+ axes[axis] = static_cast<int>(num_launch_values++);
+ }
+ }
+ TVM_FFI_CHECK_GE(op->args.size(), num_launch_values + 1, ValueError)
+ << "cuda_host kernel call is missing launch arguments";
+ size_t launch_begin = op->args.size() - num_launch_values;
+ for (size_t i = launch_begin; i < op->args.size(); ++i) {
+ auto type = op->args[i]->ty.as<PrimType>();
+ TVM_FFI_CHECK(
+ type.has_value() && type.value().IsScalar() &&
type.value().MatchesCode(kDLInt, kDLUInt),
+ ValueError)
+ << "cuda_host launch arguments must be scalar integers";
+ }
+
+ // Evaluate arguments once in their original order, including launch
values.
+ // Buffer expressions retain their pointer type for the direct kernel call.
+ std::vector<std::string> arguments;
+ for (size_t i = 1; i < op->args.size(); ++i) {
+ std::string value = PrintExpr(op->args[i]);
+ if (i < launch_begin) {
+ if (auto* ptr = op->args[i]->ty.as<PointerTypeNode>()) {
+ if (ptr->element_type.as<tirx::TensorMapTypeNode>()) {
+ // The device ABI takes a grid-constant descriptor by value, while
+ // the host IR carries a pointer to its aligned stack storage.
+ value = "*(" + value + ")";
+ }
+ }
+ }
+ std::string name = name_supply_->FreshName("cuda_arg");
+ PrintIndent();
+ stream << "auto " << name << " = " << value << ";\n";
+ arguments.push_back(name);
+ }
+ auto launch_arg = [&](int index) { return arguments[launch_begin - 1 +
index]; };
+ std::array<std::string, 6> dimensions;
+ for (size_t i = 0; i < axes.size(); ++i) {
+ if (axes[i] < 0) {
+ dimensions[i] = "1";
+ } else {
+ dimensions[i] = name_supply_->FreshName("cuda_dim");
+ PrintIndent();
+ stream << "size_t " << dimensions[i] << " = static_cast<size_t>(" <<
launch_arg(axes[i])
+ << ");\n";
+ // Match LaunchParamConfig's handling of empty dynamic dimensions.
+ PrintIndent();
+ stream << "if (" << dimensions[i] << " == 0) " << dimensions[i] << " =
1;\n";
+ }
+ }
+ std::string device = name_supply_->FreshName("cuda_device");
+ std::string cuda_stream = name_supply_->FreshName("cuda_stream");
+ PrintIndent();
+ stream << "int " << device << ";\n";
+ CheckError("cudaGetDevice(&" + device + ")");
+ PrintIndent();
+ stream << "cudaStream_t " << cuda_stream
+ << " = static_cast<cudaStream_t>(TVMFFIEnvGetStream(kDLCUDA, " <<
device << "));\n";
+ std::string bytes = shared_memory < 0 ? "0" : launch_arg(shared_memory);
+ if (shared_memory >= 0) {
+ PrintIndent();
+ stream << "if (" << bytes << " >= 49152) {\n";
+ int scope = BeginScope();
+ CheckError("cudaFuncSetAttribute(::" + std::string(symbol) +
+ ", cudaFuncAttributeMaxDynamicSharedMemorySize, " + bytes +
")");
+ EndScope(scope);
+ PrintIndent();
+ stream << "}\n";
+ }
+ PrintIndent();
+ stream << "::" << symbol << "<<<dim3(" << dimensions[0] << ", " <<
dimensions[1] << ", "
+ << dimensions[2] << "), dim3(" << dimensions[3] << ", " <<
dimensions[4] << ", "
+ << dimensions[5] << "), " << bytes << ", " << cuda_stream << ">>>(";
+ for (size_t i = 0; i + 1 < launch_begin; ++i) {
+ if (i != 0) stream << ", ";
+ stream << arguments[i];
+ }
+ stream << ");\n";
+ CheckError("cudaGetLastError()", true);
+ os << "0";
+ }
+
+ private:
+ void CheckError(const std::string& call, bool allow_unloading = false) {
+ PrintIndent();
+ stream << "if (::" << check_error_ << "(" << call;
+ if (allow_unloading) stream << ", true";
+ stream << ") != 0) return -1;\n";
+ }
+ void PrintTensorMapEncode(const CallNode* op) {
+ const auto* attr = op->attrs.as<tirx::TensorMapEncodeTiledAttr>();
+ TVM_FFI_CHECK(attr && attr->rank >= 1 && attr->rank <= 5 &&
+ op->args.size() == static_cast<size_t>(4 * attr->rank +
1),
+ ValueError)
+ << "Invalid tensormap_encode_tiled attributes or operands";
+ static const std::unordered_map<std::string, std::string> dtype_names{
+ {"int8", "UINT8"}, {"uint8", "UINT8"},
+ {"uint16", "UINT16"}, {"uint32", "UINT32"},
+ {"uint64", "UINT64"}, {"int32", "INT32"},
+ {"int64", "INT64"}, {"float16", "FLOAT16"},
+ {"float32", "FLOAT32"}, {"float64", "FLOAT64"},
+ {"bfloat16", "BFLOAT16"}, {"float8_e4m3fn", "UINT8"},
+ {"float8_e5m2", "UINT8"}, {"float4_e2m1fn", "16U4_ALIGN16B"}};
+ auto dtype =
dtype_names.find(ffi::DLDataTypeToString(attr->descriptor_dtype));
+ TVM_FFI_CHECK(dtype != dtype_names.end(), ValueError)
+ << "Unsupported cuda_host tensor-map descriptor dtype: "
+ << ffi::DLDataTypeToString(attr->descriptor_dtype);
+ std::string cuda_dtype = "CU_TENSOR_MAP_DATA_TYPE_" + dtype->second;
+ if (attr->force_cu_dtype != -1) {
+ TVM_FFI_CHECK(attr->force_cu_dtype == 11 && attr->descriptor_dtype.code
== kDLFloat &&
+ attr->descriptor_dtype.bits == 32 &&
attr->descriptor_dtype.lanes == 1,
+ ValueError)
+ << "cuda_host only supports a TFLOAT32 tensor-map dtype override";
+ cuda_dtype = "CU_TENSOR_MAP_DATA_TYPE_TFLOAT32";
+ }
+ for (int64_t value : {attr->interleave, attr->swizzle, attr->l2_promotion,
attr->oob_fill}) {
+ TVM_FFI_CHECK(value >= 0 && value <= std::numeric_limits<int>::max(),
ValueError)
+ << "cuda_host tensor-map options must fit a nonnegative CUDA enum";
+ }
+ // Preserve operand evaluation order and reject narrowing that could turn
an
+ // invalid dynamic dimension or stride into a valid but different CUDA
input.
+ std::vector<std::string> values;
+ for (size_t i = 0; i < op->args.size(); ++i) {
+ std::string value = PrintExpr(op->args[i]);
+ std::string name = name_supply_->FreshName("tensormap_arg");
+ PrintIndent();
+ stream << "auto " << name << " = " << value << ";\n";
+ values.push_back(name);
+ }
+ for (size_t i = 2; i < op->args.size(); ++i) {
+ const std::string& name = values[i];
+ auto type = op->args[i]->ty.as<PrimType>();
+ TVM_FFI_CHECK(type && type.value().IsScalar() && type.value().bits() <=
64 &&
+ type.value().MatchesCode(kDLInt, kDLUInt),
+ ValueError)
+ << "cuda_host tensor-map dimensions and strides must be scalar
integers";
+ if (type.value().MatchesCode(kDLInt)) {
+ PrintIndent();
+ stream << "TVM_FFI_CHECK(" << name << " >= 0, ValueError) "
+ << "<< \"Negative tensor-map dimension or stride\";\n";
+ }
+ if (i >= static_cast<size_t>(2 * attr->rank + 1)) {
+ PrintIndent();
+ stream << "TVM_FFI_CHECK(static_cast<uint64_t>(" << name
+ << ") <= 4294967295ULL, ValueError) "
+ << "<< \"Tensor-map dimension or stride exceeds uint32\";\n";
+ }
+ }
+ size_t index = 2;
+ auto array = [&](const char* type, int64_t count) {
+ std::string name = name_supply_->FreshName("tensormap_values");
+ PrintIndent();
+ stream << type << " " << name << "[" << std::max<int64_t>(count, 1) <<
"] = {";
+ for (int64_t i = 0; i < count; ++i) {
+ if (i) stream << ", ";
+ stream << "static_cast<" << type << ">(" << values[index++] << ")";
+ }
+ stream << "};\n";
+ return name;
+ };
+ std::string shape = array("cuuint64_t", attr->rank);
+ std::string strides = array("cuuint64_t", attr->rank - 1);
+ std::string box = array("cuuint32_t", attr->rank);
+ std::string element_strides = array("cuuint32_t", attr->rank);
+ std::string error = name_supply_->FreshName("tensormap_error");
+ PrintIndent();
+ stream << "CUresult " << error << " =
cuTensorMapEncodeTiled(static_cast<CUtensorMap*>("
+ << values[0] << "), " << cuda_dtype << ", " << attr->rank << ", "
<< values[1] << ", "
+ << shape << ", " << strides << ", " << box << ", " <<
element_strides
+ << ", static_cast<CUtensorMapInterleave>(" << attr->interleave
+ << "), static_cast<CUtensorMapSwizzle>(" << attr->swizzle
+ << "), static_cast<CUtensorMapL2promotion>(" << attr->l2_promotion
+ << "), static_cast<CUtensorMapFloatOOBfill>(" << attr->oob_fill <<
"));\n";
+ PrintIndent();
+ stream << "if (" << error << " != CUDA_SUCCESS) {\n"
+ << " const char* message = \"cuTensorMapEncodeTiled failed\";\n"
+ << " cuGetErrorString(" << error << ", &message);\n"
+ << " TVMFFIErrorSetRaisedFromCStr(\"CUDAError\", message);\n"
+ << " return -1;\n}\n";
+ }
+
+ std::string check_error_;
+ ffi::Array<ffi::String> function_names_;
+};
+
+ffi::Module BuildCUDAHost(IRModule mod, Target target) {
+ CodeGenCUDAHost cg;
+ cg.Init(target);
+ std::vector<std::pair<GlobalVar, PrimFunc>> functions;
+ for (auto [gvar, base_func] : mod->functions) {
+ functions.emplace_back(gvar, base_func.as_or_throw<PrimFunc>());
+ }
+ std::sort(functions.begin(), functions.end(), [](const auto& lhs, const
auto& rhs) {
+ return lhs.first->name_hint < rhs.first->name_hint;
+ });
+ for (const auto& [gvar, func] : functions) cg.DeclareFunction(gvar, func);
+ for (const auto& [gvar, func] : functions) cg.AddFunction(gvar, func);
+ return CSourceModuleCreate(cg.Finish(), "cu", cg.GetFunctionNames());
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ ffi::reflection::GlobalDef().def("target.build.cuda_host", BuildCUDAHost);
+}
+
+} // namespace codegen
+} // namespace tvm
diff --git a/src/backend/cuda/codegen/target_kind.cc
b/src/backend/cuda/codegen/target_kind.cc
index 3624983cef..b63f1f6e96 100644
--- a/src/backend/cuda/codegen/target_kind.cc
+++ b/src/backend/cuda/codegen/target_kind.cc
@@ -114,6 +114,8 @@ ffi::Map<ffi::String, ffi::Any>
UpdateNVPTXAttrs(ffi::Map<ffi::String, ffi::Any>
void RegisterTargetKinds() {
namespace refl = tvm::ffi::reflection;
+ TVM_REGISTER_TARGET_KIND("cuda_host", kDLCPU).set_default_keys({"cpu"});
+
TVM_REGISTER_TARGET_KIND("cuda", kDLCUDA)
.add_attr_option<ffi::String>("mcpu")
.add_attr_option<ffi::String>("arch")
diff --git a/src/tirx/op/builtin.cc b/src/tirx/op/builtin.cc
index 84db3c9749..f2856276aa 100644
--- a/src/tirx/op/builtin.cc
+++ b/src/tirx/op/builtin.cc
@@ -24,6 +24,7 @@
*/
#include <tvm/ffi/function.h>
#include <tvm/ir/prim/builtin.h>
+#include <tvm/tirx/attrs.h>
#include <tvm/tirx/builtin.h>
#include <tvm/tirx/op.h>
#include <tvm/tirx/op_attr_types.h>
@@ -32,6 +33,31 @@ namespace tvm {
namespace tirx {
namespace builtin {
+TVM_FFI_STATIC_INIT_BLOCK() {
+ TensorMapEncodeTiledAttr::RegisterReflection();
+ ffi::reflection::GlobalDef().def(
+ "tirx.TensorMapEncodeTiledAttr",
+ [](DLDataType descriptor_dtype, int64_t rank, int64_t interleave,
int64_t swizzle,
+ int64_t l2_promotion, int64_t oob_fill, int64_t force_cu_dtype) {
+ auto attrs = ffi::make_object<TensorMapEncodeTiledAttr>();
+ attrs->descriptor_dtype = descriptor_dtype;
+ attrs->rank = rank;
+ attrs->interleave = interleave;
+ attrs->swizzle = swizzle;
+ attrs->l2_promotion = l2_promotion;
+ attrs->oob_fill = oob_fill;
+ attrs->force_cu_dtype = force_cu_dtype;
+ return Attrs(attrs);
+ });
+ CallFFIKernelAttr::RegisterReflection();
+ ffi::reflection::GlobalDef().def("tirx.CallFFIKernelAttr",
+ [](ffi::Array<ffi::String> launch_params) {
+ auto attrs =
ffi::make_object<CallFFIKernelAttr>();
+ attrs->launch_params =
std::move(launch_params);
+ return Attrs(attrs);
+ });
+}
+
// Script metadata extends the canonical primitive operators registered by IR.
TVM_FFI_STATIC_INIT_BLOCK() {
#define PRIM_SCRIPT_BUILTIN(OpName) \
@@ -207,6 +233,12 @@ TIR_DEFINE_BUILTIN_FUNC(tvm_call_packed)
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque))
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("call_packed"), /*plevel=*/20);
+TIR_DEFINE_BUILTIN_FUNC(tensormap_encode_tiled)
+ .set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
+
+TIR_DEFINE_BUILTIN_FUNC(call_ffi_kernel)
+ .set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
+
TIR_DEFINE_BUILTIN_FUNC(tvm_call_cpacked)
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque))
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("call_cpacked"), /*plevel=*/20);
diff --git a/src/tirx/script/printer/expr.cc b/src/tirx/script/printer/expr.cc
index e3a77b19ab..f694b844ed 100644
--- a/src/tirx/script/printer/expr.cc
+++ b/src/tirx/script/printer/expr.cc
@@ -18,6 +18,7 @@
*/
#include <tvm/ir/prim/builtin.h>
#include <tvm/te/operation.h>
+#include <tvm/tirx/attrs.h>
#include <tvm/tirx/builtin.h>
#include <tvm/tirx/type.h>
@@ -377,6 +378,32 @@ Doc PrintTIRCall(Call call, AccessPath call_p, IRDocsifier
d) {
for (int i = 0; i < n_args; ++i) {
call_args.push_back(d->AsDoc<ExprDoc>(call->args[i],
call_p->Attr("args")->ArrayItem(i)));
}
+ if (call->op.same_as(tirx::builtin::tensormap_encode_tiled())) {
+ const auto* attrs = call->attrs.as<tirx::TensorMapEncodeTiledAttr>();
+ TVM_FFI_ICHECK(attrs);
+ auto attr_p = call_p->Attr("attrs");
+ return TIR(d, "tensormap_encode_tiled")
+ ->Call(call_args,
+ {"descriptor_dtype", "rank", "interleave", "swizzle",
"l2_promotion", "oob_fill",
+ "force_cu_dtype"},
+
{LiteralDoc::Str(ffi::DLDataTypeToString(attrs->descriptor_dtype),
+ attr_p->Attr("descriptor_dtype")),
+ LiteralDoc::Int(attrs->rank, attr_p->Attr("rank")),
+ LiteralDoc::Int(attrs->interleave,
attr_p->Attr("interleave")),
+ LiteralDoc::Int(attrs->swizzle, attr_p->Attr("swizzle")),
+ LiteralDoc::Int(attrs->l2_promotion,
attr_p->Attr("l2_promotion")),
+ LiteralDoc::Int(attrs->oob_fill, attr_p->Attr("oob_fill")),
+ LiteralDoc::Int(attrs->force_cu_dtype,
attr_p->Attr("force_cu_dtype"))});
+ }
+ if (call->op.same_as(tirx::builtin::call_ffi_kernel())) {
+ const auto* attrs = call->attrs.as<tirx::CallFFIKernelAttr>();
+ TVM_FFI_ICHECK(attrs);
+ return TIR(d, "call_ffi_kernel")
+ ->Call(call_args, {"launch_params", "ret_ty"},
+ {d->AsDoc<ExprDoc>(attrs->launch_params,
+
call_p->Attr("attrs")->Attr("launch_params")),
+ get_call_return_type_doc()});
+ }
ExprDoc op_doc = call->op.as<Op>()
? LiteralDoc::Str(call->op.as<Op>().value()->name,
call_p->Attr("op"))
: d->AsDoc<ExprDoc>(call->op, call_p->Attr("op"));
diff --git a/src/tirx/transform/lower_tirx_dedup_tensormap.cc
b/src/tirx/transform/lower_tirx_dedup_tensormap.cc
index 23f9a1cbab..5a9c5c8111 100644
--- a/src/tirx/transform/lower_tirx_dedup_tensormap.cc
+++ b/src/tirx/transform/lower_tirx_dedup_tensormap.cc
@@ -53,10 +53,12 @@ inline bool IsTensorMapAlloca(const BindNode* bind) {
return false;
}
-// Is an Evaluate of tvm_call_packed("runtime.cuTensorMapEncodeTiled", ...)?
+// Recognize typed encoding and legacy manually authored packed encoding.
inline const CallNode* AsCuTensorMapEncode(const EvaluateNode* eval) {
const CallNode* call = eval->value.as<CallNode>();
- if (!call || !call->op.same_as(builtin::tvm_call_packed())) return nullptr;
+ if (!call) return nullptr;
+ if (call->op.same_as(builtin::tensormap_encode_tiled())) return call;
+ if (!call->op.same_as(builtin::tvm_call_packed())) return nullptr;
if (call->args.empty()) return nullptr;
if (const auto* s = call->args[0].as<StringImmNode>()) {
if (s->value == "runtime.cuTensorMapEncodeTiled") return call;
@@ -64,23 +66,17 @@ inline const CallNode* AsCuTensorMapEncode(const
EvaluateNode* eval) {
return nullptr;
}
-// Extract the tensormap var and the key (arguments after the tensormap var)
-inline std::pair<ffi::Optional<Var>, ffi::Array<Expr>> ExtractEncodeKey(const
CallNode* call) {
- TVM_FFI_ICHECK(call->op.same_as(builtin::tvm_call_packed()));
- // args[0] is function name, args[1] is tensormap handle, rest are parameters
- if (call->args.size() < 2) return {ffi::Optional<Var>(), ffi::Array<Expr>()};
- ffi::Optional<Var> tensormap;
- if (auto v = call->args[1].as<Var>()) {
- tensormap = v.value();
- } else {
- tensormap = ffi::Optional<Var>();
+// Exclude only the output pointer; retain op, attributes and all input
operands
+// so descriptor dtype, forced dtype and encoding modes participate in
equality.
+inline std::pair<ffi::Optional<Var>, Call> ExtractEncodeKey(const CallNode*
call) {
+ size_t output_index = call->op.same_as(builtin::tensormap_encode_tiled()) ?
0 : 1;
+ TVM_FFI_ICHECK_GT(call->args.size(), output_index);
+ ffi::Optional<Var> tensormap = call->args[output_index].as<Var>();
+ ffi::Array<Expr> args;
+ for (size_t i = 0; i < call->args.size(); ++i) {
+ if (i != output_index) args.push_back(call->args[i]);
}
- ffi::Array<Expr> key;
- key.reserve(call->args.size() - 2);
- for (size_t i = 2; i < call->args.size(); ++i) {
- key.push_back(call->args[i]);
- }
- return {tensormap, key};
+ return {tensormap, Call(call->ty, call->op, args, call->attrs)};
}
} // namespace
@@ -88,14 +84,12 @@ inline std::pair<ffi::Optional<Var>, ffi::Array<Expr>>
ExtractEncodeKey(const Ca
// First pass: Analyze encode calls and decide canonical tensormap
per-parameter set
class CuTensorMapDedupAnalyzer : public StmtExprVisitor {
public:
- CuTensorMapDedupAnalyzer() {
- canonical_list_.emplace_back(std::vector<std::pair<ffi::Array<Expr>,
Var>>());
- }
+ CuTensorMapDedupAnalyzer() {
canonical_list_.emplace_back(std::vector<std::pair<Call, Var>>()); }
ffi::Optional<VisitInterrupt> Visit_(const ForNode* op) final {
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit(op->min));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit(op->extent));
- canonical_list_.emplace_back(std::vector<std::pair<ffi::Array<Expr>,
Var>>());
+ canonical_list_.emplace_back(std::vector<std::pair<Call, Var>>());
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit(op->body));
canonical_list_.pop_back();
return std::nullopt;
@@ -103,7 +97,7 @@ class CuTensorMapDedupAnalyzer : public StmtExprVisitor {
ffi::Optional<VisitInterrupt> Visit_(const WhileNode* op) final {
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit(op->condition));
- canonical_list_.emplace_back(std::vector<std::pair<ffi::Array<Expr>,
Var>>());
+ canonical_list_.emplace_back(std::vector<std::pair<Call, Var>>());
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit(op->body));
canonical_list_.pop_back();
return std::nullopt;
@@ -111,11 +105,11 @@ class CuTensorMapDedupAnalyzer : public StmtExprVisitor {
ffi::Optional<VisitInterrupt> Visit_(const IfThenElseNode* op) final {
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit(op->condition));
- canonical_list_.emplace_back(std::vector<std::pair<ffi::Array<Expr>,
Var>>());
+ canonical_list_.emplace_back(std::vector<std::pair<Call, Var>>());
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit(op->then_case));
canonical_list_.pop_back();
if (op->else_case) {
- canonical_list_.emplace_back(std::vector<std::pair<ffi::Array<Expr>,
Var>>());
+ canonical_list_.emplace_back(std::vector<std::pair<Call, Var>>());
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit(op->else_case.value()));
canonical_list_.pop_back();
}
@@ -153,7 +147,7 @@ class CuTensorMapDedupAnalyzer : public StmtExprVisitor {
}
private:
- std::vector<std::vector<std::pair<ffi::Array<Expr>, Var>>> canonical_list_;
+ std::vector<std::vector<std::pair<Call, Var>>> canonical_list_;
std::unordered_map<Var, Var, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>
tensormap_var_remap_;
};
@@ -165,7 +159,7 @@ class CuTensorMapDedupRewriter : public StmtExprMutator {
CuTensorMapDedupRewriter(
std::unordered_map<Var, Var, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>
var_remap) {
for (const auto& [source, target] : var_remap) VarRemapSet(source, target);
- emitted_keys_.emplace_back(std::vector<ffi::Array<Expr>>());
+ emitted_keys_.emplace_back(std::vector<Call>());
}
private:
@@ -297,7 +291,7 @@ class CuTensorMapDedupRewriter : public StmtExprMutator {
}
// Track which parameter keys have already emitted an encode call
- std::vector<std::vector<ffi::Array<Expr>>> emitted_keys_;
+ std::vector<std::vector<Call>> emitted_keys_;
};
namespace transform {
diff --git a/src/tirx/transform/lower_tvm_builtin.cc
b/src/tirx/transform/lower_tvm_builtin.cc
index cd3a88da88..9f14fe4d95 100644
--- a/src/tirx/transform/lower_tvm_builtin.cc
+++ b/src/tirx/transform/lower_tvm_builtin.cc
@@ -29,6 +29,7 @@
#include <tvm/ir/prim/expr.h>
#include <tvm/ir/scope_stack.h>
#include <tvm/runtime/logging.h>
+#include <tvm/tirx/attrs.h>
#include <tvm/tirx/builtin.h>
#include <tvm/tirx/stmt_functor.h>
#include <tvm/tirx/transform.h>
@@ -57,17 +58,21 @@ class BuiltinLower : public StmtExprMutator {
using StmtExprMutator::Mutate_;
static PrimFunc Build(PrimFunc func) {
ffi::Optional<PrimExpr> device_type = std::nullopt;
+ bool preserve_ffi_kernel = false;
if (auto target = func->GetAttr<Target>(tvm::attr::kTarget)) {
device_type = IntImm::Int32(target.value()->kind->default_device_type);
+ auto host = target.value()->GetHost().value_or(target.value());
+ preserve_ffi_kernel = host->kind->name == "cuda_host";
}
- auto mutator = ffi::make_object<BuiltinLower>(device_type);
+ auto mutator = ffi::make_object<BuiltinLower>(device_type,
preserve_ffi_kernel);
func.CopyOnWrite()->body = mutator->VisitBodyAndRealizeAlloca(func->body);
return func;
}
- explicit BuiltinLower(ffi::Optional<PrimExpr> device_type = std::nullopt)
- : device_type_(device_type) {}
+ explicit BuiltinLower(ffi::Optional<PrimExpr> device_type = std::nullopt,
+ bool preserve_ffi_kernel = false)
+ : device_type_(device_type), preserve_ffi_kernel_(preserve_ffi_kernel) {}
// NOTE: Right now, we make the following scoping requirement
// for memory allocated by the following primitives
@@ -405,7 +410,27 @@ class BuiltinLower : public StmtExprMutator {
}
UnchangedOr<Expr> Mutate_(const CallNode* op, InplaceMode inplace_mode)
final {
- if (op->op.same_as(builtin::tvm_call_packed())) {
+ if (op->op.same_as(builtin::tensormap_encode_tiled()) &&
!preserve_ffi_kernel_) {
+ const auto* attr = op->attrs.as<TensorMapEncodeTiledAttr>();
+ TVM_FFI_CHECK(attr && attr->rank >= 1 && attr->rank <= 5 &&
+ op->args.size() == static_cast<size_t>(4 * attr->rank
+ 1),
+ ValueError)
+ << "Invalid tensormap_encode_tiled attributes or operands";
+ ffi::Array<Expr> args{StringImm("runtime.cuTensorMapEncodeTiled"),
op->args[0],
+
StringImm(ffi::DLDataTypeToString(attr->descriptor_dtype)),
+ IntImm(PrimType::Int(32), attr->rank),
op->args[1]};
+ for (size_t i = 2; i < op->args.size(); ++i) args.push_back(op->args[i]);
+ for (int64_t value : {attr->interleave, attr->swizzle,
attr->l2_promotion, attr->oob_fill}) {
+ args.push_back(IntImm(PrimType::Int(32), value));
+ }
+ if (attr->force_cu_dtype != -1) {
+ args.push_back(IntImm(PrimType::Int(32), attr->force_cu_dtype));
+ }
+ Call packed(op->ty, builtin::tvm_call_packed(), args);
+ return MakeCallPackedGeneric(packed.get(), 0,
builtin::tvm_call_packed_lowered(), false);
+ }
+ if (op->op.same_as(builtin::tvm_call_packed()) ||
+ (op->op.same_as(builtin::call_ffi_kernel()) && !preserve_ffi_kernel_))
{
return MakeCallPackedGeneric(op, 0, builtin::tvm_call_packed_lowered(),
/* use_last_value_as_traced_value*/ false);
} else if (op->op.same_as(builtin::tvm_call_cpacked())) {
@@ -797,6 +822,8 @@ class BuiltinLower : public StmtExprMutator {
ffi::Optional<PrimExpr> device_type_{std::nullopt};
ffi::Optional<PrimExpr> device_id_{std::nullopt};
+ // CUDA host codegen consumes explicit launches after ordinary builtin
lowering.
+ bool preserve_ffi_kernel_{false};
bool is_precheck_{false};
// Record all stack frames.
diff --git a/src/tirx/transform/split_host_device.cc
b/src/tirx/transform/split_host_device.cc
index b700b963c0..a2ab6f41a1 100644
--- a/src/tirx/transform/split_host_device.cc
+++ b/src/tirx/transform/split_host_device.cc
@@ -31,6 +31,7 @@
#include <tvm/ir/unique_name_supply.h>
#include <tvm/target/target.h>
#include <tvm/tirx/analysis.h>
+#include <tvm/tirx/attrs.h>
#include <tvm/tirx/builtin.h>
#include <tvm/tirx/op.h>
#include <tvm/tirx/stmt_functor.h>
@@ -859,7 +860,10 @@ class DeviceKernelMutator : public StmtExprMutator {
PrimType node_ty = IsVoidType(node->ty) ? PrimType::Void() :
node->ty.as_or_throw<PrimType>();
PrimType ret_ty = node_ty.IsVoid() ? PrimType::Int(32) : node_ty;
- return Call(ret_ty, builtin::tvm_call_packed(),
call_args).as_or_throw<PrimExpr>();
+ auto attrs = ffi::make_object<CallFFIKernelAttr>();
+ attrs->launch_params = dev_info.launch_params;
+ return Call(ret_ty, builtin::call_ffi_kernel(), call_args, Attrs(attrs))
+ .as_or_throw<PrimExpr>();
}
ffi::Optional<Target> current_target_;
diff --git a/tests/python/codegen/test_target_codegen_cuda.py
b/tests/python/codegen/test_target_codegen_cuda.py
index d0459f3e95..6c3aa00edd 100644
--- a/tests/python/codegen/test_target_codegen_cuda.py
+++ b/tests/python/codegen/test_target_codegen_cuda.py
@@ -54,6 +54,47 @@ def setup_cuda_compile_mode(request):
tvm.register_global_func("tvm_callback_cuda_compile", orig_func,
override=True)
[email protected]
[email protected](not env.has_cuda(), reason="need cuda")
+def test_cuda_host_bundle(tmp_path):
+ from shutil import which
+
+ import tvm_ffi.cpp
+
+ from tvm.backend.cuda import export_cuda_host
+
+ if which("nvcc") is None:
+ pytest.skip("CUDA-host compilation requires NVCC")
+
+ @T.prim_func(s_tir=True)
+ def add_one(A: T.Buffer((32,), "float32"), B: T.Buffer((32,), "float32")):
+ for tx in T.thread_binding(32, "threadIdx.x"):
+ B[tx] = A[tx] + T.float32(1)
+
+ target = tvm.target.Target("cuda", host="cuda_host")
+ built = tvm.compile(add_one, target=target).mod
+ source = export_cuda_host(built)
+ assert source.index("__global__") < source.index("<<<")
+ library = tvm_ffi.cpp.build_inline(
+ name="cuda_host_add_one",
+ cuda_sources=source,
+ extra_cuda_cflags=[f"-arch={target.arch}"],
+ build_directory=str(tmp_path),
+ backend="cuda",
+ )
+ loaded = tvm_ffi.load_module(library)
+
+ def run_and_check():
+ dev = tvm.cuda(0)
+ values = np.arange(32, dtype="float32")
+ a = tvm.runtime.tensor(values, dev)
+ b = tvm.runtime.empty((32,), "float32", dev)
+ loaded["add_one"](a, b)
+ tvm.testing.assert_allclose(b.numpy(), values + 1)
+
+ tvm.testing.run_with_gpu_lock(run_and_check)
+
+
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
def test_cuda_vectorize_add():
diff --git
a/tests/python/tirx-transform/test_tir_transform_lower_tvm_builtin.py
b/tests/python/tirx-transform/test_tir_transform_lower_tvm_builtin.py
index bacbe70d42..c1b4ef510c 100644
--- a/tests/python/tirx-transform/test_tir_transform_lower_tvm_builtin.py
+++ b/tests/python/tirx-transform/test_tir_transform_lower_tvm_builtin.py
@@ -110,13 +110,19 @@ def test_lower_call_packed():
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
-def test_lower_call_packed_raw_string():
[email protected]("call", [tvm.tirx.call_packed,
tvm.tirx.call_ffi_kernel])
+def test_lower_call_packed_raw_string(call):
+ def invoke(*args):
+ if call is tvm.tirx.call_ffi_kernel:
+ return call(*args, launch_params=[])
+ return call(*args)
+
@I.ir_module
class Before:
@T.prim_func(s_tir=True)
def main():
T.func_attr({"target": tvm.target.Target("llvm")})
- T.call_packed("testing.echo", "payload")
+ T.evaluate(invoke("testing.echo", "payload"))
@I.ir_module
class Expected:
diff --git
a/tests/python/tirx-transform/test_tir_transform_split_host_device.py
b/tests/python/tirx-transform/test_tir_transform_split_host_device.py
index ed5ce0f9a3..4cbe5cba48 100644
--- a/tests/python/tirx-transform/test_tir_transform_split_host_device.py
+++ b/tests/python/tirx-transform/test_tir_transform_split_host_device.py
@@ -70,7 +70,7 @@ def test_split_host_device():
@T.prim_func(s_tir=True)
def main(n: T.int32):
T.func_attr({"target": T.target("cuda", host={"kind": "llvm",
"opt-level": 0})})
- T.call_packed("main_kernel", n)
+ T.call_ffi_kernel("main_kernel", n, launch_params=[])
@T.prim_func(s_tir=True)
def main_kernel(n: T.int32):
@@ -164,7 +164,7 @@ def test_split_host_device_without_func_host_attribute():
@T.prim_func(s_tir=True)
def main(n: T.int32):
T.func_attr({"target": T.target("llvm")})
- T.call_packed("main_kernel", n)
+ T.call_ffi_kernel("main_kernel", n, launch_params=[])
@T.prim_func(s_tir=True)
def main_kernel(n: T.int32):
@@ -229,7 +229,7 @@ def test_split_host_device_name_collision():
@T.prim_func(s_tir=True)
def main(n: T.int32):
T.func_attr({"target": T.target("cuda", host={"kind": "llvm",
"opt-level": 0})})
- T.call_packed("main_kernel_1", n)
+ T.call_ffi_kernel("main_kernel_1", n, launch_params=[])
@T.prim_func(s_tir=True)
def main_kernel_1(n: T.int32):
@@ -384,7 +384,7 @@ def test_thread_extent_region_extracted_as_device_kernel():
@T.prim_func(s_tir=True)
def main(A: T.Buffer(16, "float32")):
T.func_attr({"target": T.target("cuda", host="llvm")})
- T.call_packed("main_kernel", A.data, 16)
+ T.call_ffi_kernel("main_kernel", A.data, 16,
launch_params=["threadIdx.x"])
@T.prim_func(s_tir=True)
def main_kernel(A_data: T.handle("float32")):
@@ -433,9 +433,14 @@ def test_cuda_launch_preserves_flag_metadata():
launch = after["main"].body.value
assert isinstance(launch, tvm.ir.Call)
+ assert launch.op == tvm.ir.Op.get("tirx.call_ffi_kernel")
+ assert isinstance(launch.attrs, tvm.tirx.CallFFIKernelAttr)
+ assert list(launch.attrs.launch_params) ==
list(kernel.attrs["tirx.kernel_launch_params"])
# Programmatic launch is flag-only and therefore adds no packed operand.
assert len(launch.args) == 3
assert int(launch.args[-1]) == 16
+ tvm.ir.assert_structural_equal(after,
tvm.ir.load_json(tvm.ir.save_json(after)))
+ tvm.ir.assert_structural_equal(after,
tvm.script.from_source(after.script()))
def test_cuda_required_block_size_coexists_with_launch_bounds():
@@ -515,7 +520,7 @@ def test_device_scope_region_extracted_as_device_kernel():
@T.prim_func(s_tir=True)
def main(A: T.Buffer(1, "float32")):
T.func_attr({"target": T.target("cuda", host="llvm")})
- T.call_packed("main_kernel", A.data)
+ T.call_ffi_kernel("main_kernel", A.data, launch_params=[])
@T.prim_func(s_tir=True)
def main_kernel(A_data: T.handle("float32")):
@@ -558,7 +563,7 @@ def test_lower_device_kernel_launch():
@T.prim_func(s_tir=True)
def main(A: T.Buffer(1, "float32")):
T.func_attr({"target": T.target("llvm")})
- T.call_packed("kernel", A.data)
+ T.call_ffi_kernel("kernel", A.data, launch_params=[])
@T.prim_func(s_tir=True)
def kernel(A_data: T.handle("float32")):
@@ -599,7 +604,7 @@ def test_externally_visible_kernel_launch():
@T.prim_func(s_tir=True)
def main(A: T.Buffer(1, "float32")):
T.func_attr({"target": T.target("llvm")})
- T.call_packed("kernel_by_another_name", A.data)
+ T.call_ffi_kernel("kernel_by_another_name", A.data,
launch_params=[])
@T.prim_func(s_tir=True)
def kernel(A_data: T.handle("float32")):
@@ -646,7 +651,7 @@ def test_collect_launch_parameter():
@T.prim_func(s_tir=True)
def main(A: T.Buffer(16, "float32")):
T.func_attr({"target": T.target("llvm")})
- T.call_packed("kernel", A.data, 16)
+ T.call_ffi_kernel("kernel", A.data, 16,
launch_params=["threadIdx.x"])
@T.prim_func(s_tir=True)
def kernel(A_data: T.handle("float32")):
@@ -729,7 +734,7 @@ def test_bind_before_thread_extent():
@T.prim_func(s_tir=True)
def main(A: T.Buffer(16, "float32"), n: T.int32):
T.func_attr({"target": T.target("llvm")})
- T.call_packed("kernel", A.data, n, n + 1)
+ T.call_ffi_kernel("kernel", A.data, n, n + 1,
launch_params=["threadIdx.x"])
@T.prim_func(s_tir=True)
def kernel(A_data: T.handle("float32"), n: T.int32):
diff --git
a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tma.py
b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tma.py
index f2fef0a861..cde0816723 100644
--- a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tma.py
+++ b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tma.py
@@ -26,7 +26,7 @@ import tvm_ffi
import tvm
import tvm.testing
-from tvm.ir import PointerType, PrimType, Range, StringImm
+from tvm.ir import PointerType, PrimType, Range
from tvm.script import tirx as T
from tvm.script.tirx import tile as Tx
from tvm.sym import Analyzer
@@ -123,12 +123,7 @@ class _EncodeCollector:
self.calls = []
def _visit_call(self, op):
- if (
- isinstance(op.op, tvm.ir.Op)
- and op.op.name == "tirx.tvm_call_packed"
- and isinstance(op.args[0], StringImm)
- and op.args[0].value == "runtime.cuTensorMapEncodeTiled"
- ):
+ if isinstance(op.op, tvm.ir.Op) and op.op.name ==
"tirx.tensormap_encode_tiled":
self.calls.append(op)
def visit_stmt(self, stmt):
@@ -271,8 +266,8 @@ def _collect_encodes(stmts):
def _encode_signature(call):
- rank = int(call.args[3])
- cursor = 5
+ rank = call.attrs.rank
+ cursor = 2
dims = tuple(call.args[cursor : cursor + rank])
cursor += rank
strides = tuple(call.args[cursor : cursor + rank - 1])
@@ -281,13 +276,17 @@ def _encode_signature(call):
cursor += rank
element_strides = tuple(call.args[cursor : cursor + rank])
cursor += rank
- enums = tuple(call.args[cursor : cursor + 4])
- cursor += 4
- forced_dtype = call.args[cursor] if cursor < len(call.args) else None
+ enums = (
+ call.attrs.interleave,
+ call.attrs.swizzle,
+ call.attrs.l2_promotion,
+ call.attrs.oob_fill,
+ )
+ forced_dtype = call.attrs.force_cu_dtype if call.attrs.force_cu_dtype >= 0
else None
return {
- "dtype": call.args[2].value,
+ "dtype": str(call.attrs.descriptor_dtype),
"rank": rank,
- "base": call.args[4],
+ "base": call.args[1],
"dims": dims,
"strides": strides,
"boxes": boxes,
@@ -1226,7 +1225,7 @@ def
test_auto_maximum_prefix_and_mixed_radix_issue_pointer():
assert _count_tma(impl).total == 512
-def test_copy_tma_host_init_dtype_is_string():
+def test_copy_tma_host_init_dtype_is_attribute():
"""The host-init encode call must carry the dtype as a StringImm, not a
packed enum -- ``_encode_signature`` reads ``args[2].value`` as a str."""
_, host_init_stmts, _ = _lower_direct(
@@ -1238,8 +1237,7 @@ def test_copy_tma_host_init_dtype_is_string():
dtype="float16",
)
encode_call = _collect_encodes(host_init_stmts)[0]
- assert isinstance(encode_call.args[2], StringImm)
- assert encode_call.args[2].value == "float16"
+ assert str(encode_call.attrs.descriptor_dtype) == "float16"
@pytest.mark.parametrize(
