This is an automated email from the ASF dual-hosted git repository.
spectrometerHBH 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 448041da62 [TIRx][CUDA] Allow launch bounds with required block size
(#20226)
448041da62 is described below
commit 448041da62e7fedfa077d462f9e7c3f841567602
Author: Bohan Hou <[email protected]>
AuthorDate: Sat Aug 29 10:33:49 2026 -0400
[TIRx][CUDA] Allow launch bounds with required block size (#20226)
This change allows `tirx.required_block_size` to be combined with CUDA
launch-bounds attributes while preserving the existing incompatibility
with `tirx.max_registers`.
The combination is needed for kernels that require exact block/cluster
dimensions and also need a minimum blocks-per-SM launch bound. The CUDA
code generator now emits both attributes, and SplitHostDevice retains
the supported combination.
Tests:
- `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 src/backend/cuda/codegen/codegen_cuda.cc
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`
---
src/backend/cuda/codegen/codegen_cuda.cc | 14 ++++++++------
src/tirx/transform/split_host_device.cc | 7 ++-----
.../tirx-transform/test_tir_transform_split_host_device.py | 12 +++++++-----
tests/python/tirx/codegen/test_codegen_cuda.py | 14 ++++++++------
4 files changed, 25 insertions(+), 22 deletions(-)
diff --git a/src/backend/cuda/codegen/codegen_cuda.cc
b/src/backend/cuda/codegen/codegen_cuda.cc
index ef5f13c0a5..be2bfe9bae 100644
--- a/src/backend/cuda/codegen/codegen_cuda.cc
+++ b/src/backend/cuda/codegen/codegen_cuda.cc
@@ -247,11 +247,8 @@ void CodeGenCUDA::PrintExtraAttrs(const PrimFunc& f,
std::ostream& os) {
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";
+ TVM_FFI_ICHECK(!max_registers.has_value())
+ << tirx::attr::kRequiredBlockSize << " cannot be combined with 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>();
@@ -262,7 +259,12 @@ void CodeGenCUDA::PrintExtraAttrs(const PrimFunc& f,
std::ostream& os) {
<< 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
(!f->GetAttr<int64_t>(tirx::attr::kLaunchBoundsMinBlocksPerSM).has_value()) {
+
TVM_FFI_ICHECK(!f->GetAttr<int64_t>(tirx::attr::kLaunchBoundsMaxBlocksPerCluster).has_value())
+ << tirx::attr::kLaunchBoundsMaxBlocksPerCluster << " requires "
+ << tirx::attr::kLaunchBoundsMinBlocksPerSM;
+ return;
+ }
}
if (max_registers.has_value()) {
TVM_FFI_ICHECK_GT(max_registers.value(), 0);
diff --git a/src/tirx/transform/split_host_device.cc
b/src/tirx/transform/split_host_device.cc
index afb98b1690..8aba699d1f 100644
--- a/src/tirx/transform/split_host_device.cc
+++ b/src/tirx/transform/split_host_device.cc
@@ -99,11 +99,8 @@ 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";
+ TVM_FFI_ICHECK(!required_block_size_.has_value() ||
!max_registers_.has_value())
+ << tirx::attr::kRequiredBlockSize << " cannot be combined with maximum
registers";
return result;
}
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 ddc11bb1d4..2990efc725 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,7 +437,7 @@ def test_cuda_launch_preserves_flag_metadata():
assert int(launch.args[-1]) == 16
-def test_cuda_required_block_size_becomes_flag_only_launch_metadata():
+def test_cuda_required_block_size_coexists_with_launch_bounds():
@I.ir_module
class Before:
@T.prim_func(s_tir=True)
@@ -445,14 +445,16 @@ def
test_cuda_required_block_size_becomes_flag_only_launch_metadata():
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
+ with T.attr(0, "tirx.launch_bounds_min_blocks_per_sm", 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 int(kernel.attrs["tirx.launch_bounds_min_blocks_per_sm"]) == 1
assert list(kernel.attrs["tirx.kernel_launch_params"]) == [
"blockIdx.x",
"threadIdx.x",
diff --git a/tests/python/tirx/codegen/test_codegen_cuda.py
b/tests/python/tirx/codegen/test_codegen_cuda.py
index 1a79db2b63..46cec9d303 100644
--- a/tests/python/tirx/codegen/test_codegen_cuda.py
+++ b/tests/python/tirx/codegen/test_codegen_cuda.py
@@ -287,7 +287,7 @@ def test_tirx_required_block_size_emits_cuda_block_size():
assert "tirx.required_block_size" not in src
-def test_tirx_required_block_size_rejects_launch_controls():
+def test_tirx_required_block_size_emits_launch_bounds_when_requested():
@T.prim_func
def main(A: T.Buffer((4,), "int32")):
T.device_entry()
@@ -302,11 +302,13 @@ def
test_tirx_required_block_size_rejects_launch_controls():
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)
+ src, _ = _get_source(main)
+ assert (
+ 'extern "C" __global__ void __block_size__((128, 1, 1), (1, 1, 1)) '
+ "__launch_bounds__(128, 1) main_kernel" in src
+ )
+ assert "tirx.required_block_size" not in src
+ assert "tirx.launch_bounds_min_blocks_per_sm" not in src
def test_tirx_cuda_kernel_return_zero_codegen_is_void_early_return():
