This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-tirx.git
The following commit(s) were added to refs/heads/main by this push:
new 00163899f [FIX][CUDA] Register void result types for several functions
(#184)
00163899f is described below
commit 00163899f872f413a0fd6678755fb4147f7b7372
Author: Linzhang Li <[email protected]>
AuthorDate: Tue Nov 3 20:57:25 2026 -0500
[FIX][CUDA] Register void result types for several functions (#184)
## Summary
This PR registers a void result type for six CUDA and NVSHMEM statement
intrinsics. Without it,
calling them from Python under an IRBuilder without the parser silently
drops the statement.
| Op | Same-file op that already registers void |
| --- | --- |
| `tirx.cuda.wait_until` | `tirx.nvshmem.wait_until` |
| `tirx.cuda.timer_init` | `tirx.cuda.cluster_sync` |
| `tirx.cuda.timer_start` | `tirx.cuda.cluster_sync` |
| `tirx.cuda.timer_end` | `tirx.cuda.cluster_sync` |
| `tirx.cuda.timer_finalize` | `tirx.cuda.cluster_sync` |
| `tirx.nvshmem.getmem_nbi_block` | `getmem_nbi`, `getmem_nbi_warp` |
### Why the statements were dropped
An op with no `TFixedReturnType` and no `FInferType` builds a Call whose
type is `Missing`.
`MaybeEmitVoidCall` (`src/script/ir_builder/ir.cc`) emits a Call as an
`Evaluate` only when its
type is void, and otherwise returns the Call to the caller. The
TVMScript parser wraps expression
statements in `Evaluate` itself, so parsed functions were unaffected.
Python code that calls the
script namespace directly under an IRBuilder (for example TIRx-lite in
tirx-kernels) discards the
returned Call, so the statement never reaches the IR or the generated
CUDA.
For `wait_until` this removes a cross-CTA wait: kernels that use it to
wait for other CTAs (a
grid barrier, a publish flag) were compiled without the wait.
---
src/backend/cuda/op/target_builtin.cc | 6 ++
tests/python/backend/cuda/test_op_result_types.py | 82 +++++++++++++++++++++++
2 files changed, 88 insertions(+)
diff --git a/src/backend/cuda/op/target_builtin.cc
b/src/backend/cuda/op/target_builtin.cc
index f5a1ba310..8c2be2eeb 100644
--- a/src/backend/cuda/op/target_builtin.cc
+++ b/src/backend/cuda/op/target_builtin.cc
@@ -137,6 +137,7 @@ void RegisterCudaTargetBuiltins() {
sig::arg<IntExpr>("num_groups", "The number of groups."),
sig::arg<IntExpr>("group_id", "The group identifier."))
.set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
.expose();
OpDef("tirx.cuda.timer_start")
@@ -147,6 +148,7 @@ void RegisterCudaTargetBuiltins() {
sig::arg<IntExpr>("profiler_write_stride", "The profiler
write stride."),
sig::arg("leader_cond", "The leader condition."))
.set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
.expose();
OpDef("tirx.cuda.timer_end")
@@ -157,6 +159,7 @@ void RegisterCudaTargetBuiltins() {
sig::arg<IntExpr>("profiler_write_stride", "The profiler
write stride."),
sig::arg("leader_cond", "The leader condition."))
.set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
.expose();
OpDef("tirx.cuda.timer_finalize")
@@ -166,6 +169,7 @@ void RegisterCudaTargetBuiltins() {
sig::arg<IntExpr>("profiler_write_stride", "The profiler
write stride."),
sig::arg("leader_cond", "The leader condition."))
.set_attr(tvm::op_attr::kCallEffectKind, CallEffectKind::kOpaque)
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
.expose();
RegisterDeviceIntrinsics();
@@ -213,6 +217,7 @@ void RegisterDeviceIntrinsics() {
RegisterDeviceIntrinsic(OpDef("tirx.cuda.wait_until"),
CallEffectKind::kOpaque, sig::arg("dst"),
sig::arg("ptr"), sig::arg("condition"),
sig::arg("scope"),
sig::arg("space"), sig::arg("ptx_type"),
sig::arg<IntExpr>("backoff_ns"))
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
.expose();
RegisterDeviceIntrinsic(OpDef("tirx.cuda.ballot_sync"),
CallEffectKind::kOpaque,
sig::arg<IntExpr>("mask"), sig::arg("pred"))
@@ -415,6 +420,7 @@ void RegisterDeviceIntrinsics() {
RegisterDeviceIntrinsic(OpDef("tirx.nvshmem.getmem_nbi_block"),
CallEffectKind::kOpaque,
sig::arg("dst"), sig::arg("src"),
sig::arg<IntExpr>("nelems"),
sig::arg<IntExpr>("pe"))
+ .set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void())
.expose();
RegisterDeviceIntrinsic(OpDef("tirx.nvshmem.getmem_nbi_warp"),
CallEffectKind::kOpaque,
sig::arg("dst"), sig::arg("src"),
sig::arg<IntExpr>("nelems"),
diff --git a/tests/python/backend/cuda/test_op_result_types.py
b/tests/python/backend/cuda/test_op_result_types.py
new file mode 100644
index 000000000..2d35aca16
--- /dev/null
+++ b/tests/python/backend/cuda/test_op_result_types.py
@@ -0,0 +1,82 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Result types of the CUDA and NVSHMEM intrinsics.
+
+A builder emits a Call as a statement only when its type is void, so a
statement
+intrinsic registered without a result type is silently dropped when Python code
+calls it under an IRBuilder without the parser.
+"""
+
+import pytest
+from tvm_ffi import structural_equal, structural_walk
+
+from tvm_tirx import ir
+from tvm_tirx.ir.op import Op
+from tvm_tirx.ir.type import PrimType
+from tvm_tirx.script import tirx as T
+from tvm_tirx.script.ir_builder import AlreadyEmitted, IRBuilder
+from tvm_tirx.tirx.script import ir_builder as I
+
+STATEMENT_INTRINSICS = [
+ "tirx.cuda.wait_until",
+ "tirx.cuda.timer_init",
+ "tirx.cuda.timer_start",
+ "tirx.cuda.timer_end",
+ "tirx.cuda.timer_finalize",
+ "tirx.nvshmem.getmem_nbi_block",
+]
+
+# These constructors require the caller to pass the result type.
+CALLER_TYPED = {"tirx.cuda.mma_fill", "tirx.cuda.mma_store"}
+
+
[email protected]("name", STATEMENT_INTRINSICS)
+def test_statement_intrinsic_registers_void(name):
+ fixed = Op.get_attr_map("TFixedReturnType")
+ op = Op.get(name)
+ assert op in fixed and structural_equal(fixed[op], PrimType("void"))
+
+
+def test_every_intrinsic_has_a_result_type():
+ fixed = Op.get_attr_map("TFixedReturnType")
+ infer = Op.get_attr_map("FInferType")
+ untyped = {
+ name
+ for name in Op.list_op_names()
+ if name.startswith(("tirx.cuda.", "tirx.nvshmem."))
+ and Op.get(name) not in fixed
+ and Op.get(name) not in infer
+ }
+ assert untyped == CALLER_TYPED
+
+
+def test_statement_intrinsics_emit_under_a_builder():
+ with IRBuilder() as builder:
+ with I.function_():
+ flag = I.arg_("flag", T.Tensor((1,), "int32"))
+ spin = T.bind_(T.alloc_local((1,), "int32"))[0]
+ ptr = flag.ptr_to([0])
+ receipts = [
+ T.cuda.wait_until(spin, ptr, spin >= 1, scope="gpu",
ptx_type="b32"),
+ T.cuda.timer_init(ptr, 0, 0, 1, 0),
+ T.nvshmem.getmem_nbi.block(ptr, ptr, 1, 0),
+ ]
+ assert all(isinstance(receipt, AlreadyEmitted) for receipt in receipts)
+ names = []
+ structural_walk(builder.get(), (ir.Call, lambda call:
names.append(call.op.name)))
+ for name in ("tirx.cuda.wait_until", "tirx.cuda.timer_init",
"tirx.nvshmem.getmem_nbi_block"):
+ assert names.count(name) == 1
