This is an automated email from the ASF dual-hosted git repository.
MasterJH5574 pushed a commit to branch v0.27.0
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/v0.27.0 by this push:
new 7b5acb32bd [FIX][TIRx][CUDA] Support low-precision, tensor-map and
cluster kernels in CUDA-host bundles (#20489)
7b5acb32bd is described below
commit 7b5acb32bd271342a24130f4ef439ff6c2e52c6c
Author: Bohan Hou <[email protected]>
AuthorDate: Tue Sep 29 15:07:09 2026 -0400
[FIX][TIRx][CUDA] Support low-precision, tensor-map and cluster kernels in
CUDA-host bundles (#20489)
A CUDA-host bundle (`tvm.backend.cuda.export_cuda_host`, #20395)
compiles device and host code as one NVCC C++ translation unit. Building
the native kernels of mlc-ai/TIRx-kernels this way exposed several
patterns that the bundle could not compile or launch:
- **Low-precision types in the host pass.** The CUDA device header
included `cuda_fp16.h`, `cuda_bf16.h`, `cuda_fp8.h`, `cuda_fp6.h` and
`cuda_fp4.h` (and defined the `fp8_e4_t`-style aliases) only under
`defined(__CUDA_ARCH__)`. NVCC's host pass then fails on kernel
signatures that use `half` or `nv_bfloat16`. The guards now also admit
the host pass (`!defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= N`); NVRTC
and device-only NVCC compilation always define `__CUDA_ARCH__` and are
unchanged.
- **Tensor-map parameters.** The host wrapper bound a `T.TensorMap()`
parameter as `CUtensorMap* x = ((void*)...)`, which C accepts and C++
rejects. It now casts explicitly.
- **Argument types at the launch.** Host C code spells some device types
differently (bfloat16 is `uint16_t*` on the host, `nv_bfloat16*` in the
kernel), so the direct `<<<>>>` launch did not type-check. Launches now
go through a small helper that converts each argument to the kernel's
own parameter type and calls `cudaLaunchKernelEx`.
- **Launch attributes.** `clusterCtaIdx.*`, `preferredClusterCtaIdx.*`,
programmatic dependent launch and cooperative launch were rejected. They
now become the corresponding `cudaLaunchAttribute`s, following
`CUDAWrappedFunc` in the CUDA runtime module (including
`cudaFuncAttributeNonPortableClusterSizeAllowed` for cluster launches
and omitting a unit preferred cluster). Required block dimensions remain
unsupported and are still diagnosed.
---
python/tvm/backend/cuda/codegen/header.py | 13 ++-
src/backend/cuda/codegen/codegen_cuda_host.cc | 127 ++++++++++++++++++---
tests/python/codegen/test_target_codegen_cuda.py | 138 ++++++++++++++++++++++-
3 files changed, 256 insertions(+), 22 deletions(-)
diff --git a/python/tvm/backend/cuda/codegen/header.py
b/python/tvm/backend/cuda/codegen/header.py
index 55b9839426..edb682af06 100644
--- a/python/tvm/backend/cuda/codegen/header.py
+++ b/python/tvm/backend/cuda/codegen/header.py
@@ -113,9 +113,12 @@ def header_generator(tags):
#endif
"""
+ # The low-precision type headers below are also visible to NVCC's host pass
+ # (no __CUDA_ARCH__), which parses kernel signatures when this source
shares
+ # a translation unit with host code, e.g.
tvm.backend.cuda.export_cuda_host.
if "fp16" in tags:
header += R"""
-#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 530)
+#if !defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 530)
#include <cuda_fp16.h>
__device__ half max(half a, half b)
{
@@ -173,7 +176,7 @@ CUDA_UNSUPPORTED_HALF_MATH_UNARY(hexp, exp)
if "bf16" in tags:
header += R"""
-#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
+#if !defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 800)
#include <cuda_bf16.h>
__device__ nv_bfloat16 max(nv_bfloat16 a, nv_bfloat16 b)
{
@@ -224,7 +227,7 @@ CUDA_UNSUPPORTED_HALF_MATH_UNARY(herf, erf)
if "fp8" in tags:
header += R"""
-#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 890)
+#if !defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 890)
#include <cuda_fp8.h>
using fp8_e4_t = __nv_fp8_e4m3;
using fp8_e4x2_t = __nv_fp8x2_e4m3;
@@ -258,7 +261,7 @@ struct fp8_e8x16_t {
if "fp6" in tags:
header += R"""
-#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
+#if !defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 1000)
#include <cuda_fp6.h>
using fp6_e2_t = __nv_fp6_e2m3;
using fp6_e2x2_t = __nv_fp6x2_e2m3;
@@ -283,7 +286,7 @@ struct fp6_e3x16_t {
if "fp4" in tags:
header += R"""
-#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
+#if !defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 800)
#include <cuda_fp4.h>
using fp4_e2_t = __nv_fp4_e2m1;
using fp4_e2x2_t = __nv_fp4x2_e2m1;
diff --git a/src/backend/cuda/codegen/codegen_cuda_host.cc
b/src/backend/cuda/codegen/codegen_cuda_host.cc
index 5fb4aee7e5..c910d84234 100644
--- a/src/backend/cuda/codegen/codegen_cuda_host.cc
+++ b/src/backend/cuda/codegen/codegen_cuda_host.cc
@@ -62,6 +62,16 @@ class CodeGenCUDAHost : public CodeGenCHost {
<< " TVMFFIErrorSetRaisedFromCStrParts(\"CUDAError\", parts,
2);\n"
<< " cudaGetLastError();\n"
<< " return -1;\n}\n";
+ // Host C types spell some device types differently (bfloat16 is uint16_t
on
+ // the host), so convert each argument to the kernel's own parameter type.
+ launch_ = name_supply_->FreshName("tvm_cuda_host_launch");
+ decl_stream << "template <typename... Params, typename... Args>\n"
+ << "static cudaError_t " << launch_
+ << "(void (*kernel)(Params...), const cudaLaunchConfig_t*
config, "
+ << "Args... args) {\n"
+ << " static_assert(sizeof...(Params) == sizeof...(Args), "
+ << "\"kernel argument count mismatch\");\n"
+ << " return cudaLaunchKernelEx(config, kernel,
((Params)args)...);\n}\n";
}
using CodeGenCHost::PrintType;
@@ -102,6 +112,20 @@ class CodeGenCUDAHost : public CodeGenCHost {
ffi::Array<ffi::String> GetFunctionNames() { return function_names_; }
+ void Dispatch_(const BindNode* op) override {
+ // C accepts the implicit void* conversion emitted for descriptor handles;
+ // the CUDA translation unit is C++, which requires an explicit cast.
+ auto* ptr = op->var->ty.as<PointerTypeNode>();
+ if (print_ssa_form_ || !ptr ||
!ptr->element_type.as<tirx::TensorMapTypeNode>()) {
+ CodeGenC::Dispatch_(op);
+ return;
+ }
+ std::string value = PrintExpr(op->value);
+ PrintIndent();
+ stream << "CUtensorMap* " << AllocVarID(op->var.get()) << " =
(CUtensorMap*)(" << value
+ << ");\n";
+ }
+
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") {
@@ -168,10 +192,13 @@ class CodeGenCUDAHost : public CodeGenCHost {
<< "cuda_host kernel calls require a string kernel symbol";
const auto& symbol = op->args[0].as<StringImmNode>()->value;
- std::array<int, 6> axes;
+ // Match LaunchParamConfig's layout: grid, block, cluster and preferred
cluster.
+ std::array<int, 12> axes;
axes.fill(-1);
int shared_memory = -1;
size_t num_launch_values = 0;
+ bool use_programmatic_dependent_launch = false;
+ bool use_cooperative_launch = false;
std::unordered_set<std::string> seen;
for (size_t i = 0; i < attr->launch_params.size(); ++i) {
std::string tag = attr->launch_params[i];
@@ -179,19 +206,24 @@ class CodeGenCUDAHost : public CodeGenCHost {
<< "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) {
+ if (tag == runtime::launch_param::kUseProgramaticDependentLaunch) {
+ use_programmatic_dependent_launch = true;
+ } else if (tag == runtime::launch_param::kUseCooperativeLaunch) {
+ use_cooperative_launch = true;
+ } else if (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 {
+ static const char* kScopes[] = {"blockIdx.", "threadIdx.",
"clusterCtaIdx.",
+ "preferredClusterCtaIdx."};
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;
+ for (int rank = 0; rank < 4; ++rank) {
+ for (int j = 0; j < 3; ++j) {
+ if (tag == std::string(kScopes[rank]) + "xyz"[j]) axis = rank * 3
+ j;
+ }
}
TVM_FFI_CHECK_GE(axis, 0, ValueError)
<< "cuda_host does not support launch parameter: " << tag;
@@ -229,7 +261,7 @@ class CodeGenCUDAHost : public CodeGenCHost {
arguments.push_back(name);
}
auto launch_arg = [&](int index) { return arguments[launch_begin - 1 +
index]; };
- std::array<std::string, 6> dimensions;
+ std::array<std::string, 12> dimensions;
for (size_t i = 0; i < axes.size(); ++i) {
if (axes[i] < 0) {
dimensions[i] = "1";
@@ -262,16 +294,78 @@ class CodeGenCUDAHost : public CodeGenCHost {
PrintIndent();
stream << "}\n";
}
+ auto has_axis = [&](int rank) {
+ return axes[rank * 3] >= 0 || axes[rank * 3 + 1] >= 0 || axes[rank * 3 +
2] >= 0;
+ };
+ bool use_cluster = has_axis(2);
+ bool use_preferred_cluster = has_axis(3);
+ if (use_cluster) {
+ CheckError("cudaFuncSetAttribute(::" + std::string(symbol) +
+ ", cudaFuncAttributeNonPortableClusterSizeAllowed, 1)");
+ }
+ std::string config = name_supply_->FreshName("cuda_config");
+ PrintIndent();
+ stream << "cudaLaunchConfig_t " << config << " = {};\n";
+ PrintIndent();
+ stream << config << ".gridDim = dim3(" << dimensions[0] << ", " <<
dimensions[1] << ", "
+ << dimensions[2] << ");\n";
+ PrintIndent();
+ stream << config << ".blockDim = dim3(" << dimensions[3] << ", " <<
dimensions[4] << ", "
+ << dimensions[5] << ");\n";
+ PrintIndent();
+ stream << config << ".dynamicSmemBytes = " << bytes << ";\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 << config << ".stream = " << cuda_stream << ";\n";
+ int max_attrs = use_cluster + use_preferred_cluster +
use_programmatic_dependent_launch +
+ use_cooperative_launch;
+ if (max_attrs > 0) {
+ std::string attrs = name_supply_->FreshName("cuda_attrs");
+ PrintIndent();
+ stream << "cudaLaunchAttribute " << attrs << "[" << max_attrs << "] =
{};\n";
+ PrintIndent();
+ stream << config << ".attrs = " << attrs << ";\n";
+ auto set_dimension = [&](const std::string& id, const std::string&
field, int rank) {
+ PrintIndent();
+ stream << attrs << "[" << config << ".numAttrs].id = " << id << ";\n";
+ for (int j = 0; j < 3; ++j) {
+ PrintIndent();
+ stream << attrs << "[" << config << ".numAttrs].val." << field <<
"." << "xyz"[j] << " = "
+ << dimensions[rank * 3 + j] << ";\n";
+ }
+ PrintIndent();
+ stream << "++" << config << ".numAttrs;\n";
+ };
+ if (use_cluster) {
+ set_dimension("cudaLaunchAttributeClusterDimension", "clusterDim", 2);
+ }
+ if (use_preferred_cluster) {
+ // Match the CUDA runtime module: a unit preferred cluster adds no
attribute.
+ PrintIndent();
+ stream << "if (" << dimensions[9] << " != 1 || " << dimensions[10] <<
" != 1 || "
+ << dimensions[11] << " != 1) {\n";
+ int scope = BeginScope();
+ set_dimension("cudaLaunchAttributePreferredClusterDimension",
"preferredClusterDim", 3);
+ EndScope(scope);
+ PrintIndent();
+ stream << "}\n";
+ }
+ auto set_flag = [&](const std::string& id, const std::string& field) {
+ PrintIndent();
+ stream << attrs << "[" << config << ".numAttrs].id = " << id << ";\n";
+ PrintIndent();
+ stream << attrs << "[" << config << ".numAttrs++].val." << field << "
= 1;\n";
+ };
+ if (use_programmatic_dependent_launch) {
+ set_flag("cudaLaunchAttributeProgrammaticStreamSerialization",
+ "programmaticStreamSerializationAllowed");
+ }
+ if (use_cooperative_launch) {
+ set_flag("cudaLaunchAttributeCooperative", "cooperative");
+ }
}
- stream << ");\n";
- CheckError("cudaGetLastError()", true);
+ std::string launch = launch_ + "(::" + std::string(symbol) + ", &" +
config;
+ for (size_t i = 0; i + 1 < launch_begin; ++i) launch += ", " +
arguments[i];
+ CheckError(launch + ")", true);
os << "0";
}
@@ -375,6 +469,7 @@ class CodeGenCUDAHost : public CodeGenCHost {
}
std::string check_error_;
+ std::string launch_;
ffi::Array<ffi::String> function_names_;
};
diff --git a/tests/python/codegen/test_target_codegen_cuda.py
b/tests/python/codegen/test_target_codegen_cuda.py
index 6c3aa00edd..ff8bd00876 100644
--- a/tests/python/codegen/test_target_codegen_cuda.py
+++ b/tests/python/codegen/test_target_codegen_cuda.py
@@ -74,7 +74,7 @@ def test_cuda_host_bundle(tmp_path):
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("<<<")
+ assert source.index("__global__") < source.index("cudaLaunchKernelEx")
library = tvm_ffi.cpp.build_inline(
name="cuda_host_add_one",
cuda_sources=source,
@@ -95,6 +95,142 @@ def test_cuda_host_bundle(tmp_path):
tvm.testing.run_with_gpu_lock(run_and_check)
[email protected]
[email protected](not env.has_cuda(), reason="need cuda")
+def test_cuda_host_bundle_bf16_cluster(tmp_path):
+ """bfloat16 pointers and cluster launch attributes in a CUDA-host
bundle."""
+ from shutil import which
+
+ import tvm_ffi.cpp
+
+ from tvm.backend.cuda import export_cuda_host
+
+ ml_dtypes = pytest.importorskip("ml_dtypes")
+ if which("nvcc") is None:
+ pytest.skip("CUDA-host compilation requires NVCC")
+ if int(tvm.cuda(0).compute_version.split(".")[0]) < 9:
+ pytest.skip("thread block clusters require SM90 or newer")
+
+ @T.prim_func
+ def main(A_ptr: T.handle, B_ptr: T.handle, R_ptr: T.handle):
+ A = T.match_buffer(A_ptr, (128,), "bfloat16")
+ B = T.match_buffer(B_ptr, (128,), "bfloat16")
+ R = T.match_buffer(R_ptr, (4,), "int32")
+ T.device_entry()
+ for cx in T.thread_binding(2, thread="clusterCtaIdx.x"):
+ for bx in T.thread_binding(4, thread="blockIdx.x"):
+ for tx in T.thread_binding(32, thread="threadIdx.x"):
+ B[bx * 32 + tx] = A[bx * 32 + tx] + T.bfloat16(1)
+ if tx == 0:
+ R[bx] = cx
+
+ arch = env.cuda_arch()
+ target = tvm.target.Target({"kind": "cuda", "arch": arch},
host="cuda_host")
+ with target:
+ built = tvm.compile(tvm.IRModule({"main": main}), target=target,
tir_pipeline="tirx").mod
+ source = export_cuda_host(built)
+ assert "cudaLaunchAttributeClusterDimension" in source
+ library = tvm_ffi.cpp.build_inline(
+ name="cuda_host_bf16_cluster",
+ cuda_sources=source,
+ extra_cuda_cflags=[f"-arch={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(128).astype(ml_dtypes.bfloat16)
+ a = tvm.runtime.tensor(values, dev)
+ b = tvm.runtime.empty((128,), "bfloat16", dev)
+ r = tvm.runtime.empty((4,), "int32", dev)
+ loaded["main"](a, b, r)
+ tvm.testing.assert_allclose(b.numpy().astype("float32"),
np.arange(128) + 1)
+ # Each CTA reports its rank within a two-CTA cluster.
+ np.testing.assert_array_equal(r.numpy(), [0, 1, 0, 1])
+
+ tvm.testing.run_with_gpu_lock(run_and_check)
+
+
[email protected]
[email protected](not env.has_cuda(), reason="need cuda")
+def test_cuda_host_bundle_programmatic_dependent_launch(tmp_path):
+ """A programmatic-dependent-launch flag becomes a launch attribute."""
+ 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")
+ if int(tvm.cuda(0).compute_version.split(".")[0]) < 9:
+ pytest.skip("programmatic dependent launch requires SM90 or newer")
+
+ @T.prim_func(s_tir=True)
+ def add_one(A: T.Buffer((32,), "float32"), B: T.Buffer((32,), "float32")):
+ T.func_attr({"tirx.kernel_launch_params":
["tirx.use_programtic_dependent_launch"]})
+ for tx in T.thread_binding(32, "threadIdx.x"):
+ B[tx] = A[tx] + T.float32(1)
+
+ target = tvm.target.Target("cuda", host="cuda_host")
+ source = export_cuda_host(tvm.compile(add_one, target=target).mod)
+ assert "cudaLaunchAttributeProgrammaticStreamSerialization" in source
+ library = tvm_ffi.cpp.build_inline(
+ name="cuda_host_programmatic_dependent_launch",
+ 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)
+
+
[email protected]
[email protected](not env.has_cuda(), reason="need cuda")
+def test_cuda_host_bundle_tensor_map_parameter(tmp_path, monkeypatch):
+ """A tensor-map parameter handle compiles in the C++ host wrapper."""
+ 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
+ def main(A_map: T.TensorMap()):
+ T.device_entry()
+ tx = T.thread_id([32])
+ if tx == 0:
+ T.evaluate(T.address_of(A_map))
+
+ target = tvm.target.Target({"kind": "cuda", "arch": "sm_90a"},
host="cuda_host")
+ with target:
+ built = tvm.compile(tvm.IRModule({"main": main}), target=target,
tir_pipeline="tirx").mod
+ source = export_cuda_host(built)
+ assert "(CUtensorMap*)" in source
+ monkeypatch.setenv("TVM_FFI_CUDA_ARCH_LIST", "9.0a")
+ tvm_ffi.cpp.build_inline(
+ name="cuda_host_tensor_map_parameter",
+ cuda_sources=source,
+ build_directory=str(tmp_path),
+ backend="cuda",
+ )
+
+
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
def test_cuda_vectorize_add():
