[ARITH] Add optional Z3-backed proving to Analyzer (#19667)

## Summary

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

The integration is guarded by `USE_Z3`, which defaults to `AUTO`. 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::CanProve` runs 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
is `kSymbolicBound`. Z3 is linked statically from the PyPI `z3-static`
package, so `libtvm` does not need a runtime `libz3` dependency.

## Features

- Z3 build support through `USE_Z3`, defaulting to `AUTO`.
- A new `arith::Z3Prover` sub-analyzer owned by `arith::Analyzer`.
- SMT-LIB2 export for debugging and external solver reproduction.
- Python debug/config APIs: `Analyzer.get_smtlib2`,
`Analyzer.set_z3_timeout_ms`, `Analyzer.set_z3_rlimit`, and
`Analyzer.get_z3_stats`.
- C++ APIs for proving, binding, constraints, stats, model inspection,
and satisfying-value counting.
- 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 real and stub implementations live in `src/arith/z3_prover.cc`,
selected by the `TVM_USE_Z3` macro from
`cmake/modules/contrib/Z3.cmake`.
- `cmake/modules/contrib/Z3.cmake` first resolves the PIC static `libz3`
layout provided by `z3-static` using its `z3_static.get_cmake_dir()`
helper, then falls back to a custom `Z3_DIR` or `CMAKE_PREFIX_PATH`
installation.
- `USE_Z3=ON` requires Z3 to be found, while `USE_Z3=AUTO` allows source
builds and CI jobs without Z3 artifacts to continue with the stub.
- The Z3 fallback is exception-safe and gated behind `kSymbolicBound`,
so the common `kDefault` path does not pay solver cost.
- TVM `Div` and `Mod` are translated with truncating helpers rather than
Z3's Euclidean operators 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.

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

- [tile-ai/tilelang#1367](tile-ai/tilelang#1367)
- [tile-ai/tilelang#1458](tile-ai/tilelang#1458)
- [tile-ai/tilelang#2216](tile-ai/tilelang#2216)
- [tile-ai#22](tile-ai#22)
- [tile-ai#24](tile-ai#24)
- [Original TileLang TVM
commit](tile-ai@e633295)

---------

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

Author: Ubospica

CMakeLists.txt

@@ -89,6 +89,7 @@ tvm_option(COMPILER_RT_PATH "Path to COMPILER-RT" "3rdparty/compiler-rt")
 # Contrib library options
 tvm_option(USE_BLAS "The blas library to be linked" none)
 tvm_option(USE_AMX "Enable Intel AMX" OFF)
+tvm_option(USE_Z3 "Build with Z3 SMT solver support" AUTO)
 tvm_option(USE_MKL "MKL root path when use MKL blas" OFF)
 tvm_option(USE_DNNL "Enable DNNL codegen" OFF)
 tvm_option(USE_CUDNN "Build with cuDNN" OFF)
@@ -459,6 +460,7 @@ include(cmake/modules/contrib/AMX.cmake)
 include(cmake/modules/contrib/CUTLASS.cmake)
 include(cmake/modules/contrib/Random.cmake)
 include(cmake/modules/contrib/Sort.cmake)
+include(cmake/modules/contrib/Z3.cmake)
 include(cmake/modules/contrib/CoreML.cmake)
 include(cmake/modules/contrib/TensorRT.cmake)
 include(cmake/modules/contrib/NNAPI.cmake)
@@ -545,6 +547,9 @@ add_library(tvm_objs OBJECT ${COMPILER_SRCS})
 add_library(tvm_runtime_objs OBJECT ${RUNTIME_SRCS})
 target_link_libraries(tvm_objs PUBLIC tvm_ffi_header)
 target_link_libraries(tvm_runtime_objs PUBLIC tvm_ffi_header)
+if(TARGET tvm_llvm_header)
+  target_link_libraries(tvm_objs PUBLIC tvm_llvm_header)
+endif()
 
 include(GNUInstallDirs)
 

cmake/modules/LLVM.cmake

@@ -34,6 +34,19 @@ if(NOT ${USE_LLVM} MATCHES ${IS_FALSE_PATTERN})
   endif()
   include_directories(SYSTEM ${LLVM_INCLUDE_DIRS})
   add_definitions(${LLVM_DEFINITIONS})
+  add_library(tvm_llvm_header INTERFACE)
+  if(MSVC)
+    # MSVC treats GCC-style -isystem operands as source files.
+    target_include_directories(tvm_llvm_header SYSTEM INTERFACE ${LLVM_INCLUDE_DIRS})
+    target_compile_options(tvm_llvm_header INTERFACE ${LLVM_DEFINITIONS})
+  else()
+    set(TVM_LLVM_INCLUDE_FLAGS "")
+    foreach(__llvm_include_dir IN LISTS LLVM_INCLUDE_DIRS)
+      string(STRIP "${__llvm_include_dir}" __llvm_include_dir)
+      list(APPEND TVM_LLVM_INCLUDE_FLAGS "-isystem" "${__llvm_include_dir}")
+    endforeach()
+    target_compile_options(tvm_llvm_header INTERFACE ${TVM_LLVM_INCLUDE_FLAGS} ${LLVM_DEFINITIONS})
+  endif()
   message(STATUS "Build with LLVM " ${LLVM_PACKAGE_VERSION})
   message(STATUS "Set TVM_LLVM_VERSION=" ${TVM_LLVM_VERSION})
   # Set flags that are only needed for LLVM target

cmake/modules/contrib/Z3.cmake

@@ -0,0 +1,93 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+# src/arith/z3_prover.cc is always part of COMPILER_SRCS (picked up by the
+# src/arith/*.cc glob). It compiles a conservative stub by default and switches
+# to the real Z3 implementation only when the TVM_USE_Z3 macro is defined below.
+if(${USE_Z3} MATCHES ${IS_FALSE_PATTERN})
+  return()
+endif()
+
+set(TVM_Z3_REQUIRED TRUE)
+if("${USE_Z3}" MATCHES "^[Aa][Uu][Tt][Oo]$")
+  set(TVM_Z3_REQUIRED FALSE)
+endif()
+
+# Default lookup: the PIC static Z3 library shipped by the PyPI `z3-static`
+# package (headers + libz3.a + Z3 CMake package files). Linking it statically
+# keeps libtvm free of a runtime libz3 dependency. Users can override the
+# lookup by setting Z3_DIR/CMAKE_PREFIX_PATH to any Z3 installation (e.g. a
+# shared system Z3).
+if(NOT Z3_DIR)
+  find_package(Python3 COMPONENTS Interpreter QUIET)
+  if(Python3_EXECUTABLE)
+    execute_process(
+      COMMAND
+        "${Python3_EXECUTABLE}" -m z3_static.config --cmake-dir
+      OUTPUT_VARIABLE Z3_STATIC_CMAKE_DIR
+      OUTPUT_STRIP_TRAILING_WHITESPACE
+      ERROR_QUIET
+      RESULT_VARIABLE Z3_STATIC_RESULT
+    )
+    if(Z3_STATIC_RESULT EQUAL 0 AND EXISTS "${Z3_STATIC_CMAKE_DIR}")
+      set(Z3_DIR "${Z3_STATIC_CMAKE_DIR}")
+    endif()
+  endif()
+endif()
+
+find_package(Z3 CONFIG QUIET)
+if(NOT Z3_FOUND AND NOT TARGET z3::libz3 AND NOT TARGET Z3::libz3)
+  find_package(Z3 QUIET)
+endif()
+
+if(TARGET z3::libz3 OR TARGET Z3::libz3)
+  if(TARGET z3::libz3)
+    set(Z3_TARGET z3::libz3)
+  else()
+    set(Z3_TARGET Z3::libz3)
+  endif()
+  get_target_property(Z3_TARGET_INCLUDE_DIRS ${Z3_TARGET} INTERFACE_INCLUDE_DIRECTORIES)
+  if(Z3_TARGET_INCLUDE_DIRS)
+    include_directories(SYSTEM ${Z3_TARGET_INCLUDE_DIRS})
+  endif()
+  list(APPEND TVM_LINKER_LIBS ${Z3_TARGET})
+elseif(Z3_FOUND OR (Z3_INCLUDE_DIR AND Z3_LIBRARY))
+  if(NOT Z3_INCLUDE_DIR AND Z3_CXX_INCLUDE_DIRS)
+    set(Z3_INCLUDE_DIR ${Z3_CXX_INCLUDE_DIRS})
+  endif()
+  if(NOT Z3_LIBRARY AND Z3_LIBRARIES)
+    set(Z3_LIBRARY ${Z3_LIBRARIES})
+  endif()
+  if(NOT Z3_INCLUDE_DIR OR NOT Z3_LIBRARY)
+    message(FATAL_ERROR "USE_Z3 is ON, but Z3 include directory or library was not found.")
+  endif()
+  include_directories(SYSTEM ${Z3_INCLUDE_DIR})
+  list(APPEND TVM_LINKER_LIBS ${Z3_LIBRARY})
+else()
+  if(TVM_Z3_REQUIRED)
+    message(FATAL_ERROR
+      "USE_Z3 is ON, but Z3 was not found. Install the static Z3 development "
+      "package with `pip install 'z3-static>=4.16.0.post1'`, or point "
+      "Z3_DIR/CMAKE_PREFIX_PATH at a Z3 installation.")
+  endif()
+  message(STATUS "Build without Z3 SMT solver support")
+  return()
+endif()
+
+# Enable the real Z3 implementation inside the single src/arith/z3_prover.cc file.
+add_compile_definitions(TVM_USE_Z3)
+message(STATUS "Build with Z3 SMT solver support")

include/tvm/arith/analyzer.h

@@ -27,6 +27,7 @@
 #include <tvm/arith/int_set.h>
 #include <tvm/ffi/cast.h>
 #include <tvm/ffi/reflection/registry.h>
+#include <tvm/ffi/string.h>
 #include <tvm/ir/expr.h>
 #include <tvm/ir/with_context.h>
 
@@ -588,6 +589,110 @@ class IntSetAnalyzer {
   Impl* impl_;
 };
 
+class Z3Prover {
+ public:
+  /*!
+   * \brief Update binding of var to a new expression.
+   *
+   * \param var The variable of interest.
+   * \param new_range The range of allowed values for this var.
+   * \param allow_override whether we allow override of existing information.
+   */
+  TVM_DLL void Bind(const Var& var, const Range& new_range, bool allow_override = false);
+
+  /*!
+   * \brief Update binding of var to a new expression.
+   *
+   * \param var The variable of interest.
+   * \param expr The bound expression.
+   * \param allow_override whether we allow override of existing information.
+   */
+  TVM_DLL void Bind(const Var& var, const PrimExpr& expr, bool allow_override = false);
+
+  /*!
+   * \brief Whether the Z3 backend is compiled into this build (USE_Z3=ON).
+   *
+   * \return true if the real Z3 prover is available, false for the stub.
+   */
+  TVM_DLL bool IsEnabled() const;
+
+  /*!
+   * \brief Whether can we prove expr is always true.
+   *
+   * \param expr The expression.
+   * \return Whether we can prove it.
+   */
+  TVM_DLL bool CanProve(const PrimExpr& expr);
+
+  /*!
+   * \brief Update the internal state to enter constraint.
+   *
+   * \param constraint A constraint expression.
+   * \return an exit function that must be called to cleanup the constraint can be nullptr.
+   */
+  std::function<void()> EnterConstraint(const PrimExpr& constraint);
+
+  /*!
+   * \brief Get the SMTLIB2 representation of the current context.
+   *
+   * \param expr The optional expression to check.
+   * \return The SMTLIB2 string.
+   */
+  ffi::String GetSMTLIB2(const ffi::Optional<PrimExpr> expr);
+
+  /*!
+   * \brief Get statistics about Z3 prover.
+   *
+   * \return The statistics string.
+   */
+  ffi::String GetStats();
+
+  /*!
+   * \brief Set timeout in milliseconds for Z3 prover.
+   *
+   * \param timeout_ms The timeout in milliseconds.
+   */
+  void SetTimeoutMs(unsigned timeout_ms);
+
+  /*!
+   * \brief Set resource limitation for Z3 prover.
+   *
+   * \param rlimit the resource limitation.
+   */
+  void SetRLimit(unsigned rlimit);
+
+  /*!
+   * \brief Get the Z3 model for the given expression if satisfiable.
+   *
+   * \param expr The expression to get the model for.
+   * \return The model as a string.
+   */
+  ffi::String GetModel(const PrimExpr& expr);
+
+  /*!
+   * \brief Count the number of integer values that satisfy the current constraints.
+   *
+   * This method uses Z3's model enumeration to count how many distinct values of
+   * the given variable satisfy all current constraints.
+   *
+   * \param var The variable to count satisfying values for.
+   * \param max_count Maximum number of solutions to enumerate.
+   * \param min_consecutive Minimum consecutive count requirement.
+   * \return The number of distinct values that satisfy the constraints, or a negative error code.
+   */
+  TVM_DLL int64_t CountSatisfyingValues(const Var& var, int64_t max_count = 2048,
+                                        int64_t min_consecutive = 1);
+
+ private:
+  friend class AnalyzerObj;
+  friend class Analyzer;
+  explicit Z3Prover(AnalyzerObj* parent);
+  TVM_DLL ~Z3Prover();
+  void CopyFrom(const Z3Prover& other);
+  class Impl;
+  Impl* impl_;
+};
+
 /*!
  * \brief Analyzer that contains bunch of sub-analyzers.
  *
@@ -612,6 +717,8 @@ class TVM_DLL AnalyzerObj : public ffi::Object {
   IntSetAnalyzer int_set;
   /*! \brief sub-analyzer transitive comparisons */
   TransitiveComparisonAnalyzer transitive_comparisons;
+  /*! \brief sub-analyzer using Z3 */
+  Z3Prover z3_prover;
   /*! \brief constructor */
   AnalyzerObj();
   /*!
@@ -810,7 +917,16 @@ class ConstraintContext {
    * \param constraint The constraint to be applied.
    */
   ConstraintContext(const Analyzer& analyzer, PrimExpr constraint)
-      : analyzer_(analyzer), constraint_(constraint) {}
+      : ConstraintContext(analyzer, std::move(constraint), false) {}
+  /*!
+   * \brief Construct a constraint context.
+   * \param analyzer The analyzer whose context is updated. The context
+   *        keeps a reference to the analyzer while the scope is active.
+   * \param constraint The constraint to be applied.
+   * \param is_assume Whether the constraint comes from an assumption.
+   */
+  ConstraintContext(const Analyzer& analyzer, PrimExpr constraint, bool is_assume)
+      : analyzer_(analyzer), constraint_(std::move(constraint)), is_assume_(is_assume) {}
   /*!
    * \brief Construct a constraint context from a borrowed analyzer object.
    * \param analyzer The borrowed analyzer object.
@@ -819,7 +935,15 @@ class ConstraintContext {
    * This overload is for internal callers that already operate on AnalyzerObj*.
    */
   ConstraintContext(AnalyzerObj* analyzer, PrimExpr constraint)
-      : ConstraintContext(ffi::GetRef<Analyzer>(analyzer), std::move(constraint)) {}
+      : ConstraintContext(ffi::GetRef<Analyzer>(analyzer), std::move(constraint), false) {}
+  /*!
+   * \brief Construct a constraint context from a borrowed analyzer object.
+   * \param analyzer The borrowed analyzer object.
+   * \param constraint The constraint to be applied.
+   * \param is_assume Whether the constraint comes from an assumption.
+   */
+  ConstraintContext(AnalyzerObj* analyzer, PrimExpr constraint, bool is_assume)
+      : ConstraintContext(ffi::GetRef<Analyzer>(analyzer), std::move(constraint), is_assume) {}
   // enter the scope.
   void EnterWithScope();
   // exit the scope.
@@ -830,6 +954,8 @@ class ConstraintContext {
   PrimExpr constraint_;
   /*! \brief functions to be called in recovery */
   std::vector<std::function<void()>> recovery_functions_;
+  /*! \brief Whether the constraint comes from an assumption. */
+  bool is_assume_;
 };
 
 }  // namespace arith

pyproject.toml

@@ -16,7 +16,12 @@
 # under the License.
 
 [build-system]
-requires = ["scikit-build-core>=0.11", "setuptools-scm>=8"]
+# z3-static ships the PIC static libz3 + headers consumed by USE_Z3=ON.
+requires = [
+  "scikit-build-core>=0.11",
+  "setuptools-scm>=8",
+  "z3-static>=4.16.0.post1",
+]
 build-backend = "scikit_build_core.build"
 
 [project]
@@ -141,6 +146,8 @@ logging.level = "INFO"
 [tool.scikit-build.cmake.define]
 TVM_BUILD_PYTHON_MODULE = "ON"
 USE_CUDA = "OFF"
+# Statically link Z3 from the z3-static build dependency by default.
+USE_Z3 = "ON"
 BUILD_TESTING = "OFF"
 
 [tool.setuptools_scm]