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 40c2f54908 [TIR] Ignore None-valued pragma annotations (#20265)
40c2f54908 is described below
commit 40c2f54908e96875fa19e92db83f123854d97991
Author: T1Gang <[email protected]>
AuthorDate: Wed Sep 9 11:46:27 2026 +0800
[TIR] Ignore None-valued pragma annotations (#20265)
Fixes #18393.
None-valued pragma annotations represent an unspecified pragma and
should not be lowered into `AttrStmt` nodes with undefined `PrimExpr`
values.
Previously, `pragma_unroll_explicit=None` was lowered into an `AttrStmt`
with an undefined value. During `FlattenBuffer`, `StmtExprMutator`
attempted to visit that value and triggered a segmentation fault.
This patch skips `None`-valued pragma annotations during opaque lowering
while preserving the existing behavior for valid pragma values.
For `pragma_unroll`, the existing upstream behavior is preserved:
non-null values remain attached to the loop, while null values are
ignored.
Tests:
- `tests/python/tirx-transform/test_tir_transform_flatten_buffer.py`: 11
passed
-
`tests/python/s_tir/transform/test_s_tir_transform_lower_opaque_block.py`:
13 passed
- `git diff --check upstream/main..HEAD`: passed
---
src/s_tir/transform/lower_opaque_block.cc | 8 +++---
src/tirx/transform/lower_tirx_opaque.cc | 14 ++++++----
tests/python/relax/test_frontend_onnx.py | 7 ++---
.../test_s_tir_transform_lower_opaque_block.py | 20 ++++++++++++++
.../test_tir_transform_flatten_buffer.py | 20 ++++++++++++++
.../tirx/transform/test_transform_lower_tirx.py | 31 ++++++++++++++++++++++
6 files changed, 87 insertions(+), 13 deletions(-)
diff --git a/src/s_tir/transform/lower_opaque_block.cc
b/src/s_tir/transform/lower_opaque_block.cc
index 586eb727b7..d3dcf2c3a9 100644
--- a/src/s_tir/transform/lower_opaque_block.cc
+++ b/src/s_tir/transform/lower_opaque_block.cc
@@ -158,9 +158,7 @@ class OpaqueBlockLower : public StmtExprMutator {
/*! \brief Convert attr value from annotation map into PrimExpr. */
PrimExpr ConvertAttrValue(const ffi::String& key, const Any& obj) {
- if (obj == nullptr) {
- return PrimExpr();
- } else if (auto expr = obj.try_cast<PrimExpr>()) {
+ if (auto expr = obj.try_cast<PrimExpr>()) {
return expr.value();
} else if (auto str = obj.try_cast<ffi::String>()) {
return std::move(StringImm(str.value()));
@@ -187,6 +185,10 @@ class OpaqueBlockLower : public StmtExprMutator {
for (const auto& kv : annotations) {
const ffi::String& key = kv.first;
if (tirx::attr::IsPragmaKey(key)) {
+ if (kv.second == nullptr) {
+ continue;
+ }
+
pragma_attrs->emplace_back(key, ConvertAttrValue(key, kv.second));
} else if (!is_block) {
// the loop annotation is preserved
diff --git a/src/tirx/transform/lower_tirx_opaque.cc
b/src/tirx/transform/lower_tirx_opaque.cc
index 7c6ba2b2c9..99fcd7793e 100644
--- a/src/tirx/transform/lower_tirx_opaque.cc
+++ b/src/tirx/transform/lower_tirx_opaque.cc
@@ -120,9 +120,7 @@ class TIRxOpaqueLower : public StmtExprMutator {
/*! \brief Convert attr value from annotation map into PrimExpr. */
PrimExpr ConvertAttrValue(const ffi::String& key, const Any& obj) {
- if (obj == nullptr) {
- return PrimExpr();
- } else if (auto expr = obj.try_cast<PrimExpr>()) {
+ if (auto expr = obj.try_cast<PrimExpr>()) {
return expr.value();
} else if (auto str = obj.try_cast<ffi::String>()) {
return std::move(prim::StringImm(str.value()));
@@ -149,11 +147,17 @@ class TIRxOpaqueLower : public StmtExprMutator {
for (const auto& kv : annotations) {
const ffi::String& key = kv.first;
if (key == "pragma_unroll") {
- preserved_annotations.Set(key, kv.second);
+ if (kv.second != nullptr) {
+ preserved_annotations.Set(key, kv.second);
+ }
} else if (tirx::attr::IsPragmaKey(key)) {
+ if (kv.second == nullptr) {
+ continue;
+ }
+
pragma_attrs->emplace_back(key, ConvertAttrValue(key, kv.second));
} else {
- // loop annotations are always preserved (no SBlock annotation
dropping here)
+ // Loop annotations are always preserved
preserved_annotations.Set(key, kv.second);
}
}
diff --git a/tests/python/relax/test_frontend_onnx.py
b/tests/python/relax/test_frontend_onnx.py
index 335694ed76..e2de5e02a3 100644
--- a/tests/python/relax/test_frontend_onnx.py
+++ b/tests/python/relax/test_frontend_onnx.py
@@ -1230,7 +1230,7 @@ def test_multi_input_unknown_static_shape(op_name,
num_inputs):
"Slice", ["x", "starts", "ends", "axes"], ["sliced"], name="slice0"
)
other_names = [f"y{i}" for i in range(num_inputs - 1)]
- op_node = helper.make_node(op_name, ["sliced"] + other_names, ["output"],
name="op0")
+ op_node = helper.make_node(op_name, ["sliced", *other_names], ["output"],
name="op0")
graph = helper.make_graph(
[slice_node, op_node],
@@ -1241,10 +1241,7 @@ def test_multi_input_unknown_static_shape(op_name,
num_inputs):
helper.make_tensor_value_info("ends", TensorProto.INT64, [1]),
helper.make_tensor_value_info("axes", TensorProto.INT64, [1]),
]
- + [
- helper.make_tensor_value_info(name, TensorProto.FLOAT, [2, 3])
- for name in other_names
- ],
+ + [helper.make_tensor_value_info(name, TensorProto.FLOAT, [2, 3]) for
name in other_names],
outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT,
[2, 3])],
)
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("",
13)])
diff --git
a/tests/python/s_tir/transform/test_s_tir_transform_lower_opaque_block.py
b/tests/python/s_tir/transform/test_s_tir_transform_lower_opaque_block.py
index 441074128e..273fa0bdc5 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_lower_opaque_block.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_lower_opaque_block.py
@@ -421,6 +421,26 @@ def test_preserved_annotations():
tvm.ir.assert_structural_equal(mod["main"],
after.with_attr("global_symbol", "main"))
+def test_none_pragma_annotation():
+ @T.prim_func(s_tir=True)
+ def before(A: T.Buffer(8, "float32"), B: T.Buffer(8, "float32")):
+ for i in T.serial(8, annotations={"pragma_unroll_explicit": None}):
+ with T.sblock("block"):
+ B[i] = A[i] + 1.0
+
+ @T.prim_func(s_tir=True)
+ def after(A: T.Buffer(8, "float32"), B: T.Buffer(8, "float32")):
+ for i in T.serial(8):
+ B[i] = A[i] + 1.0
+
+ mod = tvm.IRModule.from_expr(before.with_attr("global_symbol", "main"))
+ mod = tvm.s_tir.transform.LowerOpaqueBlock()(mod)
+ tvm.ir.assert_structural_equal(
+ mod["main"],
+ after.with_attr("global_symbol", "main"),
+ )
+
+
def test_boolean_handling():
_check(boolean_handling_before, boolean_handling_after)
diff --git a/tests/python/tirx-transform/test_tir_transform_flatten_buffer.py
b/tests/python/tirx-transform/test_tir_transform_flatten_buffer.py
index d01400d3a5..7c2943879c 100644
--- a/tests/python/tirx-transform/test_tir_transform_flatten_buffer.py
+++ b/tests/python/tirx-transform/test_tir_transform_flatten_buffer.py
@@ -368,5 +368,25 @@ def test_flatten_inside_block():
tvm.ir.assert_structural_equal(After, Expected)
+def test_build_with_optional_pragma_unroll_explicit():
+ def check(value):
+ @I.ir_module(s_tir=True)
+ class Module:
+ @T.prim_func(s_tir=True)
+ def main(A: T.Buffer((4, 5, 6), "int16"), B: T.Buffer((4, 5, 6),
"int16")):
+ for ax0 in T.serial(4, annotations={"pragma_unroll_explicit":
value}):
+ for ax1, ax2 in T.grid(5, 6):
+ with T.sblock("copy"):
+ v0, v1, v2 = T.axis.remap("SSS", [ax0, ax1, ax2])
+ T.reads(A[v0, v1, v2])
+ T.writes(B[v0, v1, v2])
+ B[v0, v1, v2] = A[v0, v1, v2]
+
+ tvm.build(Module, target="llvm")
+
+ for value in [None, True, False]:
+ check(value)
+
+
if __name__ == "__main__":
tvm.testing.main()
diff --git a/tests/python/tirx/transform/test_transform_lower_tirx.py
b/tests/python/tirx/transform/test_transform_lower_tirx.py
index 72909c77bb..d9a571942d 100644
--- a/tests/python/tirx/transform/test_transform_lower_tirx.py
+++ b/tests/python/tirx/transform/test_transform_lower_tirx.py
@@ -62,6 +62,37 @@ def _launch_thread_extents(func):
L_LANE = T.TileLayout(T.S[32 : 1 @ laneid])
+def test_lower_tirx_opaque_optional_pragma_annotations():
+ @T.prim_func(private=True)
+ def before(A: T.Buffer(8, "float32"), B: T.Buffer(8, "float32")):
+ for i in T.serial(8, annotations={"pragma_unroll": None}):
+ B[i] = A[i] + 1.0
+ for i in T.serial(8, annotations={"pragma_unroll_explicit": None}):
+ B[i] = A[i] + 2.0
+ for i in T.serial(8, annotations={"pragma_unroll": False}):
+ B[i] = A[i] + 3.0
+ for i in T.serial(8, annotations={"pragma_unroll_explicit": 0}):
+ B[i] = A[i] + 4.0
+
+ @T.prim_func(private=True)
+ def after(A: T.Buffer(8, "float32"), B: T.Buffer(8, "float32")):
+ for i in T.serial(8):
+ B[i] = A[i] + 1.0
+ for i in T.serial(8):
+ B[i] = A[i] + 2.0
+ for i in T.serial(8, annotations={"pragma_unroll": False}):
+ B[i] = A[i] + 3.0
+ for i in T.serial(8):
+ B[i] = A[i] + 4.0
+
+ # The pragma refers to the variable bound by the loop inside its body.
+ loop = after.body.seq[3]
+ pragma = tvm.tirx.AttrStmt(loop.loop_var, "pragma_unroll_explicit", 0,
loop)
+ after = after.with_body(tvm.tirx.SeqStmt([*after.body.seq[:3], pragma]))
+ lowered = tvm.tirx.transform.LowerTIRxOpaque()(tvm.IRModule({"main":
before}))
+ tvm.ir.assert_structural_equal(lowered["main"], after, map_free_vars=True)
+
+
def test_lower_view_get():
@T.prim_func(private=True)
def before1(in_buf: T.Buffer(64, "float32"), out: T.Buffer(64, "float32"))
-> None:
