This is an automated email from the ASF dual-hosted git repository.
jinhongyii 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 0d9c9384b1 [TIRx][CUDA] Support exact required block dimensions
(#20223)
0d9c9384b1 is described below
commit 0d9c9384b1473e2a2e7d84f077f151bb3b9130bc
Author: Bohan Hou <[email protected]>
AuthorDate: Fri Aug 28 21:42:32 2026 -0400
[TIRx][CUDA] Support exact required block dimensions (#20223)
This adds an exact CUDA launch contract for TIRx kernels that need PTX
`.reqntid` rather than advisory launch bounds.
- preserve `tirx.required_block_size` through host/device splitting
- emit CUDA 13 `__block_size__` with static thread and cluster
dimensions
- pass the required-block flag through runtime metadata without adding a
packed operand
- launch with `CU_LAUNCH_KERNEL_REQUIRED_BLOCK_DIM` and avoid duplicate
cluster attributes
- reject incompatible launch-bound and max-register controls
Validation:
- `cmake --build build --parallel`
- `python -m pytest -q
tests/python/tirx-transform/test_tir_transform_split_host_device.py
tests/python/tirx/codegen/test_codegen_cuda.py` (225 passed)
- `pre-commit run --files include/tvm/tirx/function.h
src/backend/cuda/codegen/codegen_cuda.cc
src/backend/cuda/runtime/cuda_module.cc src/runtime/metadata.h
src/runtime/thread_storage_scope.h
src/tirx/transform/split_host_device.cc
tests/python/tirx-transform/test_tir_transform_split_host_device.py
tests/python/tirx/codegen/test_codegen_cuda.py`
- CUDA 13.2 SM100 runtime launch through a TIRx kernel with generated
PTX `.reqntid 128,1,1`
---
include/tvm/tirx/function.h | 13 ++++++
src/backend/cuda/codegen/codegen_cuda.cc | 20 +++++++++
src/backend/cuda/runtime/cuda_module.cc | 50 +++++++++++++++++-----
src/runtime/metadata.h | 2 +
src/runtime/thread_storage_scope.h | 6 +++
src/tirx/transform/split_host_device.cc | 34 ++++++++++++++-
.../test_tir_transform_split_host_device.py | 30 +++++++++++++
tests/python/tirx/codegen/test_codegen_cuda.py | 39 +++++++++++++++++
8 files changed, 183 insertions(+), 11 deletions(-)
diff --git a/include/tvm/tirx/function.h b/include/tvm/tirx/function.h
index 58912aefe4..523ca3b3d1 100644
--- a/include/tvm/tirx/function.h
+++ b/include/tvm/tirx/function.h
@@ -271,6 +271,7 @@ namespace attr {
*
* - tvm::runtime::launch_param::kUseProgramaticDependentLaunch
* - tvm::runtime::launch_param::kUseCooperativeLaunch
+ * - tvm::runtime::launch_param::kUseRequiredBlockDimension
*
* Flag-only launch attributes. These tags add no packed operand.
*
@@ -303,6 +304,18 @@ constexpr const char* kLaunchBoundsMaxBlocksPerCluster =
*/
constexpr const char* kMaxRegisters = "tirx.max_registers";
+/*!
+ * \brief Require CUDA to use the statically-declared block and cluster
dimensions.
+ *
+ * Emits the CUDA 13 ``__block_size__`` kernel qualifier. Unlike
+ * ``__launch_bounds__``, this is an exact launch contract: CUDA derives the
+ * PTX ``.reqntid`` directive from the thread extents, and interprets the
+ * launch grid in clusters using the cluster-CTA extents.
+ *
+ * Type: IntImm (must be 1)
+ */
+constexpr const char* kRequiredBlockSize = "tirx.required_block_size";
+
/*!
* \brief Whether to set noalias rule on the function arguments.
*
diff --git a/src/backend/cuda/codegen/codegen_cuda.cc
b/src/backend/cuda/codegen/codegen_cuda.cc
index 15484391ca..ef5f13c0a5 100644
--- a/src/backend/cuda/codegen/codegen_cuda.cc
+++ b/src/backend/cuda/codegen/codegen_cuda.cc
@@ -244,6 +244,26 @@ void CodeGenCUDA::PrintExtraAttrs(const PrimFunc& f,
std::ostream& os) {
cluster_cta_x_is_linear_rank_ = false;
}
auto max_registers = f->GetAttr<int64_t>(tirx::attr::kMaxRegisters);
+ auto required_block_size =
f->GetAttr<int64_t>(tirx::attr::kRequiredBlockSize);
+ if (required_block_size.has_value()) {
+ TVM_FFI_ICHECK_EQ(required_block_size.value(), 1);
+ TVM_FFI_ICHECK(!max_registers.has_value() &&
+
!f->GetAttr<int64_t>(tirx::attr::kLaunchBoundsMinBlocksPerSM).has_value() &&
+
!f->GetAttr<int64_t>(tirx::attr::kLaunchBoundsMaxBlocksPerCluster).has_value())
+ << tirx::attr::kRequiredBlockSize
+ << " cannot be combined with CUDA launch bounds or maximum registers";
+ const auto* tx = extractor.threadIdx_x_ext.as<IntImmNode>();
+ const auto* ty = extractor.threadIdx_y_ext.as<IntImmNode>();
+ const auto* tz = extractor.threadIdx_z_ext.as<IntImmNode>();
+ const auto* cx = extractor.clusterCtaIdx_x_ext.as<IntImmNode>();
+ const auto* cy = extractor.clusterCtaIdx_y_ext.as<IntImmNode>();
+ const auto* cz = extractor.clusterCtaIdx_z_ext.as<IntImmNode>();
+ TVM_FFI_ICHECK(tx && ty && tz && cx && cy && cz)
+ << tirx::attr::kRequiredBlockSize << " requires static thread and
cluster dimensions";
+ os << " __block_size__((" << tx->value << ", " << ty->value << ", " <<
tz->value << "), ("
+ << cx->value << ", " << cy->value << ", " << cz->value << "))";
+ return;
+ }
if (max_registers.has_value()) {
TVM_FFI_ICHECK_GT(max_registers.value(), 0);
TVM_FFI_ICHECK(!f->GetAttr<int64_t>(tirx::attr::kLaunchBoundsMinBlocksPerSM).has_value()
&&
diff --git a/src/backend/cuda/runtime/cuda_module.cc
b/src/backend/cuda/runtime/cuda_module.cc
index 3984bfde39..271f62a24b 100644
--- a/src/backend/cuda/runtime/cuda_module.cc
+++ b/src/backend/cuda/runtime/cuda_module.cc
@@ -250,9 +250,11 @@ class CUDAWrappedFunc {
CUstream strm = static_cast<CUstream>(TVMFFIEnvGetStream(kDLCUDA,
device_id));
std::array<CUlaunchAttribute, 4> attrs{};
unsigned int num_attrs = 0;
+ bool use_required_block_dimension =
launch_param_config_.use_required_block_dimension();
- // 1) Cluster
- if (launch_param_config_.use_cluster_launch()) {
+ // 1) Cluster. __block_size__ fixes this at compile time, and supplying a
+ // runtime cluster attribute for the same kernel is forbidden by CUDA.
+ if (!use_required_block_dimension &&
launch_param_config_.use_cluster_launch()) {
CUlaunchAttribute attr{};
attr.id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION;
attr.value.clusterDim.x = wl.cluster_dim(0);
@@ -262,8 +264,17 @@ class CUDAWrappedFunc {
}
// 1b) Preferred cluster (CUDA 12.8+,
cudaLaunchAttributePreferredClusterDimension)
- if (wl.preferred_cluster_dim(0) != 1 || wl.preferred_cluster_dim(1) != 1 ||
- wl.preferred_cluster_dim(2) != 1) {
+ if (use_required_block_dimension) {
+ for (int i = 0; i < 3; ++i) {
+ TVM_FFI_ICHECK(wl.preferred_cluster_dim(i) == 1 ||
+ wl.preferred_cluster_dim(i) == wl.cluster_dim(i))
+ << "CUDA required block dimensions cannot be combined with a
different preferred "
+ "cluster size";
+ }
+ }
+ if (!use_required_block_dimension &&
+ (wl.preferred_cluster_dim(0) != 1 || wl.preferred_cluster_dim(1) != 1
||
+ wl.preferred_cluster_dim(2) != 1)) {
CUlaunchAttribute attr{};
attr.id = CU_LAUNCH_ATTRIBUTE_PREFERRED_CLUSTER_DIMENSION;
attr.value.clusterDim.x = wl.preferred_cluster_dim(0);
@@ -290,12 +301,31 @@ class CUDAWrappedFunc {
// 4) Launch
CUlaunchConfig config{};
- config.gridDimX = wl.grid_dim(0);
- config.gridDimY = wl.grid_dim(1);
- config.gridDimZ = wl.grid_dim(2);
- config.blockDimX = wl.block_dim(0);
- config.blockDimY = wl.block_dim(1);
- config.blockDimZ = wl.block_dim(2);
+ if (use_required_block_dimension) {
+#if CUDA_VERSION >= 13000
+ for (int i = 0; i < 3; ++i) {
+ TVM_FFI_ICHECK_EQ(wl.grid_dim(i) % wl.cluster_dim(i), 0U)
+ << "CUDA required block dimension launch needs each logical
block-grid dimension "
+ "to be divisible by its compile-time cluster dimension";
+ }
+ config.gridDimX = wl.grid_dim(0) / wl.cluster_dim(0);
+ config.gridDimY = wl.grid_dim(1) / wl.cluster_dim(1);
+ config.gridDimZ = wl.grid_dim(2) / wl.cluster_dim(2);
+ config.blockDimX = CU_LAUNCH_KERNEL_REQUIRED_BLOCK_DIM;
+ config.blockDimY = 1;
+ config.blockDimZ = 1;
+#else
+ TVM_FFI_THROW(InternalError)
+ << "CUDA required block dimensions need CUDA Toolkit 13 or newer";
+#endif
+ } else {
+ config.gridDimX = wl.grid_dim(0);
+ config.gridDimY = wl.grid_dim(1);
+ config.gridDimZ = wl.grid_dim(2);
+ config.blockDimX = wl.block_dim(0);
+ config.blockDimY = wl.block_dim(1);
+ config.blockDimZ = wl.block_dim(2);
+ }
config.sharedMemBytes = wl.dyn_shmem_size;
config.hStream = strm;
config.attrs = num_attrs == 0 ? nullptr : attrs.data();
diff --git a/src/runtime/metadata.h b/src/runtime/metadata.h
index 0a8d96261e..ce40f81570 100644
--- a/src/runtime/metadata.h
+++ b/src/runtime/metadata.h
@@ -54,6 +54,8 @@ constexpr const char* kUseDynamicSharedMemoryTag =
"tirx.use_dyn_shared_memory";
constexpr const char* kUseProgramaticDependentLaunch =
"tirx.use_programtic_dependent_launch";
/*! \brief A tag to specify whether or not use cooperative launch */
constexpr const char* kUseCooperativeLaunch = "tirx.use_cooperative_launch";
+/*! \brief A tag to launch with CUDA's statically required block dimension */
+constexpr const char* kUseRequiredBlockDimension =
"tirx.use_required_block_dimension";
} // namespace launch_param
diff --git a/src/runtime/thread_storage_scope.h
b/src/runtime/thread_storage_scope.h
index 0443fc68a0..92be05a395 100644
--- a/src/runtime/thread_storage_scope.h
+++ b/src/runtime/thread_storage_scope.h
@@ -296,6 +296,8 @@ class LaunchParamConfig {
use_programmatic_dependent_launch_ = true;
} else if (tag == launch_param::kUseCooperativeLaunch) {
use_cooperative_launch_ = true;
+ } else if (tag == launch_param::kUseRequiredBlockDimension) {
+ use_required_block_dimension_ = true;
} else {
ThreadScope ts = ThreadScope::Create(tag);
if (ts.IsClusterCtaIdx()) {
@@ -337,6 +339,8 @@ class LaunchParamConfig {
bool use_cooperative_launch() const { return use_cooperative_launch_; }
+ bool use_required_block_dimension() const { return
use_required_block_dimension_; }
+
bool use_cluster_launch() const { return use_cluster_launch_; }
private:
@@ -352,6 +356,8 @@ class LaunchParamConfig {
bool use_programmatic_dependent_launch_{false};
/*! \brief Whether or not use cooperative launch. */
bool use_cooperative_launch_{false};
+ /*! \brief Whether CUDA should use the kernel's statically required block
dimension. */
+ bool use_required_block_dimension_{false};
/*! \brief Whether the kernel declares a cluster-to-CTA scope. */
bool use_cluster_launch_{false};
};
diff --git a/src/tirx/transform/split_host_device.cc
b/src/tirx/transform/split_host_device.cc
index b46650f864..afb98b1690 100644
--- a/src/tirx/transform/split_host_device.cc
+++ b/src/tirx/transform/split_host_device.cc
@@ -91,6 +91,7 @@ class LaunchBoundsAttrExtractor : public StmtMutator {
min_blocks_per_sm_.reset();
max_blocks_per_cluster_.reset();
max_registers_.reset();
+ required_block_size_.reset();
Stmt result = operator()(std::move(stmt));
TVM_FFI_ICHECK(!max_blocks_per_cluster_.has_value() ||
min_blocks_per_sm_.has_value())
<< tirx::attr::kLaunchBoundsMaxBlocksPerCluster << " requires "
@@ -98,12 +99,18 @@ class LaunchBoundsAttrExtractor : public StmtMutator {
TVM_FFI_ICHECK(!max_registers_.has_value() ||
(!min_blocks_per_sm_.has_value() &&
!max_blocks_per_cluster_.has_value()))
<< tirx::attr::kMaxRegisters << " cannot be combined with CUDA launch
bounds";
+ TVM_FFI_ICHECK(!required_block_size_.has_value() ||
+ (!min_blocks_per_sm_.has_value() &&
!max_blocks_per_cluster_.has_value() &&
+ !max_registers_.has_value()))
+ << tirx::attr::kRequiredBlockSize
+ << " cannot be combined with CUDA launch bounds or maximum registers";
return result;
}
std::optional<int64_t> min_blocks_per_sm() const { return
min_blocks_per_sm_; }
std::optional<int64_t> max_blocks_per_cluster() const { return
max_blocks_per_cluster_; }
std::optional<int64_t> max_registers() const { return max_registers_; }
+ std::optional<int64_t> required_block_size() const { return
required_block_size_; }
private:
Stmt VisitStmt_(const AttrStmtNode* op) final {
@@ -142,6 +149,18 @@ class LaunchBoundsAttrExtractor : public StmtMutator {
}
max_registers_ = max_registers->value;
return VisitStmt(op->body);
+ } else if (op->attr_key == tirx::attr::kRequiredBlockSize) {
+ const auto* required_block_size = op->value.as<IntImmNode>();
+ TVM_FFI_ICHECK(required_block_size)
+ << tirx::attr::kRequiredBlockSize << " expects an integer value";
+ TVM_FFI_ICHECK_EQ(required_block_size->value, 1)
+ << tirx::attr::kRequiredBlockSize << " must be 1";
+ if (required_block_size_.has_value()) {
+ TVM_FFI_ICHECK_EQ(required_block_size_.value(),
required_block_size->value)
+ << "Conflicting " << tirx::attr::kRequiredBlockSize << " values";
+ }
+ required_block_size_ = required_block_size->value;
+ return VisitStmt(op->body);
}
return StmtMutator::VisitStmt_(op);
}
@@ -149,6 +168,7 @@ class LaunchBoundsAttrExtractor : public StmtMutator {
std::optional<int64_t> min_blocks_per_sm_;
std::optional<int64_t> max_blocks_per_cluster_;
std::optional<int64_t> max_registers_;
+ std::optional<int64_t> required_block_size_;
};
class HostDeviceSplitter : public StmtMutator {
@@ -276,6 +296,10 @@ class HostDeviceSplitter : public StmtMutator {
device_func = WithAttr(std::move(device_func),
tirx::attr::kMaxRegisters,
launch_bounds_attr.max_registers().value());
}
+ if (launch_bounds_attr.required_block_size().has_value()) {
+ device_func = WithAttr(std::move(device_func),
tirx::attr::kRequiredBlockSize,
+
launch_bounds_attr.required_block_size().value());
+ }
}
auto num_inputs = cur_func_->GetAttr<int64_t>(tvm::attr::kNumInputs);
if (num_inputs.has_value()) {
@@ -363,6 +387,8 @@ class DeviceInfoCollector : public StmtVisitor {
}
}
}
+ collector.use_required_block_dimension_ =
+ func->GetAttr<int64_t>(tirx::attr::kRequiredBlockSize).value_or(0) ==
1;
collector(func->body);
@@ -373,6 +399,10 @@ class DeviceInfoCollector : public StmtVisitor {
if (collector.use_cooperative_launch_) {
collector.info_.launch_params.push_back(tvm::runtime::launch_param::kUseCooperativeLaunch);
}
+ if (collector.use_required_block_dimension_) {
+ collector.info_.launch_params.push_back(
+ tvm::runtime::launch_param::kUseRequiredBlockDimension);
+ }
// The dynamic shared memory is required to be the last of the kernel
// launch parameters. An explicit tirx.dyn_smem_bytes declaration wins;
// otherwise fall back to the size inferred from the allocation extent.
@@ -397,7 +427,8 @@ class DeviceInfoCollector : public StmtVisitor {
for (const ffi::String& param : collector.info_.launch_params) {
if (param == tvm::runtime::launch_param::kUseProgramaticDependentLaunch
||
- param == tvm::runtime::launch_param::kUseCooperativeLaunch) {
+ param == tvm::runtime::launch_param::kUseCooperativeLaunch ||
+ param == tvm::runtime::launch_param::kUseRequiredBlockDimension) {
continue;
}
collector.info_.launch_args.push_back(collector.GetArgument(param));
@@ -517,6 +548,7 @@ class DeviceInfoCollector : public StmtVisitor {
// Flag-only launch attributes requested by the original PrimFunc.
bool use_programmatic_dependent_launch_{false};
bool use_cooperative_launch_{false};
+ bool use_required_block_dimension_{false};
// Accumulated Bind definitions for inlining into extent/size expressions.
ffi::Map<Var, PrimExpr> bind_map_;
};
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 d04f2338a1..ddc11bb1d4 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
@@ -437,6 +437,36 @@ def test_cuda_launch_preserves_flag_metadata():
assert int(launch.args[-1]) == 16
+def test_cuda_required_block_size_becomes_flag_only_launch_metadata():
+ @I.ir_module
+ class Before:
+ @T.prim_func(s_tir=True)
+ def main(A: T.Buffer(4, "float32")):
+ T.func_attr({"target": T.target("cuda", host="llvm")})
+ T.attr(T.target("cuda"), "target", 0)
+ T.attr(0, "tirx.required_block_size", 1)
+ bx = T.launch_thread("blockIdx.x", 4)
+ tx = T.launch_thread("threadIdx.x", 128)
+ if tx == 0:
+ A[bx] = 0.0
+
+ after = tvm.tirx.transform.SplitHostDevice()(Before)
+ kernel = after["main_kernel"]
+ assert int(kernel.attrs["tirx.required_block_size"]) == 1
+ assert list(kernel.attrs["tirx.kernel_launch_params"]) == [
+ "blockIdx.x",
+ "threadIdx.x",
+ "tirx.use_required_block_dimension",
+ ]
+
+ launch = after["main"].body.value
+ assert isinstance(launch, tvm.ir.Call)
+ # The required-block flag reaches FunctionInfo metadata but adds no packed
operand.
+ assert len(launch.args) == 4
+ assert int(launch.args[-2]) == 4
+ assert int(launch.args[-1]) == 128
+
+
def test_cuda_launch_preserves_singleton_cluster_dimensions():
@I.ir_module
class Before:
diff --git a/tests/python/tirx/codegen/test_codegen_cuda.py
b/tests/python/tirx/codegen/test_codegen_cuda.py
index f9b47645e8..1a79db2b63 100644
--- a/tests/python/tirx/codegen/test_codegen_cuda.py
+++ b/tests/python/tirx/codegen/test_codegen_cuda.py
@@ -270,6 +270,45 @@ def test_tirx_max_registers_rejects_launch_bounds():
_get_source(main)
+def test_tirx_required_block_size_emits_cuda_block_size():
+ @T.prim_func
+ def main(A: T.Buffer((8,), "int32")):
+ T.device_entry()
+ T.attr({"tirx.required_block_size": 1})
+ bx, by = T.cta_id([4, 2])
+ _, cy = T.cta_id_in_cluster([1, 2])
+ tx = T.thread_id([128])
+ if tx == 0:
+ A[bx * 2 + by] = cy
+
+ src, _ = _get_source(main)
+ assert 'extern "C" __global__ void __block_size__((128, 1, 1), (1, 2, 1))
main_kernel' in src
+ assert "__launch_bounds__" not in src
+ assert "tirx.required_block_size" not in src
+
+
+def test_tirx_required_block_size_rejects_launch_controls():
+ @T.prim_func
+ def main(A: T.Buffer((4,), "int32")):
+ T.device_entry()
+ T.attr(
+ {
+ "tirx.required_block_size": 1,
+ "tirx.launch_bounds_min_blocks_per_sm": 1,
+ }
+ )
+ bx = T.cta_id([4])
+ tx = T.thread_id([128])
+ if tx == 0:
+ A[bx] = A[bx] + 1
+
+ with pytest.raises(
+ tvm.error.InternalError,
+ match="cannot be combined with CUDA launch bounds or maximum
registers",
+ ):
+ _get_source(main)
+
+
def test_tirx_cuda_kernel_return_zero_codegen_is_void_early_return():
@T.prim_func
def main(A: T.Buffer((4,), "int32")):
