diff --git a/python/tvm/backend/cuda/codegen/header.py b/python/tvm/backend/cuda/codegen/header.py index 55b9839426d7..edb682af0682 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 __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 __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 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 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 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 5fb4aee7e5c7..c910d84234fb 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 \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 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(); + if (print_ssa_form_ || !ptr || !ptr->element_type.as()) { + 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()->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()->value; - std::array axes; + // Match LaunchParamConfig's layout: grid, block, cluster and preferred cluster. + std::array 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 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(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 dimensions; + std::array 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 << "<<>>("; - 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 function_names_; }; diff --git a/tests/python/codegen/test_target_codegen_cuda.py b/tests/python/codegen/test_target_codegen_cuda.py index 6c3aa00edddc..ff8bd00876d3 100644 --- a/tests/python/codegen/test_target_codegen_cuda.py +++ b/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():