[FIX][TIRx][CUDA] Support low-precision, tensor-map and cluster kernels in CUDA-host bundles

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 @@ def header_generator(tags):
 
     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 @@ def header_generator(tags):
 
     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 @@ def header_generator(tags):
 
     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 @@ def header_generator(tags):
 
     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;

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,30 +192,38 @@ 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];
       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) {
+      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_;
 };
 

tests/python/codegen/test_target_codegen_cuda.py

@@ -74,7 +74,7 @@ def add_one(A: T.Buffer((32,), "float32"), B: T.Buffer((32,), "float32")):
     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 run_and_check():
     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_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)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(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)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(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():