[ARITH] Add optional Z3-backed proving to Analyzer by Ubospica · Pull Request #19667 · apache/tvm

Ubospica · 2026-06-03T20:50:04Z

Summary

This PR adds an optional Z3 SMT solver backend to tvm::arith::Analyzer for stronger integer arithmetic proving.

The integration is disabled by default and guarded by a new USE_Z3 CMake option. When enabled, Analyzer::CanProve keeps the existing TVM arithmetic analysis path first, then falls back to Z3 only when the existing analyzers cannot prove the predicate. When disabled, a stub implementation is built so the C++ and Python APIs remain available without requiring Z3. When USE_Z3=ON, the CMake lookup can consume the PyPI z3-static package and statically link Z3, so libtvm can avoid a runtime libz3 dependency. Enabling this by default for wheels is intentionally deferred until the z3-static wheels are published.

Features

This PR adds:

Optional Z3 build support through USE_Z3, defaulting to OFF.
A new arith::Z3Prover sub-analyzer owned by arith::Analyzer.
Updated Analyzer::CanProve flow: it first runs the existing TVM logic (simplify, const-int-bound, int-set, transitive comparison) unchanged; only if that fails, and only at kSymbolicBound strength, does it ask z3_prover.CanProve(simplified), which proves the predicate by checking that its negation is unsatisfiable under the current constraints. Z3 is a pure fallback and a negative/unknown result keeps the original false.
SMT-LIB2 export for debugging and external solver reproduction.
Python APIs:
- Analyzer.get_smtlib2(expr=None)
- Analyzer.set_z3_timeout_ms(timeout_ms)
- Analyzer.set_z3_rlimit(rlimit)
- Analyzer.get_z3_stats()
C++ APIs:
- Z3Prover::CanProve
- Z3Prover::Bind
- Z3Prover::EnterConstraint
- Z3Prover::GetSMTLIB2
- Z3Prover::GetStats
- Z3Prover::SetTimeoutMs
- Z3Prover::SetRLimit
- Z3Prover::GetModel
- Z3Prover::CountSatisfyingValues
Scalar integer, unsigned integer, and boolean expression translation to Z3.
Support for arithmetic, comparisons, boolean operators, min, max, select, if_then_else, let, casts, truncated division/modulo, floor division/modulo, and selected bitwise/shift operations.
Deterministic solver control using Z3 rlimit, with random_seed fixed to 42.
Thread-local Z3 context sharing to reduce initialization overhead while keeping thread safety.
A disabled-mode stub implementation that returns conservative results when Z3 is not built.

Implementation notes

The on/off implementations live in a single src/arith/z3_prover.cc, selected by the TVM_USE_Z3 macro that cmake/modules/contrib/Z3.cmake defines when USE_Z3=ON. This keeps the optional backend in one file under src/arith and lets CMake control whether the real solver is compiled.
The Z3 fallback is exception-safe (z3::exception is caught and treated as "cannot prove") and gated behind kSymbolicBound, so the common kDefault path (including deeply recursive internal CanProve calls) never pays the prover cost.
TVM Div/Mod are truncated (round toward zero); they are translated with explicit truncating helpers instead of Z3's Euclidean operator/ / operator% to stay sound for negative dividends.
Shift handling relies on Z3's native bit-vector semantics and does not add hard assertions to the shared solver, which would otherwise permanently poison it.

TileLang References

The implementation is based on the Z3 analyzer integration used in TileLang's TVM fork, with the upstream port kept optional and scoped to TVM's arithmetic analyzer.

Acknowledgements

The original Z3 integration in TileLang's TVM fork was authored by @kurisu6912 and @SiriusNEO in tile-ai/tvm@e633295. This PR ports that work to upstream TVM as an optional backend. Thanks to @kurisu6912 and the TileLang authors for the original design and implementation.

Differences From TileLang

