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.git
The following commit(s) were added to refs/heads/main by this push:
new 46bc8c6a38 [Fix][Arith] Give each materialized Z3 solver a private
context (#20221)
46bc8c6a38 is described below
commit 46bc8c6a38e41b55fbc41f796e77773cb7b4f887
Author: Shushi Hong <[email protected]>
AuthorDate: Sat Aug 29 10:38:40 2026 -0400
[Fix][Arith] Give each materialized Z3 solver a private context (#20221)
This PR gives each lazily materialized Z3 solver a private context,
preventing shared AST state from affecting results near the
deterministic rlimit. The legacy Z3ContextScope APIs are retained as
no-ops for compatibility.
Co-authored-by: LeiWang1999 <[email protected]>
---
include/tvm/arith/analyzer.h | 8 +++--
python/tvm/arith/analyzer.py | 6 ++--
src/arith/z3_prover.cc | 83 ++++++++++++++++++--------------------------
3 files changed, 42 insertions(+), 55 deletions(-)
diff --git a/include/tvm/arith/analyzer.h b/include/tvm/arith/analyzer.h
index dbde2cf0df..4d9c8b1365 100644
--- a/include/tvm/arith/analyzer.h
+++ b/include/tvm/arith/analyzer.h
@@ -599,14 +599,16 @@ class IntSetAnalyzer {
/*!
* \brief Enter a thread-local Z3 context scope.
*
- * The outermost scope creates a fresh Z3 context. Nested scopes on the same
- * thread reuse that context so all Analyzers created during one compilation
- * can share it.
+ * Deprecated no-op. Every materialized Z3 solver owns a private context,
+ * which subsumes the per-compilation isolation these scopes provided; kept
+ * only until the remaining downstream call sites are removed.
*/
TVM_DLL void EnterZ3ContextScope();
/*!
* \brief Exit the current thread-local Z3 context scope.
+ *
+ * Deprecated no-op, see EnterZ3ContextScope.
*/
TVM_DLL void ExitZ3ContextScope();
diff --git a/python/tvm/arith/analyzer.py b/python/tvm/arith/analyzer.py
index 453d1aa33c..ad086518fb 100644
--- a/python/tvm/arith/analyzer.py
+++ b/python/tvm/arith/analyzer.py
@@ -118,10 +118,10 @@ class ConstraintScope:
class Z3ContextScope:
- """Share one fresh Z3 context across Analyzers created in this scope.
+ """Deprecated no-op retained for compatibility.
- The outermost scope creates the context. Nested scopes on the same thread
- reuse it, and exiting the outermost scope releases its ownership.
+ Every materialized Z3 solver now owns a private context, which provides
+ per-analyzer isolation without an explicit compilation scope.
"""
def __enter__(self):
diff --git a/src/arith/z3_prover.cc b/src/arith/z3_prover.cc
index 083bc99381..49644c78ce 100644
--- a/src/arith/z3_prover.cc
+++ b/src/arith/z3_prover.cc
@@ -95,48 +95,17 @@ struct Namespace {
}
};
-struct Z3ContextState {
- std::shared_ptr<z3::context> fallback_context;
- std::shared_ptr<z3::context> scoped_context;
- size_t scope_depth{0};
-};
-
-Z3ContextState& GetZ3ContextState() {
- static thread_local Z3ContextState state;
- return state;
-}
-
-std::shared_ptr<z3::context> GetCurrentZ3Context() {
- auto& state = GetZ3ContextState();
- if (state.scope_depth != 0) {
- TVM_FFI_ICHECK(state.scoped_context != nullptr);
- return state.scoped_context;
- }
- if (state.fallback_context == nullptr) {
- state.fallback_context = std::make_shared<z3::context>();
- }
- return state.fallback_context;
-}
-
} // namespace
-void EnterZ3ContextScope() {
- auto& state = GetZ3ContextState();
- if (state.scope_depth == 0) {
- state.scoped_context = std::make_shared<z3::context>();
- }
- ++state.scope_depth;
-}
+// Deprecated no-ops, kept only for callers of
+// arith.EnterZ3ContextScope / arith.ExitZ3ContextScope (Z3ContextScope in
+// Python). Shared per-compile Z3 contexts are gone: every materialized solver
+// owns a private context (see Z3Prover::Impl::Materialize), which subsumes the
+// per-compilation isolation these scopes provided. Remove together with the
+// remaining call sites.
+void EnterZ3ContextScope() {}
-void ExitZ3ContextScope() {
- auto& state = GetZ3ContextState();
- TVM_FFI_ICHECK_GT(state.scope_depth, 0U)
- << "ExitZ3ContextScope called without a matching EnterZ3ContextScope";
- --state.scope_depth;
- if (state.scope_depth == 0) {
- state.scoped_context.reset();
- }
-}
+void ExitZ3ContextScope() {}
class Z3Prover::Impl : ExprFunctor<z3::expr(const Expr&)> {
public:
@@ -144,8 +113,10 @@ class Z3Prover::Impl : ExprFunctor<z3::expr(const Expr&)> {
using Self = Z3Prover::Impl;
AnalyzerObj* analyzer;
- // Analyzers created in one compile scope share a context. Keeping the
pointer
- // on each prover also lets Analyzers and their clones outlive that scope.
+ /// @brief Z3 context owning every handle of this prover.
+ // Null until the solver materializes; Materialize creates a context private
+ // to this prover, which clones of a materialized prover then share so their
+ // copied handles stay valid.
std::shared_ptr<z3::context> ctx;
/// @brief Z3 solver instance
@@ -207,14 +178,14 @@ class Z3Prover::Impl : ExprFunctor<z3::expr(const Expr&)>
{
return result;
}
- Impl(AnalyzerObj* parent) : analyzer(parent), ctx(GetCurrentZ3Context()) {
- // The solver is created lazily (see Materialize). Analyzers are
- // constructed in large numbers as cheap scratch objects, and only a tiny
- // fraction ever reaches the Z3 fallback of CanProve; creating the solver
- // (and translating every Bind into z3 constraints) eagerly made every
- // Analyzer pay Z3's initialization cost up front. Until the first
- // operation that actually needs the solver, Bind/EnterConstraint only
- // journal their arguments into scope_stack_.
+ Impl(AnalyzerObj* parent) : analyzer(parent) {
+ // The solver (and its context) is created lazily (see Materialize).
+ // Analyzers are constructed in large numbers as cheap scratch objects,
+ // and only a tiny fraction ever reaches the Z3 fallback of CanProve;
+ // creating the solver (and translating every Bind into Z3 constraints)
+ // eagerly made every Analyzer pay Z3's initialization cost up front.
+ // Until the first operation that actually needs the solver,
+ // Bind/EnterConstraint only journal their arguments into scope_stack_.
scope_stack_.push_back({});
scope_side_effects_.push_back({});
// use rlimit, not timeout to ensure deterministic behavior
@@ -233,6 +204,14 @@ class Z3Prover::Impl : ExprFunctor<z3::expr(const Expr&)> {
/// for these entries, so CanProve answers are unchanged.
void Materialize() {
if (solver) return;
+ // The solver gets its own context: verdicts under the deterministic
+ // rlimit budget depend on the context's accumulated AST state (ids feed
+ // the search heuristics), so sharing a context couples every solver's
+ // borderline answers to whichever analyzers happened to materialize
+ // before it. A private context makes each analyzer's answers a pure
+ // function of its own journal and query history. Unmaterialized provers
+ // hold no context (and no Z3 state) at all.
+ ctx = std::make_shared<z3::context>();
solver.emplace(CreateSolver(*ctx));
if (timeout_ms != UINT_MAX) {
solver->set("timeout", timeout_ms);
@@ -410,6 +389,12 @@ class Z3Prover::Impl : ExprFunctor<z3::expr(const Expr&)> {
/// @brief Binded
/// @brief Bind a variable to a value or a range
+ /// A bind journaled inside a live constraint scope is dropped with that
+ /// scope. The eager path differed only cosmetically: it kept the var's memo
+ /// translation after the scope's solver frame (and with it the var's range
+ /// assertions) was popped, leaving an unconstrained placeholder --
+ /// observably the same as translating the var as a fresh free variable on
+ /// demand, which is what a later query on the journal does.
void Bind(const Var& var, const PrimExpr& value, bool allow_override =
false) {
if (!IsZ3SupportedExpr(var.get())) return;
scope_stack_.back().push_back(Scope{Scope::BindValue, var, value});
