[ARITH] Add optional Z3-backed proving to Analyzer#19667
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request integrates Z3 SMT solver support into TVM's arithmetic analyzer, adding a new Z3Prover sub-analyzer, updating CMake configurations, and exposing Z3-related APIs to Python. Feedback on the implementation highlights a critical soundness issue where shift assertions permanently poison the global solver context, an optimization failure in range binding due to unfolded expressions, and thread-safety limitations with the static thread-local Z3 context. Additionally, minor improvements are suggested to resolve an implicit double-to-unsigned conversion and remove redundant header includes.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| // Add constraint that shift amount should be non-negative | ||
| // This is a common assumption in many programming languages | ||
| solver.add(b_expr >= 0); | ||
|
|
||
| // Also limit shift amount to avoid unrealistic large shifts | ||
| // We'll limit to 64 bits (reasonable for most use cases) | ||
| solver.add(b_expr < 64); |
There was a problem hiding this comment.
Soundness Issue with Shift Amount Assertions
Calling solver.add() inside VisitShiftOp permanently adds the assertions b_expr >= 0 and b_expr < 64 to the global solver context. Because CanProve does not use solver.push() and solver.pop(), these assertions persist across different CanProve calls.
This poisons the solver for all subsequent unrelated proofs involving the same variable b, leading to unsound results (e.g., proving b < 100 is always true even when b is completely unbounded in other contexts).
To fix this, we should rely on Z3's native bit-vector shift behavior without adding hard assertions to the solver.
// Rely on Z3's native bit-vector shift behavior without adding hard assertions
// to the solver, which would otherwise permanently poison the global solver context.
There was a problem hiding this comment.
This makes sense. Could you have a correction?
| if (tirx::is_const_int(range->min) && tirx::is_const_int(range->min + range->extent)) { | ||
| int64_t min_value = *tirx::as_const_int(range->min); | ||
| int64_t max_value = *tirx::as_const_int(range->min + range->extent); | ||
| if (min_value < max_value) { | ||
| solver.add(ctx->int_val(min_value) <= var_expr); | ||
| solver.add(var_expr < ctx->int_val(max_value)); | ||
| } |
There was a problem hiding this comment.
Constant Range Optimization is Never Triggered
range->min + range->extent is constructed as a new AddNode and is not automatically folded. Therefore, tirx::is_const_int(range->min + range->extent) will always return false, meaning the optimized path that adds simple bounds directly to the solver is never executed.
Instead, it always falls back to the else branch, generating a much more complex Z3 expression involving ConvertBool and or/and operators, which slows down Z3 solving.
We should check is_const_int on range->min and range->extent individually, and sum their values in C++.
if (tirx::is_const_int(range->min) && tirx::is_const_int(range->extent)) {
int64_t min_value = *tirx::as_const_int(range->min);
int64_t extent_value = *tirx::as_const_int(range->extent);
int64_t max_value = min_value + extent_value;
if (min_value < max_value) {
solver.add(ctx->int_val(min_value) <= var_expr);
solver.add(var_expr < ctx->int_val(max_value));
}
}| void CopyFrom(const Self& other_) { | ||
| // 1. create a new solver | ||
| // because this->solver depends on this->ctx | ||
| // we need to deconstruct the old solver, and create a new one depending on other_.ctx | ||
| solver = CreateSolver(*other_.ctx); | ||
| // 2. copy the context | ||
| // the context is a shared_ptr, we can just copy the pointer | ||
| ctx = other_.ctx; |
There was a problem hiding this comment.
Thread-Safety and Cross-Thread Copying Limitation
ctx is declared as inline static thread_local std::shared_ptr<z3::context>. Because it is static thread_local, any access to other_.ctx (even if other_ was created on a different thread) will resolve to the thread-local variable of the current thread executing the code.
Therefore, ctx = other_.ctx; is a redundant self-assignment on the same thread, and fails to copy or reference the actual context of other_ if it was created on a different thread. If Z3Prover is copied across threads, attempting to add assertions/expressions from the other thread's context to the current thread's solver will cause Z3 to crash or throw an exception.
We should remove the redundant self-assignment and add a warning comment about this limitation.
void CopyFrom(const Self& other_) {
// 1. create a new solver
// because this->solver depends on this->ctx
// we need to deconstruct the old solver, and create a new one depending on other_.ctx
solver = CreateSolver(*ctx);
// 2. Note: ctx is static thread_local, so other_.ctx is the same as ctx on the current thread.
// Cross-thread copying of Z3Prover is not supported because Z3 expressions cannot be
// shared across different thread-local contexts without translation.There was a problem hiding this comment.
I think we can remove it as gemini suggested
| // Z3's implementation of timeout, when setting timeout T ms, it will stop at T - 1 ms | ||
| // SetTimeoutMs(5); | ||
| // use rlimit, not timeout to ensure determinstic behavior | ||
| SetRLimit(1e4); |
There was a problem hiding this comment.
Implicit Double-to-Unsigned Conversion
SetRLimit(1e4) passes a double literal 1e4 to a function expecting unsigned rlimit. This triggers an implicit double-to-unsigned conversion.
We should use an integer literal like 10000U to avoid implicit conversion warnings.
SetRLimit(10000U);| #include <unordered_set> | ||
| #include <vector> | ||
|
|
||
| #include "tvm/arith/analyzer.h" |
There was a problem hiding this comment.
Redundant Include
#include "tvm/arith/analyzer.h" is redundant because <tvm/arith/analyzer.h> is already included at the top of the file.
| #include <tvm/tirx/expr.h> | ||
| #include <tvm/tirx/op.h> | ||
|
|
||
| #include "tvm/arith/analyzer.h" |
There was a problem hiding this comment.
Redundant Include
#include "tvm/arith/analyzer.h" is redundant because <tvm/arith/analyzer.h> is already included at the top of the file.
Ubospica
changed the title
Ubospica
changed the title
There was a problem hiding this comment.
Maybe it's better to remove z3_prover_off.cc and z3_prover_on.cc to /src/arith?
And we should have ASF license header on those two cc files
There was a problem hiding this comment.
There is no negative test currently, could you add one?
And seems the new tests don't actually run in CI. USE_Z3 defaults to OFF and no buildconfig enables it , so every job builds the Z3-off stub and _require_z3() skips all four real tests. Do we need to enable it during CI? cc @tqchen
| z3::expr VisitExpr_(const SubNode* op) override { | ||
| return VisitArith(z3::operator-, op, op->a, op->b); | ||
| } | ||
| z3::expr VisitExpr_(const MulNode* op) override { |
There was a problem hiding this comment.
VisitExpr_(DivNode) and VisitExpr_(ModNode) map TVM's truncated div/mod (round toward zero) onto Z3's native operator/ / operator%, which are Euclidean (floor / non-negative remainder for a positive divisor). These disagree once the dividend can be negative.
CanProve can return true for a predicate that is false under TVM semantics (its negation comes out unsat under Euclidean div) which is unsound. Just like FloorDiv/FloorMod via the custom floordiv/floormod helpers. Div/Mod need the same treatment with truncating helpers rather than the raw Z3 operators I think.
I saw the same place on Tilelang, seems it's the same but still think it's an issue
| } | ||
| } | ||
|
|
||
| z3::expr VisitExprDefault_(const Object* op) override { |
There was a problem hiding this comment.
Here VisitExprDefault_ does a LOG(FATAL) whenever it hits a node type it doesn't handle, which kills the whole build. That's a problem here because Z3 is just a fallback — Analyzer::CanProve only reaches it after the native analyzers have already given up. So if an expression happens to contain some unsupported node, what used to be a harmless "couldn't prove it" now turns into a hard crash. Could we just return a fresh free variable here (return Create(op);)?
| } | ||
| } | ||
|
|
||
| if (z3_prover.CanProve(simplified)) { |
There was a problem hiding this comment.
The Z3 fallback isn't exception-safe — Z3 can throw z3::exception and nothing catches it, so it escapes Analyzer::CanProve and breaks the calling pass. A fallback shouldn't turn a failed proof into a crash. it's good to wrap the body of Z3Prover::CanProve in try { ... } catch (const z3::exception&) { return false; }.
|
Z3 fallback currently ignores the strength gate, so enabling it at all (especially in deeply recursive kDefault CanProve calls) causes unnecessary Z3 invocations and a measurable perf hit; it should be gated behind a higher threshold like kSymbolicBound to avoid paying prover cost on cheap paths I think. |
|
BTW, the original Z3 integration is in this commit (Not a PR): tile-ai@e633295 |
There was a problem hiding this comment.
For the tests, do we need to cover all tests in the original PR?: tile-ai/tilelang#1367
There was a problem hiding this comment.
Using two source files to control z3_off/z3_on is a bit strange. Can we use macros to achieve this so that we can merge them into a single z3_prover.cc?
|
given this is optional registration, single file and cmake control addition of the file perhaps is the best approach, great to add testcases to cover things when z3 is detected |
Ubospica
force-pushed
the
2026-06-03-arith-z3-analyzer
branch
from
6211940 to
39f808e
Compare
Move TVM's core C++ and CUDA dialects to C++20 and update first-party helper build paths that compile TVM-generated C++ code. This keeps the codebase ready for C++20 native dependencies while leaving vendored third-party projects to manage their own minimum standards.
Limit the C++20 baseline update to TVM's core CMake configuration and docs. Leave app, JVM, web, and Python helper compile flags at their existing C++17 settings so those downstream build surfaces can migrate separately.
Replace deprecated C++ traits and make captures explicit so the C++20 baseline builds cleanly under -Werror.
Add an optional Z3 SMT solver backend to tvm::arith::Analyzer for stronger integer arithmetic proving. The integration is guarded by a new USE_Z3 CMake option (default OFF). When enabled, Analyzer::CanProve runs the existing analysis path first and only falls back to Z3 when the existing analyzers cannot prove the predicate. When disabled, a stub implementation keeps the C++ and Python APIs available without Z3.
- Remove python/tvm/_version.py: it is auto-generated by the build's VCS
versioning ("don't track in version control") and is not imported at
runtime; the package __version__ comes from python/tvm/libinfo.py.
Also add it to .gitignore so it is not re-committed.
- Remove four stray 3rdparty gitlinks (cnpy, dmlc-core, rang, zlib) that
are not real submodules (absent from .gitmodules).
USE_Z3=ON now resolves Z3 through the PyPI z3-staticlib package (PIC static libz3 + headers + CMake package files) so libtvm carries no runtime libz3 dependency. The z3-solver shared-library fallback is removed; a custom Z3 installation can still be selected via Z3_DIR/CMAKE_PREFIX_PATH. Wheel builds enable USE_Z3 by default and pull z3-staticlib as a build requirement.
Use the z3-static CMake package path first and make the rebased Z3 analyzer build cleanly with LLVM and CUDA enabled.
Ubospica
force-pushed
the
2026-06-03-arith-z3-analyzer
branch
from
ce0b276 to
98b2aec
Compare
Keep the LLVM header interface target on explicit compile options so object targets inherit llvm-config include paths reliably.
Default USE_Z3 to AUTO so source builds use z3-static when installed while CI environments without the package still build the conservative stub.
Avoid passing GCC-style -isystem options to cl while keeping explicit LLVM include propagation for non-MSVC builds.
Clarify why the LLVM header target uses include directories instead of GCC-style -isystem options on MSVC.
4 participants
Summary
This PR adds a Z3 SMT solver backend to
tvm::arith::Analyzerfor stronger integer arithmetic proving.The integration is guarded by
USE_Z3, which defaults toAUTO. In the default mode, TVM enables Z3 when the static Z3 development artifacts are available and otherwise builds the conservative stub implementation. When Z3 is enabled,Analyzer::CanProveruns the existing TVM arithmetic analysis path first, then falls back to Z3 only when the existing analyzers cannot prove the predicate and the requested strength iskSymbolicBound. Z3 is linked statically from the PyPIz3-staticpackage, solibtvmdoes not need a runtimelibz3dependency.Features
USE_Z3, defaulting toAUTO.arith::Z3Proversub-analyzer owned byarith::Analyzer.Analyzer.get_smtlib2,Analyzer.set_z3_timeout_ms,Analyzer.set_z3_rlimit, andAnalyzer.get_z3_stats.min,max,select,if_then_else,let, casts, truncated division/modulo, floor division/modulo, and selected bitwise/shift operations.rlimit, withrandom_seedfixed to42.Implementation Notes
src/arith/z3_prover.cc, selected by theTVM_USE_Z3macro fromcmake/modules/contrib/Z3.cmake.cmake/modules/contrib/Z3.cmakefirst resolves the PIC staticlibz3layout provided byz3-staticusing itsz3_static.get_cmake_dir()helper, then falls back to a customZ3_DIRorCMAKE_PREFIX_PATHinstallation.USE_Z3=ONrequires Z3 to be found, whileUSE_Z3=AUTOallows source builds and CI jobs without Z3 artifacts to continue with the stub.kSymbolicBound, so the commonkDefaultpath does not pay solver cost.DivandModare translated with truncating helpers rather than Z3's Euclidean operators to stay sound for negative dividends.References
The implementation is based on the Z3 analyzer integration used in TileLang's TVM fork, with the upstream port kept scoped to TVM's arithmetic analyzer.