This is an automated email from the ASF dual-hosted git repository.
tlopex 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 e1d85f64c0 [REFACTOR][TIRx] Declare dynamic shared memory with a CUDA
leaf op (#20563)
e1d85f64c0 is described below
commit e1d85f64c0d8027da42713eede5574755fc7e2ed
Author: Tianqi Chen <[email protected]>
AuthorDate: Tue Oct 6 10:42:24 2026 -0400
[REFACTOR][TIRx] Declare dynamic shared memory with a CUDA leaf op (#20563)
Replace the dynamic shared-memory size attribute with
`T.cuda.dyn_smem_bytes(bytes)`, a CUDA declaration whose signature
requires an integer constant. `SMEMPool.commit()` emits the declaration,
and host/device splitting collects its byte count for the runtime launch
before removing it from the device body.
Preserve pool high-water accounting, explicit-size precedence,
allocation-extent fallback, and dynamic shared-memory launch argument
ordering. Migrate direct declaration sites to the CUDA spelling.
---
python/tvm/backend/cuda/lang/alloc_pool.py | 9 ++---
src/backend/cuda/op/target_builtin.cc | 9 +++++
src/tirx/transform/split_host_device.cc | 40 +++++++++++-----------
.../tile_primitive/cuda/copy_async/test_tma.py | 8 ++---
4 files changed, 36 insertions(+), 30 deletions(-)
diff --git a/python/tvm/backend/cuda/lang/alloc_pool.py
b/python/tvm/backend/cuda/lang/alloc_pool.py
index 119b4000d7..6f305f731b 100644
--- a/python/tvm/backend/cuda/lang/alloc_pool.py
+++ b/python/tvm/backend/cuda/lang/alloc_pool.py
@@ -477,7 +477,7 @@ class SMEMPool:
self.max_offset = max(self.max_offset, self.offset)
def commit(self, size=None):
- """Emit pool size annotation into the IR.
+ """Emit the dynamic shared memory byte-count declaration into the IR.
Must be called after all ``alloc()`` / ``move_base_to()`` calls.
@@ -495,10 +495,7 @@ class SMEMPool:
f"the pool high-water mark ({self.max_offset})"
)
import tvm.tirx
+ from tvm.backend.cuda.op import dyn_smem_bytes
from tvm.tirx.script.ir_builder.parser_protocol import add_to_parent
- add_to_parent(
- tvm.tirx.AttrStmt(
- 0, "tirx.dyn_smem_bytes", tvm.tirx.IntImm("int64", resolved),
tvm.tirx.Evaluate(0)
- )
- )
+
add_to_parent(tvm.tirx.Evaluate(dyn_smem_bytes(tvm.tirx.IntImm("int64",
resolved))))
diff --git a/src/backend/cuda/op/target_builtin.cc
b/src/backend/cuda/op/target_builtin.cc
index 7eaa19106e..0c33c27bb5 100644
--- a/src/backend/cuda/op/target_builtin.cc
+++ b/src/backend/cuda/op/target_builtin.cc
@@ -219,6 +219,15 @@ void RegisterCudaTargetBuiltins() {
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
+ OpDef("tirx.cuda.dyn_smem_bytes",
+ "Declare the dynamic shared memory size in bytes for the kernel
launch.")
+ .signature(sig::arg<IntImm>("bytes", "The constant byte count."))
+ .set_attr<TFixedReturnType>("TFixedReturnType", PrimType::Void())
+ .set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.cuda.dyn_smem_bytes"))
+ .set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
+ .set_attr<TCallEffectKind>("TCallEffectKind",
+
static_cast<int64_t>(CallEffectKind::kEmbedInfo));
+
RegisterDeviceIntrinsics();
}
diff --git a/src/tirx/transform/split_host_device.cc
b/src/tirx/transform/split_host_device.cc
index 5858c380f7..f65c79ba63 100644
--- a/src/tirx/transform/split_host_device.cc
+++ b/src/tirx/transform/split_host_device.cc
@@ -444,7 +444,7 @@ class DeviceInfoCollector : public StmtExprVisitor {
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;
+ // launch parameters. An explicit tirx.cuda.dyn_smem_bytes declaration
wins;
// otherwise fall back to the size inferred from the allocation extent.
// A zero-extent allocation is a pool-style extern placeholder, so having
// neither a declaration nor a usable extent is an authoring error.
@@ -453,7 +453,7 @@ class DeviceInfoCollector : public StmtExprVisitor {
TVM_FFI_ICHECK(!(inferred && inferred->value == 0))
<< "PrimFunc " << gvar->name_hint
<< " allocates dynamic shared memory with a placeholder extent but
does not declare "
- "its size; annotate the kernel with tirx.dyn_smem_bytes
(SMEMPool.commit() emits "
+ "its size; annotate the kernel with tirx.cuda.dyn_smem_bytes
(SMEMPool.commit() emits "
"it).";
collector->dyn_shmem_size = collector->inferred_shmem_size_;
}
@@ -482,7 +482,7 @@ class DeviceInfoCollector : public StmtExprVisitor {
if (launch_param ==
tvm::runtime::launch_param::kUseDynamicSharedMemoryTag) {
TVM_FFI_ICHECK(dyn_shmem_size.has_value())
<< "Compute kernel requires launch parameter \"" << launch_param
- << "\", but PrimFunc did not declare tirx.dyn_smem_bytes.";
+ << "\", but PrimFunc did not declare tirx.cuda.dyn_smem_bytes.";
return dyn_shmem_size.value();
}
@@ -518,17 +518,15 @@ class DeviceInfoCollector : public StmtExprVisitor {
return StmtExprVisitor::Visit_(op);
}
- ffi::Optional<VisitInterrupt> Visit_(const AttrStmtNode* op) final {
- if (op->attr_key == "tirx.dyn_smem_bytes") {
- // Kernel-level declaration of the dynamic shared memory launch size.
- // The backing shared.dyn allocation is an extern placeholder; this
- // attribute is the single source of truth for the launch parameter.
+ ffi::Optional<VisitInterrupt> Visit_(const EvaluateNode* op) final {
+ static const Op dyn_smem_bytes = Op::Get("tirx.cuda.dyn_smem_bytes");
+ if (const auto* call = op->value.as<CallNode>(); call &&
call->op.same_as(dyn_smem_bytes)) {
+ // The declaration supplies the launch size even when the backing
+ // shared.dyn allocation is an extern placeholder.
TVM_FFI_ICHECK(!dyn_shmem_size.has_value())
- << "Only one tirx.dyn_smem_bytes declaration is allowed per kernel.";
- TVM_FFI_ICHECK(op->value.as<IntImmNode>()) << "tirx.dyn_smem_bytes must
be an IntImm";
- dyn_shmem_size = op->value.as_or_throw<PrimExpr>();
+ << "Only one tirx.cuda.dyn_smem_bytes declaration is allowed per
kernel.";
+ dyn_shmem_size = call->args[0].as_or_throw<IntImm>();
}
-
return StmtExprVisitor::Visit_(op);
}
@@ -569,7 +567,7 @@ class DeviceInfoCollector : public StmtExprVisitor {
saw_dyn_shared_alloc_ = true;
// Fallback launch size inferred from the allocation extent, used when
- // no tirx.dyn_smem_bytes declaration is present (e.g. s_tir schedules
+ // no tirx.cuda.dyn_smem_bytes declaration is present (e.g. s_tir
schedules
// allocate shared.dyn with a concrete extent). A zero extent is a
// pool-style extern placeholder and carries no size information.
tvm::Tuple shape = call->args[0].as_or_throw<tvm::Tuple>();
@@ -603,7 +601,7 @@ class DeviceInfoCollector : public StmtExprVisitor {
// Whether a shared.dyn allocation was seen.
bool saw_dyn_shared_alloc_{false};
// Launch size inferred from the allocation extent (fallback when no
- // tirx.dyn_smem_bytes declaration is present).
+ // tirx.cuda.dyn_smem_bytes declaration is present).
ffi::Optional<PrimExpr> inferred_shmem_size_{std::nullopt};
// Flag-only launch attributes requested by the original PrimFunc.
bool use_programmatic_dependent_launch_{false};
@@ -729,7 +727,7 @@ class DeviceKernelMutator : public StmtExprMutator {
write_ptr->body = ReturnRemover::Apply(write_ptr->body.value(),
!preserve_early_returns);
// The dyn-smem size declaration was consumed by DeviceInfoCollector;
// it has no meaning inside the kernel body.
- class StripDynSmemAttr : public StmtExprMutator {
+ class StripDynSmemDeclaration : public StmtExprMutator {
public:
using StmtExprMutator::Mutate;
using StmtExprMutator::Mutate_;
@@ -738,14 +736,16 @@ class DeviceKernelMutator : public StmtExprMutator {
return StmtExprMutator::Mutate(input, inplace_mode);
}
- UnchangedOr<Stmt> Mutate_(const AttrStmtNode* op, InplaceMode
inplace_mode) final {
- if (op->attr_key == "tirx.dyn_smem_bytes") {
- return Mutate(op->body, inplace_mode).ValueOrUnchanged(op->body);
+ UnchangedOr<Stmt> Mutate_(const EvaluateNode* op, InplaceMode
inplace_mode) final {
+ static const Op dyn_smem_bytes =
Op::Get("tirx.cuda.dyn_smem_bytes");
+ if (const auto* call = op->value.as<CallNode>();
+ call && call->op.same_as(dyn_smem_bytes)) {
+ return Evaluate(0);
}
- return StmtExprMutator::Mutate_(op, inplace_mode);
+ return ffi::Unchanged();
}
};
- write_ptr->body = ffi::make_object<StripDynSmemAttr>()
+ write_ptr->body = ffi::make_object<StripDynSmemDeclaration>()
->Mutate(write_ptr->body, InplaceMode::kAllow)
.ValueOrUnchanged(write_ptr->body);
}
diff --git
a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tma.py
b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tma.py
index 1836719c1c..86636c20dd 100644
--- a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tma.py
+++ b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tma.py
@@ -1532,7 +1532,7 @@ def rank_change(A: T.Tensor((8, 8), 'float16')):
T.cta_id([1])
tid = T.thread_id([1])
dyn = T.alloc_tensor((65,), "uint64", scope="shared.dyn")
- T.attr({"tirx.dyn_smem_bytes": 65 * 8})
+ T.cuda.dyn_smem_bytes(65 * 8)
A_smem = T.decl_tensor((64,), "float16", dyn.data,
layout=T.TileLayout(T.S[64]))
mbar = T.decl_tensor((1,), "uint64", dyn.data, elem_offset=16)
if tid == 0:
@@ -1559,7 +1559,7 @@ def selector_gather(
T.cta_id([1])
tid = T.thread_id([128])
dyn = T.alloc_tensor((520,), "uint64", scope="shared.dyn")
- T.attr({"tirx.dyn_smem_bytes": 520 * 8})
+ T.cuda.dyn_smem_bytes(520 * 8)
A_smem = T.decl_tensor(
(4, 64), "bfloat16", dyn.data, layout=T.TileLayout(T.S[4, 64])
)
@@ -1730,7 +1730,7 @@ def _build_sparse_decode_qo_tma_regression():
T.cta_id([1])
tid = T.thread_id([128])
dyn = T.alloc_tensor((shared_bytes + 8,), "uint8", scope="shared.dyn")
- T.attr({"tirx.dyn_smem_bytes": shared_bytes + 8})
+ T.cuda.dyn_smem_bytes(shared_bytes + 8)
q_smem = T.decl_tensor(
(64, 512), "bfloat16", dyn.data, scope="shared.dyn",
layout=q_layout
)
@@ -2004,7 +2004,7 @@ def _build_selector_gather_gpu_kernel(dtype="float16"):
T.cta_id([1])
tid = T.thread_id([128])
dyn = T.alloc_tensor((shared_bytes + 64,), "uint8", scope="shared.dyn")
- T.attr({"tirx.dyn_smem_bytes": shared_bytes + 64})
+ T.cuda.dyn_smem_bytes(shared_bytes + 64)
A_smem = T.decl_tensor(
(4, cols), dtype, dyn.data, layout=T.TileLayout(T.S[4, cols])
)
