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 e45a49da0 [Fix][TIRx] Make layout analysis state explicit and reusable
(#146)
e45a49da0 is described below
commit e45a49da0bd7141f4c7ba6c6a827ab56f641b022
Author: Shushi Hong <[email protected]>
AuthorDate: Sun Nov 1 15:30:15 2026 -0500
[Fix][TIRx] Make layout analysis state explicit and reusable (#146)
`Layout.apply` creates temporary Analyzers while evaluating coordinate
mappings, including additional instances inside nested helpers. When
queries fall back to Z3, these instances may each initialize their own
solver and context. Frequent layout evaluation therefore pays repeated
setup costs.
This change requires callers to provide an Analyzer at the C++, FFI, and
Rust layout-application entrypoints. C++ helpers borrow `AnalyzerObj*`
and reuse it throughout the evaluation. CUDA tile lowering shares a
dedicated layout analyzer across each instruction’s address queries,
separate from the bounded verifier’s assumptions.
`AnalyzerObj::SetZ3DefaultFallback(bool)` controls DEFAULT-strength Z3
fallback through a private per-instance setting, disabled by default and
preserved by cloning. `StmtExprMutatorWithConstraintContext` explicitly
enables it for its owned analyzer, including constructors used by
subclasses. Native analysis still runs first; explicit stronger queries
and Z3-specific APIs retain their existing behavior. Z3Prover is
unchanged.
Z3 fallback enabled, this can produce many inconclusive SMT queries.
Optimizing that proof strategy will be handled separately from this
change to analysis-state ownership and reuse.
---------
Co-authored-by: LeiWang1999 <[email protected]>
---
include/tvm/sym/analyzer.h | 12 +++++++
include/tvm/tirx/layout.h | 19 +++++++---
.../tvm_tirx/backend/cuda/tile/patterns/common.py | 13 ++++++-
.../backend/cuda/tile/patterns/tcgen05_mma.py | 2 +-
python/tvm_tirx/backend/cuda/tile/patterns/tma.py | 12 +++++--
python/tvm_tirx/sym/analyzer.py | 25 ++++++++++---
python/tvm_tirx/tirx/layout.py | 23 +++++++++---
rust/src/sym/analyzer.rs | 10 ++++++
rust/src/tirx/layout.rs | 16 ++++++---
rust/tests/tirx/layout.rs | 10 +++---
src/backend/trn/transform/lower_trainium_layout.cc | 4 +--
src/sym/analyzer.cc | 11 ++++--
src/tirx/ir/layout/compose_layout.cc | 42 +++++++++++-----------
src/tirx/ir/layout/layout.cc | 38 ++++++++++++--------
src/tirx/ir/layout/tile_core.cc | 19 +++++-----
src/tirx/ir/layout/tile_internal.h | 20 +++++++----
src/tirx/ir/layout/tile_tile_ops.cc | 3 +-
src/tirx/ir/tensor.cc | 3 +-
src/tirx/tile/lower.cc | 4 ++-
src/tirx/transform/lower_tensor_layouts.cc | 2 +-
tests/python/sym/test_z3.py | 15 ++++++--
tests/python/tirx/test_layout.py | 16 +++++++++
.../tirx/tile/cuda/test_explicit_patterns.py | 15 ++++++++
23 files changed, 246 insertions(+), 88 deletions(-)
diff --git a/include/tvm/sym/analyzer.h b/include/tvm/sym/analyzer.h
index f59a7acbc..e7e1eee9c 100644
--- a/include/tvm/sym/analyzer.h
+++ b/include/tvm/sym/analyzer.h
@@ -569,6 +569,15 @@ class TVM_TIRX_DLL AnalyzerObj : public ffi::Object {
Z3Prover z3_prover;
/*! \brief Construct an analyzer with no accepted premises and default query
settings. */
AnalyzerObj();
+ /*!
+ * \brief Set whether DEFAULT-strength queries try Z3 after native analysis
fails.
+ *
+ * Enabled by default. This setting is local to the analyzer and preserved
by Clone.
+ * Stronger proof queries and direct Z3 APIs are unaffected. Disabling the
fallback
+ * does not clear accepted premises or existing solver state.
+ * \param enabled Whether to allow the DEFAULT-strength Z3 fallback.
+ */
+ void SetZ3DefaultFallback(bool enabled);
/*!
* \brief Synchronize all derived facts to the supplied premise environment.
*
@@ -659,6 +668,7 @@ class TVM_TIRX_DLL AnalyzerObj : public ffi::Object {
* Do not use strength beyond default in sub-analyzers and
* only use it in top-level predicate analysis.
* When built with Z3, all strengths try it after native analysis fails.
+ * SetZ3DefaultFallback(false) disables this fallback only for
DEFAULT-strength queries.
* Z3's existing resource limit and timeout apply; unknown means cannot
prove.
*/
bool CanProve(const PrimExpr& cond, ProofStrength strength =
ProofStrength::kDefault);
@@ -704,6 +714,8 @@ class TVM_TIRX_DLL AnalyzerObj : public ffi::Object {
private:
friend class detail::TempConstraintScope;
+ // Policy for DEFAULT-strength queries; explicit stronger queries can still
use Z3.
+ bool z3_default_fallback_{true};
// Record a value binding after AnalyzerObj validates and orders the
premises.
void IngestBinding(const Var& var, const PrimExpr& expr);
// Record a half-open variable range from the accepted premise prefix.
diff --git a/include/tvm/tirx/layout.h b/include/tvm/tirx/layout.h
index 40287a776..87c7336bf 100644
--- a/include/tvm/tirx/layout.h
+++ b/include/tvm/tirx/layout.h
@@ -66,11 +66,22 @@ class Layout : public ffi::ObjectRef {
/*! \brief Get the span of the layout (of some axis) */
TVM_TIRX_DLL PrimExpr GetSpan(ffi::Optional<ffi::String> axis_name =
std::nullopt) const;
- /*! \brief Apply layout on the input coordinate and get the mapped output */
- TVM_TIRX_DLL ffi::Map<ffi::String, PrimExpr> Apply(ffi::Array<PrimExpr>
coord) const;
- TVM_TIRX_DLL ffi::Map<ffi::String, PrimExpr> Apply(PrimExpr coord) const;
+ /*!
+ * \brief Apply a layout using caller-owned analysis state.
+ *
+ * The analyzer is required and must outlive the call. Its current premises
+ * may simplify the result. The layout adds no persistent premises and does
+ * not retain the analyzer.
+ * Related queries can reuse its caches and solver. Concurrent queries must
+ * use separate analyzers.
+ */
+ TVM_TIRX_DLL ffi::Map<ffi::String, PrimExpr> Apply(ffi::Array<PrimExpr>
coord,
+ sym::AnalyzerObj*
analyzer) const;
+ TVM_TIRX_DLL ffi::Map<ffi::String, PrimExpr> Apply(PrimExpr coord,
+ sym::AnalyzerObj*
analyzer) const;
TVM_TIRX_DLL ffi::Map<ffi::String, PrimExpr> Apply(const
ffi::Array<PrimExpr>& coord,
- const
ffi::Array<PrimExpr>& shape) const;
+ const
ffi::Array<PrimExpr>& shape,
+ sym::AnalyzerObj*
analyzer) const;
/*! \brief Turn the layout to canonical form */
TVM_TIRX_DLL Layout Canonicalize() const;
diff --git a/python/tvm_tirx/backend/cuda/tile/patterns/common.py
b/python/tvm_tirx/backend/cuda/tile/patterns/common.py
index b4a26242b..79de578a8 100644
--- a/python/tvm_tirx/backend/cuda/tile/patterns/common.py
+++ b/python/tvm_tirx/backend/cuda/tile/patterns/common.py
@@ -59,6 +59,13 @@ class Context:
def __init__(self, sctx, atom_scope="thread"):
self.sctx = sctx
self.analyzer = Analyzer()
+ # Repeated descriptor checks use native proofs before bounded
verification.
+ self.analyzer.set_z3_default_fallback(False)
+ # Reuse layout-query state across this instruction's address walk.
+ # 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():
@@ -281,7 +288,11 @@ class Region:
)
def physical(self, atom, repeat=None):
- result = dict(self.layout.apply(*self.indices(atom, repeat),
shape=self.ty.shape))
+ result = dict(
+ self.layout.apply(
+ *self.indices(atom, repeat), shape=self.ty.shape,
analyzer=self.ctx.layout_analyzer
+ )
+ )
return {key: self.ctx.analyzer.simplify(value) for key, value in
result.items()}
def memory(self, atom):
diff --git a/python/tvm_tirx/backend/cuda/tile/patterns/tcgen05_mma.py
b/python/tvm_tirx/backend/cuda/tile/patterns/tcgen05_mma.py
index b9a9d4e72..a1517d634 100644
--- a/python/tvm_tirx/backend/cuda/tile/patterns/tcgen05_mma.py
+++ b/python/tvm_tirx/backend/cuda/tile/patterns/tcgen05_mma.py
@@ -59,7 +59,7 @@ def _tmem_operand(ctx, operand, expected, atom_shape):
for var, size in zip(coords, atom_shape):
ctx.add_range(var, Range(0, size))
physical = operand.physical(coords)
- image = expected.apply(*coords, shape=atom_shape)
+ image = expected.apply(*coords, shape=atom_shape,
analyzer=ctx.layout_analyzer)
if set(physical) - {"TLane", "TCol"}:
raise ValueError(f"{operand.name}: expected a TMEM image")
for axis in ("TLane", "TCol"):
diff --git a/python/tvm_tirx/backend/cuda/tile/patterns/tma.py
b/python/tvm_tirx/backend/cuda/tile/patterns/tma.py
index 0f0adbc5f..7b2b9f49e 100644
--- a/python/tvm_tirx/backend/cuda/tile/patterns/tma.py
+++ b/python/tvm_tirx/backend/cuda/tile/patterns/tma.py
@@ -128,7 +128,7 @@ def _tensor_map(ctx, region, rank, box, encoding):
strides = [prod(ty.shape[i + 1 :]) for i in range(tensor_rank)]
base = {"m": sum(i * s for i, s in zip(starts, strides))}
else:
- base = layout.apply(*starts, shape=ty.shape)
+ base = layout.apply(*starts, shape=ty.shape,
analyzer=ctx.layout_analyzer)
if set(base) - {"m"}:
raise ValueError("TMA global view must use only the memory axis")
if layout is not None:
@@ -138,7 +138,9 @@ def _tensor_map(ctx, region, rank, box, encoding):
unit = list(starts)
unit[i] += 1
if layout is not None:
- stride = ctx.analyzer.simplify(layout.apply(*unit,
shape=ty.shape)["m"] - base["m"])
+ stride = ctx.analyzer.simplify(
+ layout.apply(*unit, shape=ty.shape,
analyzer=ctx.layout_analyzer)["m"] - base["m"]
+ )
# A unit dimension has no address movement; choose the smallest
# legal descriptor stride for this unobserved coordinate.
if ctx.analyzer.can_prove(size == 1):
@@ -148,7 +150,11 @@ def _tensor_map(ctx, region, rank, box, encoding):
ctx.add_range(var, Range(0, size))
coordinates.append(var)
if layout is not None:
- actual = layout.apply(*(a + b for a, b in zip(starts, coordinates)),
shape=ty.shape)["m"]
+ actual = layout.apply(
+ *(a + b for a, b in zip(starts, coordinates)),
+ shape=ty.shape,
+ analyzer=ctx.layout_analyzer,
+ )["m"]
ctx.require(
actual == base["m"] + sum(c * s for c, s in zip(coordinates,
strides)),
"TMA global view must be rectangular with fixed strides; "
diff --git a/python/tvm_tirx/sym/analyzer.py b/python/tvm_tirx/sym/analyzer.py
index 1123e7a94..5a5a1b5db 100644
--- a/python/tvm_tirx/sym/analyzer.py
+++ b/python/tvm_tirx/sym/analyzer.py
@@ -92,6 +92,21 @@ class Analyzer(tvm_ffi.Object):
def __init__(self):
self.__init_handle_by_constructor__(_ffi_api.Analyzer)
+ def set_z3_default_fallback(self, enabled: bool) -> None:
+ """Set whether DEFAULT-strength queries try Z3 after native analysis
fails.
+
+ Enabled by default. This setting is local to this analyzer and
preserved
+ by clone(). Stronger proof queries and Z3-specific APIs are unaffected.
+ Changing it does not clear accepted premises or existing solver state.
+ Z3 must also be available in the native build.
+
+ Parameters
+ ----------
+ enabled : bool
+ Whether to allow the DEFAULT-strength Z3 fallback.
+ """
+ _ffi_api.AnalyzerSetZ3DefaultFallback(self, enabled)
+
@property
def is_z3_enabled(self) -> bool:
"""Whether this build includes the Z3 backend (``USE_Z3=ON``).
@@ -324,10 +339,12 @@ class Analyzer(tvm_ffi.Object):
strength: ProofStrength
The proof strength. ``SYMBOLIC_BOUND`` additionally enables
symbolic
- bound analysis. When built with Z3 (``TVM_TIRX_USE_Z3=ON``), all
- strengths, including ``DEFAULT``, try Z3 after native analysis
fails
- to prove the predicate. The existing Z3 timeout and resource limits
- apply; an unknown solver result means the predicate cannot be
proved.
+ bound analysis. When built with Z3 (``TVM_TIRX_USE_Z3=ON``),
explicit
+ stronger queries try Z3 after native analysis fails to prove the
+ predicate. ``DEFAULT`` queries also try Z3 by default; this can be
+ disabled with :py:meth:`set_z3_default_fallback`.
+ The existing Z3 timeout and resource limits apply; an unknown
solver
+ result means the predicate cannot be proved.
Returns
-------
diff --git a/python/tvm_tirx/tirx/layout.py b/python/tvm_tirx/tirx/layout.py
index 4a297f254..6f71de117 100644
--- a/python/tvm_tirx/tirx/layout.py
+++ b/python/tvm_tirx/tirx/layout.py
@@ -71,7 +71,12 @@ class Layout(tvm_ffi.Object):
# Note: no backward-compat alias; `cosize` is removed.
- def apply(self, *coord: list[Expr], shape: list[Expr] | None = None) ->
dict[str, Expr]:
+ def apply(
+ self,
+ *coord: list[Expr],
+ shape: list[Expr] | None = None,
+ analyzer: "tvm_tirx.sym.Analyzer | None" = None,
+ ) -> dict[str, Expr]:
"""Apply the layout on the input coordinate and get the mapped output.
Input cases:
@@ -80,16 +85,26 @@ class Layout(tvm_ffi.Object):
- shape is provided -> turn the coord with shape into a 1D coordinate
- shape is not provided -> use the default shape
+ analyzer optionally supplies caller-owned analysis state for repeated
+ queries. Its premises may simplify the result; this method adds no
+ persistent premises and does not retain the analyzer. Concurrent calls
+ must use separate analyzers. Without it, a temporary analyzer with the
+ DEFAULT-strength Z3 fallback disabled is created. Pass an analyzer
+ explicitly to control its proof settings.
+
Returns
-------
Dict[str, Expr]
The mapped output (axis name -> value on the axis)
"""
+ if analyzer is None:
+ analyzer = tvm_tirx.sym.Analyzer()
+ analyzer.set_z3_default_fallback(False)
if shape is not None:
- return _ffi_api.LayoutApplyWithShape(self, coord, shape) #
pylint: disable=no-member
+ return _ffi_api.LayoutApplyWithShape(self, coord, shape, analyzer)
# pylint: disable=no-member
if len(coord) == 1:
- return _ffi_api.LayoutApplyLinear(self, coord[0]) # pylint:
disable=no-member
- return _ffi_api.LayoutApply(self, coord) # pylint: disable=no-member
+ return _ffi_api.LayoutApplyLinear(self, coord[0], analyzer) #
pylint: disable=no-member
+ return _ffi_api.LayoutApply(self, coord, analyzer) # pylint:
disable=no-member
def canonicalize(self) -> "Layout":
"""Canonicalize the layout by simplifying and fusing iterators where
possible.
diff --git a/rust/src/sym/analyzer.rs b/rust/src/sym/analyzer.rs
index edd73d73d..8c07af9d3 100644
--- a/rust/src/sym/analyzer.rs
+++ b/rust/src/sym/analyzer.rs
@@ -423,6 +423,16 @@ impl Analyzer {
)
}
+ /// Set whether default-strength queries try Z3 after native analysis
fails.
+ ///
+ /// Enabled by default and preserved by `deep_clone`. Stronger queries and
direct Z3
+ /// APIs are unaffected. This does not clear accepted premises or solver
state.
+ pub fn set_z3_default_fallback(&self, enabled: bool) -> Result<()> {
+ tvm_ffi::cached_global_func!("sym.AnalyzerSetZ3DefaultFallback")
+ .call_tuple((self, enabled))?;
+ Ok(())
+ }
+
/// A copy of the analyzer and the facts it knows (`sym.AnalyzerClone`).
pub fn deep_clone(&self) -> Result<Self> {
tvm_ffi::cached_global_func!("sym.AnalyzerClone").call_tuple((self,))?.try_into()
diff --git a/rust/src/tirx/layout.rs b/rust/src/tirx/layout.rs
index f31b293fb..0c88c7ca0 100644
--- a/rust/src/tirx/layout.rs
+++ b/rust/src/tirx/layout.rs
@@ -217,16 +217,23 @@ impl Layout {
pub fn apply(
&self,
coord: impl IntoIterator<Item = impl IntoExpr>,
+ analyzer: &Analyzer,
) -> Result<Map<String, Expr>> {
let coord =
coord.into_iter().map(IntoExpr::into_expr).collect::<Result<Array<Expr>>>()?;
- tvm_ffi::cached_global_func!("tirx.LayoutApply").call_tuple((self,
&coord))?.try_into()
+ tvm_ffi::cached_global_func!("tirx.LayoutApply")
+ .call_tuple((self, &coord, analyzer))?
+ .try_into()
}
/// The coordinate on each axis of the element at the linear index `coord`
/// (`tirx.LayoutApplyLinear`).
- pub fn apply_linear(&self, coord: impl IntoExpr) -> Result<Map<String,
Expr>> {
+ pub fn apply_linear(
+ &self,
+ coord: impl IntoExpr,
+ analyzer: &Analyzer,
+ ) -> Result<Map<String, Expr>> {
tvm_ffi::cached_global_func!("tirx.LayoutApplyLinear")
- .call_tuple((self, &coord.into_expr()?))?
+ .call_tuple((self, &coord.into_expr()?, analyzer))?
.try_into()
}
@@ -236,11 +243,12 @@ impl Layout {
&self,
coord: impl IntoIterator<Item = impl IntoExpr>,
shape: impl IntoIterator<Item = impl IntoExpr>,
+ analyzer: &Analyzer,
) -> Result<Map<String, Expr>> {
let coord =
coord.into_iter().map(IntoExpr::into_expr).collect::<Result<Array<Expr>>>()?;
let shape =
shape.into_iter().map(IntoExpr::into_expr).collect::<Result<Array<Expr>>>()?;
tvm_ffi::cached_global_func!("tirx.LayoutApplyWithShape")
- .call_tuple((self, &coord, &shape))?
+ .call_tuple((self, &coord, &shape, analyzer))?
.try_into()
}
diff --git a/rust/tests/tirx/layout.rs b/rust/tests/tirx/layout.rs
index 1a5c4b724..bec4199ae 100644
--- a/rust/tests/tirx/layout.rs
+++ b/rust/tests/tirx/layout.rs
@@ -20,6 +20,7 @@
//! Tests of `tirx::layout`: layout axes, tile layouts and swizzled layouts.
use tvm_tirx::ir::Expr;
+use tvm_tirx::sym::Analyzer;
use tvm_tirx::tirx::{Axis, ComposeLayout, Iter, Layout, TileLayout};
use tvm_tirx::tvm_ffi::{Map, String as FfiString};
@@ -38,6 +39,7 @@ fn memory_axes() {
#[test]
fn layouts() {
let _serial = serial();
+ let analyzer = Analyzer::new();
let tile = tile_8x8();
let layout = Layout::from(&tile);
layout.check_well_formed().unwrap();
@@ -45,9 +47,9 @@ fn layouts() {
assert_eq!(as_i64(&layout.span(Some("m")).unwrap()), 64);
let key = FfiString::from("m");
let at = |coord: Map<FfiString, Expr>|
as_i64(&coord.get(&key).unwrap().unwrap());
- assert_eq!(at(layout.apply([int32(1), int32(2)]).unwrap()), 10);
- assert_eq!(at(layout.apply_linear(int32(10)).unwrap()), 10);
- assert_eq!(at(layout.apply_with_shape([int32(10)], [int32(64)]).unwrap()),
10);
+ assert_eq!(at(layout.apply([int32(1), int32(2)], &analyzer).unwrap()), 10);
+ assert_eq!(at(layout.apply_linear(int32(10), &analyzer).unwrap()), 10);
+ assert_eq!(at(layout.apply_with_shape([int32(10)], [int32(64)],
&analyzer).unwrap()), 10);
let canonical = layout.canonicalize().unwrap();
assert_eq!(as_i64(&canonical.size(None).unwrap()), 64);
@@ -64,7 +66,7 @@ fn layouts() {
// Memory offset 10 becomes 10 ^ ((10 >> 3) & 3) = 11.
let swizzled = ComposeLayout::new(0, 2, 3, &tile, true).unwrap();
assert_eq!(swizzled.atom_len().unwrap(), 3);
- assert_eq!(at(Layout::from(&swizzled).apply([int32(1),
int32(2)]).unwrap()), 11);
+ assert_eq!(at(Layout::from(&swizzled).apply([int32(1), int32(2)],
&analyzer).unwrap()), 11);
}
#[test]
diff --git a/src/backend/trn/transform/lower_trainium_layout.cc
b/src/backend/trn/transform/lower_trainium_layout.cc
index 44db1f01d..1ac75d860 100644
--- a/src/backend/trn/transform/lower_trainium_layout.cc
+++ b/src/backend/trn/transform/lower_trainium_layout.cc
@@ -297,7 +297,7 @@ class TrainiumLayoutApplier : public
tirx::StmtExprMutatorWithConstraintContext
.as_or_throw<Layout>();
auto tile_layout = layout.as<TileLayoutObj>();
if (IsTrainiumLayout(tile_layout)) {
- auto coord = layout.Apply(indices, tensor.shape());
+ auto coord = layout.Apply(indices, tensor.shape(),
GetSyncedAnalyzer());
std::vector<PrimExpr> res;
for (const auto& axis : tensor.scope() == scope_trn_psum
? ffi::Array<ffi::String>{"Bank", "P", "F"}
@@ -316,7 +316,7 @@ class TrainiumLayoutApplier : public
tirx::StmtExprMutatorWithConstraintContext
<< "layout: unable to verify that the coordinate matches
the current thread. "
<< "Use .view() + .local() to decompose thread and memory
axes.";
}
- auto res = layout.Canonicalize().Apply(indices, tensor.shape());
+ auto res = layout.Canonicalize().Apply(indices, tensor.shape(),
GetSyncedAnalyzer());
TVM_FFI_ICHECK_EQ(res.size(), 1) << "Expected a single element offset";
return {GetSyncedAnalyzer()->Simplify((*res.begin()).second)};
}
diff --git a/src/sym/analyzer.cc b/src/sym/analyzer.cc
index cf96a3040..60bfcfb8e 100644
--- a/src/sym/analyzer.cc
+++ b/src/sym/analyzer.cc
@@ -57,6 +57,8 @@ AnalyzerObj::AnalyzerObj()
int_set(this),
z3_prover(this) {}
+void AnalyzerObj::SetZ3DefaultFallback(bool enabled) { z3_default_fallback_ =
enabled; }
+
ConstraintContext AnalyzerObj::GetConstraintContext() const {
// Never return the sealed internal position: callers may extend this child.
return last_snapshot_->Snapshot();
@@ -313,9 +315,10 @@ bool AnalyzerObj::CanProve(const PrimExpr& expr,
ProofStrength strength) {
}
}
- // Native analysis runs first. Z3 is the final best-effort fallback at every
- // strength, including default queries from recursive simplification.
- if (z3_prover.CanProve(simplified)) {
+ // Native analysis runs first. Explicit stronger queries retain the Z3
fallback
+ // even when it is disabled for default queries from recursive
simplification.
+ if ((z3_default_fallback_ || strength >= ProofStrength::kSymbolicBound) &&
+ z3_prover.CanProve(simplified)) {
return true;
}
return false;
@@ -344,6 +347,7 @@ PrimExpr AnalyzerObj::Simplify(const PrimExpr& expr, int
steps) {
Analyzer AnalyzerObj::Clone() const {
Analyzer cloned;
+ cloned->SetZ3DefaultFallback(z3_default_fallback_);
cloned->const_int_bound.CopyFrom(this->const_int_bound);
cloned->modular_set.CopyFrom(this->modular_set);
cloned->rewrite_simplify.CopyFrom(this->rewrite_simplify);
@@ -360,6 +364,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef()
.def("sym.Analyzer", []() { return Analyzer(); })
.def("sym.AnalyzerClone", [](Analyzer analyzer) { return
analyzer->Clone(); })
+ .def_method("sym.AnalyzerSetZ3DefaultFallback",
&AnalyzerObj::SetZ3DefaultFallback)
.def_method("sym.AnalyzerSyncConstraintContext",
&AnalyzerObj::SyncConstraintContext)
.def_method("sym.AnalyzerGetConstraintContext",
&AnalyzerObj::GetConstraintContext)
.def("sym.AnalyzerConstIntBound",
diff --git a/src/tirx/ir/layout/compose_layout.cc
b/src/tirx/ir/layout/compose_layout.cc
index ed712d83e..3bc010e55 100644
--- a/src/tirx/ir/layout/compose_layout.cc
+++ b/src/tirx/ir/layout/compose_layout.cc
@@ -29,7 +29,8 @@ using namespace tvm::prim;
namespace {
-PrimExpr ApplyFullSwizzle(const ComposeLayoutObj* layout, const PrimExpr& m) {
+PrimExpr ApplyFullSwizzle(const ComposeLayoutObj* layout, const PrimExpr& m,
+ sym::AnalyzerObj* analyzer) {
auto swizzle = [&](const PrimExpr& x) -> PrimExpr {
if (layout->swizzle_inner) {
return x ^ ((x & layout->outer_mask) >> layout->atom_len);
@@ -37,7 +38,6 @@ PrimExpr ApplyFullSwizzle(const ComposeLayoutObj* layout,
const PrimExpr& m) {
return x ^ ((x & layout->inner_mask) << layout->atom_len);
};
int base = 1 << layout->per_element;
- sym::Analyzer analyzer;
PrimVar m_once("compose_m", m.ty());
PrimExpr quotient = floordiv(m_once, base);
PrimVar quotient_once("compose_q", quotient.ty());
@@ -46,7 +46,7 @@ PrimExpr ApplyFullSwizzle(const ComposeLayoutObj* layout,
const PrimExpr& m) {
return prim::Let(m_once, m, prim::Let(quotient_once, quotient, body));
}
-void AddExpr(std::optional<PrimExpr>* sum, const PrimExpr& term, const
sym::Analyzer& analyzer) {
+void AddExpr(std::optional<PrimExpr>* sum, const PrimExpr& term,
sym::AnalyzerObj* analyzer) {
if (IsZero(term)) return;
if (sum->has_value()) {
*sum = analyzer->Simplify(sum->value() + term);
@@ -74,7 +74,7 @@ bool MulWithoutOverflow(int64_t lhs, int64_t rhs, int64_t*
result) {
}
void CollectOffsetTerms(const PrimExpr& expr, int sign, std::vector<PrimExpr>*
dynamic_terms,
- int64_t* constant, bool* valid, const sym::Analyzer&
analyzer) {
+ int64_t* constant, bool* valid, sym::AnalyzerObj*
analyzer) {
if (!*valid) return;
PrimExpr simplified = analyzer->Simplify(expr);
if (const auto* imm = simplified.as<IntImmObj>()) {
@@ -104,7 +104,7 @@ void CollectOffsetTerms(const PrimExpr& expr, int sign,
std::vector<PrimExpr>* d
}
std::optional<PrimExpr> DivideExactTerm(const PrimExpr& term, int64_t divisor,
- const sym::Analyzer& analyzer) {
+ sym::AnalyzerObj* analyzer) {
PrimExpr simplified = analyzer->Simplify(term);
if (const auto* imm = simplified.as<IntImmObj>()) {
if (imm->value % divisor != 0) return std::nullopt;
@@ -131,8 +131,8 @@ std::optional<PrimExpr> DivideExactTerm(const PrimExpr&
term, int64_t divisor,
}
ffi::Map<ffi::String, PrimExpr> ApplyStructured(const ComposeLayoutObj* layout,
- const TileLayout& tile,
- ffi::Array<PrimExpr> coord) {
+ const TileLayout& tile,
ffi::Array<PrimExpr> coord,
+ sym::AnalyzerObj* analyzer) {
// Normalize the bounded TileLayout terms as m = B + D + K, where B is a
// multiple of one atom, D is dynamic and contained in that atom, and K is
// the constant remainder. If D + K is provably carry-free, the swizzle is
@@ -146,18 +146,17 @@ ffi::Map<ffi::String, PrimExpr> ApplyStructured(const
ComposeLayoutObj* layout,
TVM_FFI_ICHECK_EQ(coord.size(), tile->shard.size())
<< "Coordinate size must match the number of shard axes";
- sym::Analyzer analyzer;
for (size_t i = 0; i < tile->shard.size(); ++i) {
if (analyzer->CanProveEqual(tile->shard[i]->extent, 1)) {
coord.Set(i, IntImm(coord[i].ty(), 0));
}
}
- auto mapped = tile.Apply(coord);
+ auto mapped = tile.Apply(coord, analyzer);
TVM_FFI_ICHECK(mapped.size() == 1 && mapped.find("m") != mapped.end());
PrimExpr m = mapped["m"];
auto fallback = [&]() {
- return ffi::Map<ffi::String, PrimExpr>{{"m", ApplyFullSwizzle(layout, m)}};
+ return ffi::Map<ffi::String, PrimExpr>{{"m", ApplyFullSwizzle(layout, m,
analyzer)}};
};
if (layout->swizzle_len == 0) return mapped;
@@ -347,13 +346,14 @@ PrimExpr layout_detail::GetSpan(const ComposeLayoutObj*
layout,
}
ffi::Map<ffi::String, PrimExpr> layout_detail::Apply(const ComposeLayoutObj*
layout,
- ffi::Array<PrimExpr>
coord) {
- return ApplyStructured(layout, layout->tile_layout, std::move(coord));
+ ffi::Array<PrimExpr>
coord,
+ sym::AnalyzerObj*
analyzer) {
+ return ApplyStructured(layout, layout->tile_layout, std::move(coord),
analyzer);
}
-ffi::Map<ffi::String, PrimExpr> layout_detail::Apply(const ComposeLayoutObj*
layout,
- PrimExpr coord) {
- auto res = layout->tile_layout.Apply(coord);
+ffi::Map<ffi::String, PrimExpr> layout_detail::Apply(const ComposeLayoutObj*
layout, PrimExpr coord,
+ sym::AnalyzerObj*
analyzer) {
+ auto res = layout->tile_layout.Apply(coord, analyzer);
TVM_FFI_ICHECK(res.size() == 1 && res.find("m") != res.end());
auto m = res["m"];
// Inline the swizzle XOR (formerly SwizzleLayoutNode::Apply): the swizzle
@@ -366,21 +366,21 @@ ffi::Map<ffi::String, PrimExpr>
layout_detail::Apply(const ComposeLayoutObj* lay
}
};
auto base = 1 << layout->per_element;
- sym::Analyzer analyzer;
return {
{"m", analyzer->Simplify((f(floordiv(m, base)) << layout->per_element) +
floormod(m, base))}};
}
ffi::Map<ffi::String, PrimExpr> layout_detail::Apply(const ComposeLayoutObj*
layout,
const
ffi::Array<PrimExpr>& coord,
- const
ffi::Array<PrimExpr>& shape) {
+ const
ffi::Array<PrimExpr>& shape,
+ sym::AnalyzerObj*
analyzer) {
TVM_FFI_ICHECK_EQ(coord.size(), shape.size())
<< "ValueError: The size of coord and shape should be equal";
- auto grouped = TryGroup(layout->tile_layout, shape);
+ auto grouped = TryGroup(layout->tile_layout, shape, analyzer);
if (!grouped.has_value()) {
- auto mapped = layout->tile_layout.Apply(coord, shape);
+ auto mapped = layout->tile_layout.Apply(coord, shape, analyzer);
TVM_FFI_ICHECK(mapped.size() == 1 && mapped.find("m") != mapped.end());
- return {{"m", ApplyFullSwizzle(layout, mapped["m"])}};
+ return {{"m", ApplyFullSwizzle(layout, mapped["m"], analyzer)}};
}
const auto& [grouped_tile, separators] = grouped.value();
@@ -402,7 +402,7 @@ ffi::Map<ffi::String, PrimExpr> layout_detail::Apply(const
ComposeLayoutObj* lay
shard_coord.Set(i, split[j]);
}
}
- return ApplyStructured(layout, grouped_tile, std::move(shard_coord));
+ return ApplyStructured(layout, grouped_tile, std::move(shard_coord),
analyzer);
}
Layout layout_detail::Canonicalize(const ComposeLayoutObj* layout) {
diff --git a/src/tirx/ir/layout/layout.cc b/src/tirx/ir/layout/layout.cc
index e8259a103..9876f5e2b 100644
--- a/src/tirx/ir/layout/layout.cc
+++ b/src/tirx/ir/layout/layout.cc
@@ -52,20 +52,27 @@ PrimExpr Layout::GetSpan(ffi::Optional<ffi::String>
axis_name) const {
[&](const auto* node) { return
layout_detail::GetSpan(node, axis_name); });
}
-ffi::Map<ffi::String, PrimExpr> Layout::Apply(ffi::Array<PrimExpr> coord)
const {
- return DispatchLayout(
- *this, [&](const auto* node) { return layout_detail::Apply(node,
std::move(coord)); });
+ffi::Map<ffi::String, PrimExpr> Layout::Apply(ffi::Array<PrimExpr> coord,
+ sym::AnalyzerObj* analyzer)
const {
+ TVM_FFI_ICHECK(analyzer != nullptr);
+ return DispatchLayout(*this, [&](const auto* node) {
+ return layout_detail::Apply(node, std::move(coord), analyzer);
+ });
}
-ffi::Map<ffi::String, PrimExpr> Layout::Apply(PrimExpr coord) const {
- return DispatchLayout(
- *this, [&](const auto* node) { return layout_detail::Apply(node,
std::move(coord)); });
+ffi::Map<ffi::String, PrimExpr> Layout::Apply(PrimExpr coord,
sym::AnalyzerObj* analyzer) const {
+ TVM_FFI_ICHECK(analyzer != nullptr);
+ return DispatchLayout(*this, [&](const auto* node) {
+ return layout_detail::Apply(node, std::move(coord), analyzer);
+ });
}
ffi::Map<ffi::String, PrimExpr> Layout::Apply(const ffi::Array<PrimExpr>&
coord,
- const ffi::Array<PrimExpr>&
shape) const {
- return DispatchLayout(*this,
- [&](const auto* node) { return
layout_detail::Apply(node, coord, shape); });
+ const ffi::Array<PrimExpr>&
shape,
+ sym::AnalyzerObj* analyzer)
const {
+ TVM_FFI_ICHECK(analyzer != nullptr);
+ return DispatchLayout(
+ *this, [&](const auto* node) { return layout_detail::Apply(node, coord,
shape, analyzer); });
}
Layout Layout::Canonicalize() const {
@@ -90,13 +97,16 @@ TVM_FFI_STATIC_INIT_BLOCK() {
return layout.GetSpan(axis_name);
});
def.def("tirx.LayoutApplyWithShape",
- [](Layout layout, ffi::Array<PrimExpr> coord, ffi::Array<PrimExpr>
shape) {
- return layout.Apply(coord, shape);
- });
+ [](Layout layout, ffi::Array<PrimExpr> coord, ffi::Array<PrimExpr>
shape,
+ const sym::Analyzer& analyzer) { return layout.Apply(coord,
shape, analyzer.get()); });
def.def("tirx.LayoutApply",
- [](Layout layout, ffi::Array<PrimExpr> coord) { return
layout.Apply(coord); });
+ [](Layout layout, ffi::Array<PrimExpr> coord, const sym::Analyzer&
analyzer) {
+ return layout.Apply(std::move(coord), analyzer.get());
+ });
def.def("tirx.LayoutApplyLinear",
- [](Layout layout, PrimExpr coord) { return layout.Apply(coord); });
+ [](Layout layout, PrimExpr coord, const sym::Analyzer& analyzer) {
+ return layout.Apply(std::move(coord), analyzer.get());
+ });
def.def("tirx.LayoutCanonicalize", [](Layout layout) { return
layout.Canonicalize(); });
def.def("tirx.LayoutTile", [](Layout layout, TileLayout outer,
ffi::Array<PrimExpr> outer_shape,
ffi::Array<PrimExpr> inner_shape) {
diff --git a/src/tirx/ir/layout/tile_core.cc b/src/tirx/ir/layout/tile_core.cc
index 2aa72d120..035dc0780 100644
--- a/src/tirx/ir/layout/tile_core.cc
+++ b/src/tirx/ir/layout/tile_core.cc
@@ -218,13 +218,16 @@ PrimExpr layout_detail::GetSpan(const TileLayoutObj*
layout, ffi::Optional<ffi::
return analyzer->Simplify(result);
}
-ffi::Map<ffi::String, PrimExpr> layout_detail::Apply(const TileLayoutObj*
layout, PrimExpr coord) {
- return Apply(layout, SplitCoord(coord,
ffi::GetRef<TileLayout>(layout).GetShardShape()));
+ffi::Map<ffi::String, PrimExpr> layout_detail::Apply(const TileLayoutObj*
layout, PrimExpr coord,
+ sym::AnalyzerObj*
analyzer) {
+ return Apply(layout, SplitCoord(coord,
ffi::GetRef<TileLayout>(layout).GetShardShape()),
+ analyzer);
}
ffi::Map<ffi::String, PrimExpr> layout_detail::Apply(const TileLayoutObj*
layout,
const
ffi::Array<PrimExpr>& coord,
- const
ffi::Array<PrimExpr>& shape) {
+ const
ffi::Array<PrimExpr>& shape,
+ sym::AnalyzerObj*
analyzer) {
TVM_FFI_ICHECK_EQ(coord.size(), shape.size())
<< "ValueError: The size of coord and shape should be equal";
// Group-first path: if this layout can be regrouped by ``shape`` (the input
@@ -234,7 +237,7 @@ ffi::Map<ffi::String, PrimExpr> layout_detail::Apply(const
TileLayoutObj* layout
// symbolic form small (local mod/divs) and avoids the cross-dim noise of the
// flatten+split-against-shard-shape round-trip. Equivalent numerical output,
// much friendlier for sym.Analyzer downstream.
- if (auto grouped_opt = TryGroup(ffi::GetRef<TileLayout>(layout), shape);
+ if (auto grouped_opt = TryGroup(ffi::GetRef<TileLayout>(layout), shape,
analyzer);
grouped_opt.has_value()) {
auto& [grouped, seps] = *grouped_opt;
ffi::Array<PrimExpr> per_shard_coords;
@@ -255,15 +258,15 @@ ffi::Map<ffi::String, PrimExpr>
layout_detail::Apply(const TileLayoutObj* layout
per_shard_coords.Set(i, split[j]);
}
}
- return grouped.Apply(per_shard_coords);
+ return grouped.Apply(per_shard_coords, analyzer);
}
// Fallback: flatten across input shape then split against shard shape.
- return Apply(layout, FlattenCoord(coord, shape));
+ return Apply(layout, FlattenCoord(coord, shape), analyzer);
}
ffi::Map<ffi::String, PrimExpr> layout_detail::Apply(const TileLayoutObj*
layout,
- Array<PrimExpr> coord) {
- sym::Analyzer analyzer;
+ Array<PrimExpr> coord,
+ sym::AnalyzerObj*
analyzer) {
TVM_FFI_ICHECK_EQ(coord.size(), layout->shard.size())
<< "Coordinate size must match the number of shard axes";
std::unordered_map<ffi::String, PrimExpr> result;
diff --git a/src/tirx/ir/layout/tile_internal.h
b/src/tirx/ir/layout/tile_internal.h
index a067e5da6..028ae7970 100644
--- a/src/tirx/ir/layout/tile_internal.h
+++ b/src/tirx/ir/layout/tile_internal.h
@@ -35,11 +35,14 @@ namespace layout_detail {
void CheckWellFormed(const TileLayoutObj* layout);
PrimExpr GetSize(const TileLayoutObj* layout, ffi::Optional<ffi::String>
axis_name);
PrimExpr GetSpan(const TileLayoutObj* layout, ffi::Optional<ffi::String>
axis_name);
-ffi::Map<ffi::String, PrimExpr> Apply(const TileLayoutObj* layout,
ffi::Array<PrimExpr> coord);
-ffi::Map<ffi::String, PrimExpr> Apply(const TileLayoutObj* layout, PrimExpr
coord);
+ffi::Map<ffi::String, PrimExpr> Apply(const TileLayoutObj* layout,
ffi::Array<PrimExpr> coord,
+ sym::AnalyzerObj* analyzer);
+ffi::Map<ffi::String, PrimExpr> Apply(const TileLayoutObj* layout, PrimExpr
coord,
+ sym::AnalyzerObj* analyzer);
ffi::Map<ffi::String, PrimExpr> Apply(const TileLayoutObj* layout,
const ffi::Array<PrimExpr>& coord,
- const ffi::Array<PrimExpr>& shape);
+ const ffi::Array<PrimExpr>& shape,
+ sym::AnalyzerObj* analyzer);
Layout Canonicalize(const TileLayoutObj* layout);
Layout Tile(const TileLayoutObj* layout, const TileLayout& outer,
const ffi::Array<PrimExpr>& outer_shape, const
ffi::Array<PrimExpr>& inner_shape);
@@ -47,11 +50,14 @@ Layout Tile(const TileLayoutObj* layout, const TileLayout&
outer,
void CheckWellFormed(const ComposeLayoutObj* layout);
PrimExpr GetSize(const ComposeLayoutObj* layout, ffi::Optional<ffi::String>
axis_name);
PrimExpr GetSpan(const ComposeLayoutObj* layout, ffi::Optional<ffi::String>
axis_name);
-ffi::Map<ffi::String, PrimExpr> Apply(const ComposeLayoutObj* layout,
ffi::Array<PrimExpr> coord);
-ffi::Map<ffi::String, PrimExpr> Apply(const ComposeLayoutObj* layout, PrimExpr
coord);
+ffi::Map<ffi::String, PrimExpr> Apply(const ComposeLayoutObj* layout,
ffi::Array<PrimExpr> coord,
+ sym::AnalyzerObj* analyzer);
+ffi::Map<ffi::String, PrimExpr> Apply(const ComposeLayoutObj* layout, PrimExpr
coord,
+ sym::AnalyzerObj* analyzer);
ffi::Map<ffi::String, PrimExpr> Apply(const ComposeLayoutObj* layout,
const ffi::Array<PrimExpr>& coord,
- const ffi::Array<PrimExpr>& shape);
+ const ffi::Array<PrimExpr>& shape,
+ sym::AnalyzerObj* analyzer);
Layout Canonicalize(const ComposeLayoutObj* layout);
Layout Tile(const ComposeLayoutObj* layout, const TileLayout& outer,
const ffi::Array<PrimExpr>& outer_shape, const
ffi::Array<PrimExpr>& inner_shape);
@@ -65,7 +71,7 @@ std::pair<TileLayout, std::vector<int64_t>> Group(TileLayout
layout,
// Same as Group but returns std::nullopt instead of fatal-checking when the
// layout cannot be regrouped by ``shape``.
std::optional<std::pair<TileLayout, std::vector<int64_t>>> TryGroup(
- TileLayout layout, const ffi::Array<PrimExpr>& shape);
+ TileLayout layout, const ffi::Array<PrimExpr>& shape, sym::AnalyzerObj*
analyzer);
} // namespace tirx
} // namespace tvm
diff --git a/src/tirx/ir/layout/tile_tile_ops.cc
b/src/tirx/ir/layout/tile_tile_ops.cc
index aa134c28d..1a5f2cb83 100644
--- a/src/tirx/ir/layout/tile_tile_ops.cc
+++ b/src/tirx/ir/layout/tile_tile_ops.cc
@@ -73,11 +73,10 @@ std::pair<TileLayout, std::vector<int64_t>>
Group(TileLayout layout,
}
std::optional<std::pair<TileLayout, std::vector<int64_t>>> TryGroup(
- TileLayout layout, const ffi::Array<PrimExpr>& shape) {
+ TileLayout layout, const ffi::Array<PrimExpr>& shape, sym::AnalyzerObj*
analyzer) {
// Same algorithm as Group but returns std::nullopt instead of ICHECK-failing
// on regroup impossibility. Used by Apply(coord, shape) to opportunistically
// pick the group-first path with a fallback to flatten+split.
- sym::Analyzer analyzer;
size_t shape_idx = 0;
PrimExpr prod = 1;
diff --git a/src/tirx/ir/tensor.cc b/src/tirx/ir/tensor.cc
index a4f10633b..f8b200210 100644
--- a/src/tirx/ir/tensor.cc
+++ b/src/tirx/ir/tensor.cc
@@ -395,8 +395,9 @@ ffi::Array<PrimExpr>
TensorTypeObj::ElemOffset(ffi::Array<PrimExpr> input_indice
<< input_indices.size() << "-dimensional indices provided.";
sym::Analyzer ana;
+ ana->SetZ3DefaultFallback(false);
if (layout.has_value()) {
- auto mapped = layout.value().Apply(input_indices, shape);
+ auto mapped = layout.value().Apply(input_indices, shape, ana.get());
if (mapped.empty() && shape.empty()) {
return {IntImm(DefaultIndexType(), 0)};
}
diff --git a/src/tirx/tile/lower.cc b/src/tirx/tile/lower.cc
index fa0b4c00a..951bc6cf9 100644
--- a/src/tirx/tile/lower.cc
+++ b/src/tirx/tile/lower.cc
@@ -174,7 +174,9 @@ class TileDispatcher : public StmtExprMutator {
public:
using StmtExprMutator::Mutate;
using StmtExprMutator::Mutate_;
- explicit TileDispatcher(const Target& target) : target_(target) {}
+ explicit TileDispatcher(const Target& target) : target_(target) {
+ analyzer_->SetZ3DefaultFallback(false);
+ }
static Stmt LowerOpCalls(const Stmt& stmt, const Target& target) {
return ffi::make_object<TileDispatcher>(target)
diff --git a/src/tirx/transform/lower_tensor_layouts.cc
b/src/tirx/transform/lower_tensor_layouts.cc
index fbe259110..28e5fef31 100644
--- a/src/tirx/transform/lower_tensor_layouts.cc
+++ b/src/tirx/transform/lower_tensor_layouts.cc
@@ -322,7 +322,7 @@ class LayoutApplier : public
StmtExprMutatorWithConstraintContext {
auto tile = layout.as_or_throw<TileLayout>();
layout = TileLayout(tile->shard, tile->replica, {});
}
- auto res = layout.Canonicalize().Apply(indices, logical_type->shape);
+ auto res = layout.Canonicalize().Apply(indices, logical_type->shape,
GetSyncedAnalyzer());
if (res.empty() && logical_type->shape.empty()) {
return {IntImm(logical_type->DefaultIndexType(), 0)};
}
diff --git a/tests/python/sym/test_z3.py b/tests/python/sym/test_z3.py
index bbde37e69..42c2e0424 100644
--- a/tests/python/sym/test_z3.py
+++ b/tests/python/sym/test_z3.py
@@ -898,13 +898,22 @@ def test_z3_default_fallback_nonlinear_inequality():
y = tirx.Var("y", "int32")
# x*x + y*y - 2*x*y == (x-y)**2 >= 0 in the mathematical integer
# model used for wide signed indices. This does not model int32 overflow.
- # A fresh analyzer must prove this without an option or explicit strength
- # when Z3 is available. Native analysis alone cannot currently prove it.
+ # A fresh analyzer uses Z3 by default; disabling DEFAULT fallback preserves
+ # the stronger query's Z3 capability and survives cloning independently.
expr = x * x + y * y >= 2 * (x * y)
assert analyzer.can_prove(expr) == analyzer.is_z3_enabled
+ analyzer.set_z3_default_fallback(False)
+ assert not analyzer.can_prove(expr)
+ assert analyzer.can_prove(expr, SB) == analyzer.is_z3_enabled
+ cloned = analyzer.clone()
+ analyzer.set_z3_default_fallback(True)
+ assert analyzer.can_prove(expr) == analyzer.is_z3_enabled
+ assert not cloned.can_prove(expr)
+ cloned.set_z3_default_fallback(True)
+ assert cloned.can_prove(expr) == analyzer.is_z3_enabled
assert analyzer.can_prove(x == x)
# Equality at x == y is a counterexample to the strict inequality.
- assert not analyzer.can_prove(x * x + y * y > 2 * (x * y))
+ assert not cloned.can_prove(x * x + y * y > 2 * (x * y))
def test_z3_default_fallback_clone_preserves_independent_resource_limit():
diff --git a/tests/python/tirx/test_layout.py b/tests/python/tirx/test_layout.py
index 5dba9394e..66b085ed5 100644
--- a/tests/python/tirx/test_layout.py
+++ b/tests/python/tirx/test_layout.py
@@ -1474,6 +1474,22 @@ def
test_compose_structured_apply_falls_back_to_full_swizzle(case):
assert _evaluate_layout_expr(actual, values) ==
_evaluate_layout_expr(expected, values)
+def test_layout_apply_with_reused_analyzer():
+ tile = TileLayout(S[(8, 8) : (64, 1)] + 8)
+ analyzer = tvm_tirx.sym.Analyzer()
+ x, y = Var("x", "int32"), Var("y", "int32")
+ for layout in (tile, ComposeLayout(3, 3, 3, tile), ComposeLayout(3, 3, 3,
tile, False)):
+ for coords, options in (
+ ((x, y), {}),
+ ((x * 8 + y,), {}),
+ ((x, y), {"shape": (8, 8)}),
+ ):
+ expected = layout.apply(*coords, **options)
+ assert_structural_equal(layout.apply(*coords, analyzer=analyzer,
**options), expected)
+ assert not analyzer.can_prove(x >= 0)
+ assert not analyzer.can_prove(y < 8)
+
+
def test_compose_structured_apply_emits_shared_base_address_chain():
x = Var("x", ty="int32")
y = Var("y", ty="int32")
diff --git a/tests/python/tirx/tile/cuda/test_explicit_patterns.py
b/tests/python/tirx/tile/cuda/test_explicit_patterns.py
index 845b7e06d..a6a4f09dc 100644
--- a/tests/python/tirx/tile/cuda/test_explicit_patterns.py
+++ b/tests/python/tirx/tile/cuda/test_explicit_patterns.py
@@ -150,6 +150,21 @@ def _dispatch(
return run_dispatch(T.cuda.tile.apply(T.ptx[opcode], dst, src,
scope=scope), ctx)
+def test_layout_analysis_is_separate_from_verification():
+ from tvm_tirx.backend.cuda.tile.patterns.common import Context
+
+ sctx = SimpleNamespace(launch_params={}, var_range_map={},
scope_kind="thread", intra={})
+ ctx = Context(sctx)
+ x = Var("x", "int32")
+ ctx.add_range(x, tvm_tirx.ir.Range(0, 4))
+ assert ctx.analyzer.can_prove(x < 4)
+ assert not ctx.layout_analyzer.can_prove(x < 4)
+ a, b = Var("a", "int32"), Var("b", "int32")
+ predicate = a * a + b * b >= 2 * a * b
+ assert not ctx.analyzer.can_prove(predicate)
+ assert not ctx.layout_analyzer.can_prove(predicate)
+
+
def test_contract_rejections():
with pytest.raises(RuntimeError, match="consecutive elements"):
_dispatch(src_layout=TileLayout(S[(2, 4) : (8, 2)]))