The core Z3Prover algorithm is kept close to the latest TileLang TVM implementation. The differences are in how it is integrated into upstream TVM:

Optional build, single file. We make Z3 optional, defaulting USE_Z3=OFF. Both the real implementation and the conservative stub live in one src/arith/z3_prover.cc, switched by the TVM_USE_Z3 macro, so Z3 stays out of the default build. In TileLang it is enabled by default.
Z3 discovery. We add cmake/modules/contrib/Z3.cmake, which, when USE_Z3=ON, first tries the PIC static libz3 layout provided by the PyPI z3-static package (headers + static library + CMake package files). A custom (e.g. shared system) Z3 can still be selected via Z3_DIR/CMAKE_PREFIX_PATH. TileLang resolves a shared Z3 from z3-solver directly via cmake/pypi-z3/FindZ3.cmake.
Packaging. This PR does not add a mandatory z3-static build-system requirement or enable USE_Z3 by default for all wheels yet, because the static Z3 wheels must be published first. Once that package is available, wheel builds can enable USE_Z3=ON and link Z3 statically without install RPATH handling, build-tree libz3 copying, or a runtime z3-solver Python dependency. TileLang carries shared-library wheel packaging for all of this.
Strength gating. The Z3 fallback is gated behind kSymbolicBound so it is only consulted from top-level predicate proving, avoiding prover cost on cheap recursive paths.
Soundness fixes. Truncated Div/Mod use truncating helpers rather than Z3's Euclidean operators, shift handling no longer poisons the shared solver, and CanProve is exception-safe.
Python surface. We expose only the debug/config APIs (get_smtlib2, set_z3_timeout_ms, set_z3_rlimit, get_z3_stats) to Python, keeping GetModel and CountSatisfyingValues as C++-only. TileLang exposes the same debug/config set, and its set_z3_rlimit is named set_z3_max_step.
No clone API. We do not port TileLang's Analyzer::Clone / Python Analyzer.clone(); Z3Prover::CopyFrom is retained so a future clone can reuse it.

gemini-code-assist

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.

gemini-code-assist · 2026-06-03T20:52:59Z

+      // 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);


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.

This makes sense. Could you have a correction?

gemini-code-assist · 2026-06-03T20:52:59Z

+    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));
+      }


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)); } }

gemini-code-assist · 2026-06-03T20:52:59Z

+  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;


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.

I think we can remove it as gemini suggested

gemini-code-assist · 2026-06-03T20:53:00Z

+    // 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);


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);

gemini-code-assist · 2026-06-03T20:53:00Z

+#include <unordered_set>
+#include <vector>
+
+#include "tvm/arith/analyzer.h"


Redundant Include

#include "tvm/arith/analyzer.h" is redundant because <tvm/arith/analyzer.h> is already included at the top of the file.

gemini-code-assist · 2026-06-03T20:53:00Z

+#include <tvm/tirx/expr.h>
+#include <tvm/tirx/op.h>
+
+#include "tvm/arith/analyzer.h"


Redundant Include

#include "tvm/arith/analyzer.h" is redundant because <tvm/arith/analyzer.h> is already included at the top of the file.

tlopex · 2026-06-03T23:34:40Z

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

tlopex

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

tlopex · 2026-06-04T00:00:09Z

+  z3::expr VisitExpr_(const SubNode* op) override {
+    return VisitArith(z3::operator-, op, op->a, op->b);
+  }
+  z3::expr VisitExpr_(const MulNode* op) override {


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

tlopex · 2026-06-04T00:13:29Z

+    }
+  }
+
+  z3::expr VisitExprDefault_(const Object* op) override {


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);)?

tlopex · 2026-06-04T00:19:33Z

    }
  }

+  if (z3_prover.CanProve(simplified)) {


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; }.

tlopex · 2026-06-04T00:29:06Z

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.
And EnterConstraint duplicates EnterWithScope without any real caller in this PR, and the new is_assume flag is effectively dead code since it is never set to true, so we should either deduplicate the entry logic and wire the flag with a real use/test or drop both changes until they’re actually needed.

