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 73e38d3f44 [TIRx][CUDA] Add register and cluster launch controls
(#20159)
73e38d3f44 is described below
commit 73e38d3f4447aa507cf795542aeffcc0948368b5
Author: Bohan Hou <[email protected]>
AuthorDate: Thu Aug 20 15:56:59 2026 -0400
[TIRx][CUDA] Add register and cluster launch controls (#20159)
## Summary
- add a `tirx.max_registers` PrimFunc attribute and lower it to CUDA 13
`__maxnreg__`
- reject combinations of the register cap with CUDA launch-bounds
attributes
- enable non-portable CUDA cluster sizes when a launch requests more
than eight CTAs
- keep launch attributes in fixed local storage and add CUDA codegen
coverage
## Motivation
SM100 TIRx kernels need an explicit register cap to preserve
source-level occupancy and instruction scheduling. Large reduction
kernels can also require 16-CTA clusters, which the CUDA driver rejects
unless the non-portable cluster-size attribute is enabled.
## Impact
TIRx kernels can opt into an exact CUDA register budget while retaining
the existing launch-bounds path for other specializations. CUDA launches
with cluster dimensions above eight are enabled only for kernels that
request them; existing launch behavior remains unchanged.
## Validation
- `cmake --build build --parallel`
- `python -m pytest tests/python/tirx/codegen/test_codegen_cuda.py -q`
(`201 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/tirx/transform/split_host_device.cc
tests/python/tirx/codegen/test_codegen_cuda.py`
---
include/tvm/tirx/function.h | 10 ++++++++
src/backend/cuda/codegen/codegen_cuda.cc | 9 +++++++
src/backend/cuda/runtime/cuda_module.cc | 24 ++++++++++++------
src/tirx/transform/split_host_device.cc | 21 ++++++++++++++++
tests/python/tirx/codegen/test_codegen_cuda.py | 35 ++++++++++++++++++++++++++
5 files changed, 92 insertions(+), 7 deletions(-)
diff --git a/include/tvm/tirx/function.h b/include/tvm/tirx/function.h
index d335ee9515..58912aefe4 100644
--- a/include/tvm/tirx/function.h
+++ b/include/tvm/tirx/function.h
@@ -293,6 +293,16 @@ constexpr const char* kLaunchBoundsMinBlocksPerSM =
"tirx.launch_bounds_min_bloc
constexpr const char* kLaunchBoundsMaxBlocksPerCluster =
"tirx.launch_bounds_max_blocks_per_cluster";
+/*!
+ * \brief CUDA maximum registers per thread.
+ *
+ * Emits the CUDA 13 ``__maxnreg__`` kernel qualifier. This attribute is
+ * mutually exclusive with the launch-bounds attributes.
+ *
+ * Type: IntImm
+ */
+constexpr const char* kMaxRegisters = "tirx.max_registers";
+
/*!
* \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 8b6f2fa2d5..15484391ca 100644
--- a/src/backend/cuda/codegen/codegen_cuda.cc
+++ b/src/backend/cuda/codegen/codegen_cuda.cc
@@ -243,6 +243,15 @@ void CodeGenCUDA::PrintExtraAttrs(const PrimFunc& f,
std::ostream& os) {
} else {
cluster_cta_x_is_linear_rank_ = false;
}
+ auto max_registers = f->GetAttr<int64_t>(tirx::attr::kMaxRegisters);
+ 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()
&&
+
!f->GetAttr<int64_t>(tirx::attr::kLaunchBoundsMaxBlocksPerCluster).has_value())
+ << tirx::attr::kMaxRegisters << " cannot be combined with CUDA launch
bounds";
+ os << " __maxnreg__(" << max_registers.value() << ")";
+ return;
+ }
if (const IntImmNode* const threadIdx_ext_int =
threadIdx_ext.as<IntImmNode>()) {
if (threadIdx_ext_int->value == 1) {
// unable to extract the number of threads per block, hence directly
return
diff --git a/src/backend/cuda/runtime/cuda_module.cc
b/src/backend/cuda/runtime/cuda_module.cc
index 5a598c87cc..0d210ac0da 100644
--- a/src/backend/cuda/runtime/cuda_module.cc
+++ b/src/backend/cuda/runtime/cuda_module.cc
@@ -231,9 +231,19 @@ class CUDAWrappedFunc {
<< "Failed to set the allowed dynamic shared memory size to " <<
wl.dyn_shmem_size;
}
}
+ if (wl.cluster_dim(0) * wl.cluster_dim(1) * wl.cluster_dim(2) > 8) {
+ CUresult result = cuFuncSetAttribute(
+ fcache_[device_id],
CU_FUNC_ATTRIBUTE_NON_PORTABLE_CLUSTER_SIZE_ALLOWED, 1);
+ if (result != CUDA_SUCCESS) {
+ TVM_FFI_THROW(InternalError)
+ << "Failed to allow non-portable CUDA cluster size (" <<
wl.cluster_dim(0) << ", "
+ << wl.cluster_dim(1) << ", " << wl.cluster_dim(2) << ")";
+ }
+ }
}
CUstream strm = static_cast<CUstream>(TVMFFIEnvGetStream(kDLCUDA,
device_id));
- std::vector<CUlaunchAttribute> attrs;
+ std::array<CUlaunchAttribute, 4> attrs{};
+ unsigned int num_attrs = 0;
// 1) Cluster
if (wl.cluster_dim(0) != 1 || wl.cluster_dim(1) != 1 || wl.cluster_dim(2)
!= 1) {
@@ -242,7 +252,7 @@ class CUDAWrappedFunc {
attr.value.clusterDim.x = wl.cluster_dim(0);
attr.value.clusterDim.y = wl.cluster_dim(1);
attr.value.clusterDim.z = wl.cluster_dim(2);
- attrs.push_back(attr);
+ attrs[num_attrs++] = attr;
}
// 1b) Preferred cluster (CUDA 12.8+,
cudaLaunchAttributePreferredClusterDimension)
@@ -253,7 +263,7 @@ class CUDAWrappedFunc {
attr.value.clusterDim.x = wl.preferred_cluster_dim(0);
attr.value.clusterDim.y = wl.preferred_cluster_dim(1);
attr.value.clusterDim.z = wl.preferred_cluster_dim(2);
- attrs.push_back(attr);
+ attrs[num_attrs++] = attr;
}
// 2) Programmatic stream serialization
@@ -261,7 +271,7 @@ class CUDAWrappedFunc {
CUlaunchAttribute attr{};
attr.id = CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION;
attr.value.programmaticStreamSerializationAllowed = 1;
- attrs.push_back(attr);
+ attrs[num_attrs++] = attr;
}
// 3) Cooperative
@@ -269,7 +279,7 @@ class CUDAWrappedFunc {
CUlaunchAttribute attr{};
attr.id = CU_LAUNCH_ATTRIBUTE_COOPERATIVE;
attr.value.cooperative = 1;
- attrs.push_back(attr);
+ attrs[num_attrs++] = attr;
}
// 4) Launch
@@ -282,8 +292,8 @@ class CUDAWrappedFunc {
config.blockDimZ = wl.block_dim(2);
config.sharedMemBytes = wl.dyn_shmem_size;
config.hStream = strm;
- config.attrs = attrs.empty() ? nullptr : attrs.data();
- config.numAttrs = static_cast<unsigned int>(attrs.size());
+ config.attrs = num_attrs == 0 ? nullptr : attrs.data();
+ config.numAttrs = num_attrs;
CUresult result = cuLaunchKernelEx(&config, fcache_[device_id], void_args,
nullptr);
diff --git a/src/tirx/transform/split_host_device.cc
b/src/tirx/transform/split_host_device.cc
index fc6890666a..b46650f864 100644
--- a/src/tirx/transform/split_host_device.cc
+++ b/src/tirx/transform/split_host_device.cc
@@ -90,15 +90,20 @@ class LaunchBoundsAttrExtractor : public StmtMutator {
Stmt Extract(Stmt stmt) {
min_blocks_per_sm_.reset();
max_blocks_per_cluster_.reset();
+ max_registers_.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 "
<< tirx::attr::kLaunchBoundsMinBlocksPerSM;
+ 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";
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_; }
private:
Stmt VisitStmt_(const AttrStmtNode* op) final {
@@ -126,12 +131,24 @@ class LaunchBoundsAttrExtractor : public StmtMutator {
}
max_blocks_per_cluster_ = max_blocks_per_cluster->value;
return VisitStmt(op->body);
+ } else if (op->attr_key == tirx::attr::kMaxRegisters) {
+ const auto* max_registers = op->value.as<IntImmNode>();
+ TVM_FFI_ICHECK(max_registers) << tirx::attr::kMaxRegisters << " expects
an integer value";
+ TVM_FFI_ICHECK_GT(max_registers->value, 0)
+ << tirx::attr::kMaxRegisters << " must be positive";
+ if (max_registers_.has_value()) {
+ TVM_FFI_ICHECK_EQ(max_registers_.value(), max_registers->value)
+ << "Conflicting " << tirx::attr::kMaxRegisters << " values";
+ }
+ max_registers_ = max_registers->value;
+ return VisitStmt(op->body);
}
return StmtMutator::VisitStmt_(op);
}
std::optional<int64_t> min_blocks_per_sm_;
std::optional<int64_t> max_blocks_per_cluster_;
+ std::optional<int64_t> max_registers_;
};
class HostDeviceSplitter : public StmtMutator {
@@ -255,6 +272,10 @@ class HostDeviceSplitter : public StmtMutator {
device_func = WithAttr(std::move(device_func),
tirx::attr::kLaunchBoundsMaxBlocksPerCluster,
launch_bounds_attr.max_blocks_per_cluster().value());
}
+ if (launch_bounds_attr.max_registers().has_value()) {
+ device_func = WithAttr(std::move(device_func),
tirx::attr::kMaxRegisters,
+ launch_bounds_attr.max_registers().value());
+ }
}
auto num_inputs = cur_func_->GetAttr<int64_t>(tvm::attr::kNumInputs);
if (num_inputs.has_value()) {
diff --git a/tests/python/tirx/codegen/test_codegen_cuda.py
b/tests/python/tirx/codegen/test_codegen_cuda.py
index daaa9f36d6..e4e2fce5ba 100644
--- a/tests/python/tirx/codegen/test_codegen_cuda.py
+++ b/tests/python/tirx/codegen/test_codegen_cuda.py
@@ -203,6 +203,41 @@ def
test_tirx_launch_bounds_max_blocks_per_cluster_emits_third_operand():
assert "tirx.launch_bounds_max_blocks_per_cluster" not in src
+def test_tirx_max_registers_attr_emits_cuda_maxnreg():
+ @T.prim_func
+ def main(A: T.Buffer((4,), "int32")):
+ T.device_entry()
+ T.attr({"tirx.max_registers": 92})
+ bx = T.cta_id([4])
+ tx = T.thread_id([128])
+ if tx == 0:
+ A[bx] = A[bx] + 1
+
+ src, _ = _get_source(main)
+ assert 'extern "C" __global__ void __maxnreg__(92) main_kernel' in src
+ assert "__launch_bounds__" not in src
+ assert "tirx.max_registers" not in src
+
+
+def test_tirx_max_registers_rejects_launch_bounds():
+ @T.prim_func
+ def main(A: T.Buffer((4,), "int32")):
+ T.device_entry()
+ T.attr(
+ {
+ "tirx.max_registers": 92,
+ "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"):
+ _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")):
