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 1c7116249 [Fix][TIRx] Restore default Z3 fallback in tile lowering
(#166)
1c7116249 is described below
commit 1c711624991329841bcc5e5bccae31b3c9b871f5
Author: Shushi Hong <[email protected]>
AuthorDate: Mon Nov 2 15:12:34 2026 -0500
[Fix][TIRx] Restore default Z3 fallback in tile lowering (#166)
This PR restores DEFAULT-strength Z3 fallback for the C++ TileDispatcher
and CUDA pattern layout analyzer, removing the temporary opt-outs
introduced in #146 while retaining Analyzer reuse. ComposeLayout now
uses existing layout bounds to classify address terms contained within
one swizzle atom before attempting atom-alignment proofs. The existing
carry checks, symbolic proof path, and full-swizzle fallback are
preserved.
---
python/tvm_tirx/backend/cuda/tile/patterns/common.py | 1 -
src/tirx/ir/layout/compose_layout.cc | 16 ++++++++++------
src/tirx/tile/lower.cc | 4 +---
tests/python/tirx/test_layout.py | 9 ++++++---
tests/python/tirx/tile/cuda/test_explicit_patterns.py | 2 +-
5 files changed, 18 insertions(+), 14 deletions(-)
diff --git a/python/tvm_tirx/backend/cuda/tile/patterns/common.py
b/python/tvm_tirx/backend/cuda/tile/patterns/common.py
index 05f46a76e..558248e7a 100644
--- a/python/tvm_tirx/backend/cuda/tile/patterns/common.py
+++ b/python/tvm_tirx/backend/cuda/tile/patterns/common.py
@@ -63,7 +63,6 @@ class Context:
# Keep it separate from the bounded verifier so address construction
# does not depend on its evolving group/repeat/descriptor premises.
self.layout_analyzer = Analyzer()
- self.layout_analyzer.set_z3_default_fallback(False)
self.constraints = ConstraintContext()
self.static_ranges = {}
for var, extent in sctx.launch_params.values():
diff --git a/src/tirx/ir/layout/compose_layout.cc
b/src/tirx/ir/layout/compose_layout.cc
index 3bc010e55..67677df5c 100644
--- a/src/tirx/ir/layout/compose_layout.cc
+++ b/src/tirx/ir/layout/compose_layout.cc
@@ -201,16 +201,20 @@ ffi::Map<ffi::String, PrimExpr> ApplyStructured(const
ComposeLayoutObj* layout,
if (!value.has_value() || !add_constant(*value)) return fallback();
continue;
}
+ PrimExpr simplified_extent = analyzer->Simplify(iter->extent);
+ const auto* extent_imm = simplified_extent.as<IntImmObj>();
+ auto extent = extent_imm ? extent_imm->value.as<int64_t>() : std::nullopt;
+ int64_t term_max = 0;
+ bool has_bound =
+ extent.has_value() && *extent > 0 && MulWithoutOverflow(*extent - 1,
*stride, &term_max);
+ // Classify terms contained in one atom before attempting symbolic
division.
+ // Their layout bounds already suffice for the existing carry check.
+ if (has_bound && term_max < atom && add_low(term, term_max)) continue;
if (auto quotient = DivideExactTerm(term, atom, analyzer);
quotient.has_value()) {
add_high(term, quotient.value());
continue;
}
- PrimExpr simplified_extent = analyzer->Simplify(iter->extent);
- const auto* extent_imm = simplified_extent.as<IntImmObj>();
- auto extent = extent_imm ? extent_imm->value.as<int64_t>() : std::nullopt;
- if (!extent.has_value() || *extent <= 0) return fallback();
- int64_t term_max;
- if (!MulWithoutOverflow(*extent - 1, *stride, &term_max)) return
fallback();
+ if (!has_bound) return fallback();
// A single structured coordinate may span multiple atoms even when its
// stride is smaller than one atom. If the stride divides the atom, split
diff --git a/src/tirx/tile/lower.cc b/src/tirx/tile/lower.cc
index 951bc6cf9..fa0b4c00a 100644
--- a/src/tirx/tile/lower.cc
+++ b/src/tirx/tile/lower.cc
@@ -174,9 +174,7 @@ class TileDispatcher : public StmtExprMutator {
public:
using StmtExprMutator::Mutate;
using StmtExprMutator::Mutate_;
- explicit TileDispatcher(const Target& target) : target_(target) {
- analyzer_->SetZ3DefaultFallback(false);
- }
+ explicit TileDispatcher(const Target& target) : target_(target) {}
static Stmt LowerOpCalls(const Stmt& stmt, const Target& target) {
return ffi::make_object<TileDispatcher>(target)
diff --git a/tests/python/tirx/test_layout.py b/tests/python/tirx/test_layout.py
index 66b085ed5..b644149f0 100644
--- a/tests/python/tirx/test_layout.py
+++ b/tests/python/tirx/test_layout.py
@@ -1302,14 +1302,17 @@ def _naive_compose_expr(tile, coords, per_element=3,
swizzle_len=3, atom_len=3,
@pytest.mark.parametrize("swizzle_len", [3, 2, 1, 0])
-def test_compose_structured_apply_uses_bounded_normal_form(swizzle_len):
[email protected]("z3_default_fallback", [False, True])
+def test_compose_structured_apply_uses_bounded_normal_form(swizzle_len,
z3_default_fallback):
x = Var("x", ty="int32")
y = Var("y", ty="int32")
tile = TileLayout(S[(8, 8) : (64, 1)] + 72)
layout = ComposeLayout(3, swizzle_len, 3, tile)
+ analyzer = tvm_tirx.sym.Analyzer()
+ analyzer.set_z3_default_fallback(z3_default_fallback)
- actual = layout.apply(x, y)["m"]
- shaped = layout.apply(x, y, shape=(8, 8))["m"]
+ actual = layout.apply(x, y, analyzer=analyzer)["m"]
+ shaped = layout.apply(x, y, shape=(8, 8), analyzer=analyzer)["m"]
if swizzle_len == 0:
expected = x * 64 + y + 72
else:
diff --git a/tests/python/tirx/tile/cuda/test_explicit_patterns.py
b/tests/python/tirx/tile/cuda/test_explicit_patterns.py
index 0e6b49e88..78e79b810 100644
--- a/tests/python/tirx/tile/cuda/test_explicit_patterns.py
+++ b/tests/python/tirx/tile/cuda/test_explicit_patterns.py
@@ -162,7 +162,7 @@ def test_layout_analysis_is_separate_from_verification():
a, b = Var("a", "int32"), Var("b", "int32")
predicate = a * a + b * b >= 2 * a * b
assert ctx.analyzer.can_prove(predicate) == ctx.analyzer.is_z3_enabled
- assert not ctx.layout_analyzer.can_prove(predicate)
+ assert ctx.layout_analyzer.can_prove(predicate) ==
ctx.layout_analyzer.is_z3_enabled
def test_contract_rejections():