SiriusNEO · 2026-06-06T05:49:32Z

BTW, the original Z3 integration is in this commit (Not a PR): tile-ai@e633295

SiriusNEO · 2026-06-06T05:45:46Z

For the tests, do we need to cover all tests in the original PR?: tile-ai/tilelang#1367

SiriusNEO · 2026-06-06T05:55:10Z

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?

tqchen · 2026-06-06T10:54:11Z

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

SiriusNEO

LGTM, cc @kurisu6912

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.

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.

Signed-off-by: Ubospica <ubospica@gmail.com>

- 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.

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.

Call the packaged z3_static.config entry point so CMake does not need to know the wheel's internal artifact layout.

Require the z3-static post release that exposes the config module used by the Z3 CMake lookup path.

Apply the formatting changes requested by taplo and clang-format after the Z3 package update.

Keep the Z3 analyzer branch aligned with the tvm-ffi revision already present on main.

Drop the relaxed access-path matching now that the PR no longer changes the tvm-ffi submodule revision.

Retry GitHub platform checks after conda-forge returned HTTP 429 during dependency setup.

Retry after the Windows GitHub job remained stuck in environment setup without logs.

tlopex

LGTM!

gemini-code-assist Bot reviewed Jun 3, 2026

View reviewed changes

Ubospica changed the title ~~[ARITH] Add optional Z3-backed proving to Analyzer~~ [ARITH] Add optional Z3-backend proving to Analyzer Jun 3, 2026

Ubospica changed the title ~~[ARITH] Add optional Z3-backend proving to Analyzer~~ [ARITH] Add optional Z3-backed proving to Analyzer Jun 3, 2026

tlopex reviewed Jun 3, 2026

View reviewed changes

tlopex requested changes Jun 3, 2026

View reviewed changes

tlopex reviewed Jun 4, 2026

View reviewed changes

SiriusNEO suggested changes Jun 6, 2026

View reviewed changes

Ubospica force-pushed the 2026-06-03-arith-z3-analyzer branch from 6211940 to 39f808e Compare June 10, 2026 20:48

Ubospica mentioned this pull request Jun 10, 2026

Add z3-static publishing and enable static Z3 wheel builds mlc-ai/package#103

Merged

Ubospica force-pushed the 2026-06-03-arith-z3-analyzer branch 2 times, most recently from b52e0c4 to 55d503c Compare June 14, 2026 23:51

SiriusNEO approved these changes Jun 15, 2026

View reviewed changes

Ubospica added 14 commits June 16, 2026 16:19

[CMAKE] Keep downstream helper standards unchanged

f69ef25

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.

[DOCS] Fix from source dependency list formatting

a93a250

update

254010f

Signed-off-by: Ubospica <ubospica@gmail.com>

Stop ignoring python/tvm/_version.py

5c9a546

[CMAKE] Enable static z3 package by default

6c114da

Use the z3-static CMake package path first and make the rebased Z3 analyzer build cleanly with LLVM and CUDA enabled.

[CMAKE] Clean up LLVM include propagation

84a2b4a

Keep the LLVM header interface target on explicit compile options so object targets inherit llvm-config include paths reliably.

[CMAKE] Auto-enable Z3 when available

048b5d0

Default USE_Z3 to AUTO so source builds use z3-static when installed while CI environments without the package still build the conservative stub.

[CMAKE] Use MSVC include dirs for LLVM headers

80fc0ef

Avoid passing GCC-style -isystem options to cl while keeping explicit LLVM include propagation for non-MSVC builds.

[CMAKE] Document MSVC LLVM include handling

f46a5bd

Clarify why the LLVM header target uses include directories instead of GCC-style -isystem options on MSVC.

[CMAKE] Clarify MSVC LLVM include handling

f211a32

Ubospica added 8 commits June 16, 2026 16:19

[CI] Retry GPU infrastructure check

010b079

[TEST] Match reflected structural equal access paths

1d995cc

[TEST] Lift rlimit for Z3 monotonicity proof

6495bd5

[TEST] Lift rlimit for Z3 index bound proof

6d25638

[ARITH] Pin Z3 context lifetime in prover

eccf6d1

[TEST] Tolerate reflected access path formatting

9c9f152

[CMAKE] Use z3-static config module for Z3 lookup

c94829f

Call the packaged z3_static.config entry point so CMake does not need to know the wheel's internal artifact layout.

[CMAKE] Require z3-static config package

a4df6ba

Require the z3-static post release that exposes the config module used by the Z3 CMake lookup path.

Ubospica force-pushed the 2026-06-03-arith-z3-analyzer branch from 1c5cedd to a4df6ba Compare June 16, 2026 20:21

Ubospica added 5 commits June 16, 2026 16:22

[LINT] Apply Z3 formatting fixes

5ddcbfc

Apply the formatting changes requested by taplo and clang-format after the Z3 package update.

[FFI] Restore tvm-ffi submodule pointer

7b89ce7

Keep the Z3 analyzer branch aligned with the tvm-ffi revision already present on main.

[TEST] Restore Relax structural path checks

f028cee

Drop the relaxed access-path matching now that the PR no longer changes the tvm-ffi submodule revision.

[CI] Retry infrastructure checks

c857a26

Retry GitHub platform checks after conda-forge returned HTTP 429 during dependency setup.

[CI] Retry stuck Windows setup

fde41c9

Retry after the Windows GitHub job remained stuck in environment setup without logs.

tlopex approved these changes Jun 18, 2026

View reviewed changes

tlopex merged commit e7b87fe into apache:main Jun 18, 2026
10 checks passed

Conversation

Ubospica commented Jun 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Features

Implementation notes

TileLang References

Acknowledgements

Differences From TileLang

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Jun 3, 2026

Choose a reason for hiding this comment

Soundness Issue with Shift Amount Assertions

Uh oh!

tlopex Jun 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jun 3, 2026

Choose a reason for hiding this comment

Constant Range Optimization is Never Triggered

Uh oh!

gemini-code-assist Bot Jun 3, 2026

Choose a reason for hiding this comment

Thread-Safety and Cross-Thread Copying Limitation

Uh oh!

tlopex Jun 4, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jun 3, 2026

Choose a reason for hiding this comment

Implicit Double-to-Unsigned Conversion

Uh oh!

gemini-code-assist Bot Jun 3, 2026

Choose a reason for hiding this comment

Redundant Include

Uh oh!

gemini-code-assist Bot Jun 3, 2026

Choose a reason for hiding this comment

Redundant Include

Uh oh!

tlopex Jun 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tlopex left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tlopex Jun 4, 2026

Choose a reason for hiding this comment

Uh oh!

tlopex Jun 4, 2026

Choose a reason for hiding this comment

Uh oh!

tlopex Jun 4, 2026

Choose a reason for hiding this comment

Uh oh!

tlopex commented Jun 4, 2026

Uh oh!

SiriusNEO commented Jun 6, 2026

Uh oh!

SiriusNEO Jun 6, 2026

Choose a reason for hiding this comment

Uh oh!

SiriusNEO Jun 6, 2026

Choose a reason for hiding this comment

Uh oh!

tqchen commented Jun 6, 2026

Uh oh!

SiriusNEO left a comment

Choose a reason for hiding this comment

Uh oh!

tlopex left a comment

Choose a reason for hiding this comment

Uh oh!

Ubospica commented Jun 3, 2026 •

edited

Loading

tlopex Jun 3, 2026 •

edited

Loading

tlopex Jun 3, 2026 •

edited

Loading

tlopex left a comment •

edited

Loading