update

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

Author: Ubospica

3rdparty/cnpy

@@ -0,0 +1 @@
+Subproject commit 4e8810b1a8637695171ed346ce68f6984e585ef4

3rdparty/dmlc-core

@@ -0,0 +1 @@
+Subproject commit 3031e4a61a98f49f07a42cfdec6242340fb2fd8c

3rdparty/rang

@@ -0,0 +1 @@
+Subproject commit cabe04d6d6b05356fa8f9741704924788f0dd762

3rdparty/zlib

@@ -0,0 +1 @@
+Subproject commit ef24c4c7502169f016dcd2a26923dbaf3216748c

cmake/modules/contrib/Z3.cmake

@@ -15,8 +15,10 @@
 # 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(NOT USE_Z3)
-  list(APPEND COMPILER_SRCS src/target/z3/z3_prover_off.cc)
   return()
 endif()
 
@@ -73,4 +75,6 @@ else()
   message(FATAL_ERROR "USE_Z3 is ON, but Z3 was not found. Install Z3 or PyPI z3-solver.")
 endif()
 
-list(APPEND COMPILER_SRCS src/target/z3/z3_prover_on.cc)
+# 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

@@ -296,7 +296,7 @@ class RewriteSimplifier {
    *
    * \return an exit function that must be called to cleanup the constraint can be nullptr.
    */
-  TVM_DLL std::function<void()> EnterConstraint(const PrimExpr& constraint, bool is_assume = false);
+  TVM_DLL std::function<void()> EnterConstraint(const PrimExpr& constraint);
 
   /*! \brief Flags to enable more computationally-intensive simplifications
    *
@@ -555,8 +555,8 @@ class ConstraintContext {
    * \param analyzer The analyzer.
    * \param constraint The constraint to be applied.
    */
-  ConstraintContext(Analyzer* analyzer, PrimExpr constraint, bool is_assume = false)
-      : analyzer_(analyzer), constraint_(constraint), is_assume_(is_assume) {}
+  ConstraintContext(Analyzer* analyzer, PrimExpr constraint)
+      : analyzer_(analyzer), constraint_(constraint) {}
   // enter the scope.
   void EnterWithScope();
   // exit the scope.
@@ -567,7 +567,6 @@ class ConstraintContext {
   PrimExpr constraint_;
   /*! \brief functions to be called in recovery */
   std::vector<std::function<void()>> recovery_functions_;
-  bool is_assume_;
 };
 
 /*!
@@ -644,6 +643,13 @@ class Z3Prover {
    */
   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.
    *
@@ -656,10 +662,9 @@ class Z3Prover {
    * \brief Update the internal state to enter constraint.
    *
    * \param constraint A constraint expression.
-   * \param is_assume Whether the constraint comes from an assumption.
    * \return an exit function that must be called to cleanup the constraint can be nullptr.
    */
-  std::function<void()> EnterConstraint(const PrimExpr& constraint, bool is_assume = false);
+  std::function<void()> EnterConstraint(const PrimExpr& constraint);
 
   /*!
    * \brief Get the SMTLIB2 representation of the current context.
@@ -886,8 +891,6 @@ class TVM_DLL Analyzer {
    * \note Analyzer will call into sub-analyzers to get the result.
    */
   PrimExpr Simplify(const PrimExpr& expr, int steps = 2);
-
-  std::function<void()> EnterConstraint(const PrimExpr& constraint, bool is_assume = false);
 };
 
 }  // namespace arith

python/tvm/_version.py

@@ -0,0 +1,24 @@
+# file generated by vcs-versioning
+# don't change, don't track in version control
+from __future__ import annotations
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+version: str
+__version__: str
+__version_tuple__: tuple[int | str, ...]
+version_tuple: tuple[int | str, ...]
+commit_id: str | None
+__commit_id__: str | None
+
+__version__ = version = '0.25.dev100'
+__version_tuple__ = version_tuple = (0, 25, 'dev100')
+
+__commit_id__ = commit_id = 'g35152f312'

python/tvm/arith/analyzer.py

@@ -123,51 +123,97 @@ def __init__(self):
         self._enter_constraint_context = _mod("enter_constraint_context")
         self._can_prove_equal = _mod("can_prove_equal")
         self._can_prove = _mod("can_prove")
+        self._is_z3_enabled = _mod("is_z3_enabled")
         self._get_smtlib2 = _mod("get_smtlib2")
         self._set_z3_timeout_ms = _mod("set_z3_timeout_ms")
         self._set_z3_rlimit = _mod("set_z3_rlimit")
         self._get_z3_stats = _mod("get_z3_stats")
         self._get_enabled_extensions = _mod("get_enabled_extensions")
         self._set_enabled_extensions = _mod("set_enabled_extensions")
 
-    def get_smtlib2(self, expr: tirx.PrimExpr = None) -> str:
+    @property
+    def is_z3_enabled(self) -> bool:
+        """Whether this build includes the Z3 backend (``USE_Z3=ON``).
+
+        The Z3-specific methods (:py:meth:`get_smtlib2`, :py:meth:`get_z3_stats`,
+        :py:meth:`set_z3_timeout_ms`, :py:meth:`set_z3_rlimit`) only work when
+        this is ``True``.
+        """
+        return bool(self._is_z3_enabled())
+
+    def _check_z3_enabled(self) -> None:
+        if not self.is_z3_enabled:
+            raise RuntimeError(
+                "The Z3 backend is not available in this build. "
+                "Rebuild TVM with USE_Z3=ON to use Z3-specific Analyzer APIs."
+            )
+
+    def get_smtlib2(self, expr: tirx.PrimExpr | None = None) -> str:
         """Get the current Z3 problem in SMT-LIB2 format.
 
+        Raises
+        ------
+        RuntimeError
+            If TVM was built without Z3 (``USE_Z3=OFF``), since there is no
+            solver state to export. Use :py:attr:`is_z3_enabled` to check first.
+
         Parameters
         ----------
         expr : Optional[PrimExpr]
             The expression to prove. If provided, its negation is added to the problem.
         """
+        self._check_z3_enabled()
         return self._get_smtlib2(expr)
 
     def set_z3_timeout_ms(self, timeout_ms: int) -> None:
         """Set Z3 timeout in milliseconds.
 
+        Raises
+        ------
+        RuntimeError
+            If TVM was built without Z3 (``USE_Z3=OFF``).
+
         Parameters
         ----------
         timeout_ms : int
             The timeout in milliseconds.
         """
+        self._check_z3_enabled()
         self._set_z3_timeout_ms(timeout_ms)
 
     def set_z3_rlimit(self, rlimit: int) -> None:
         """Set Z3 resource limit.
 
+        The resource limit gives deterministic solver budgeting (unlike a wall
+        clock timeout). A value of ``0`` disables the limit.
+
+        Raises
+        ------
+        RuntimeError
+            If TVM was built without Z3 (``USE_Z3=OFF``).
+
         Parameters
         ----------
         rlimit : int
             The resource limit.
         """
+        self._check_z3_enabled()
         self._set_z3_rlimit(rlimit)
 
     def get_z3_stats(self) -> str:
         """Get Z3 solver statistics.
 
+        Raises
+        ------
+        RuntimeError
+            If TVM was built without Z3 (``USE_Z3=OFF``).
+
         Returns
         -------
         stats : str
             The Z3 statistics.
         """
+        self._check_z3_enabled()
         return self._get_z3_stats()
 
     def const_int_bound(self, expr: tirx.PrimExpr) -> ConstIntBound:
@@ -301,7 +347,9 @@ def can_prove(
             The expression.
 
         strength: ProofStrength
-            The proof strength
+            The proof strength. When TVM is built with Z3 (``USE_Z3=ON``), the
+            optional Z3 fallback is only consulted at ``SYMBOLIC_BOUND`` or
+            higher, after the native analyzers fail to prove the predicate.
 
         Returns
         -------

src/arith/analyzer.cc

@@ -131,11 +131,10 @@ void ConstraintContext::EnterWithScope() {
   // entering the scope.
   recovery_functions_.push_back(analyzer_->const_int_bound.EnterConstraint(constraint_));
   recovery_functions_.push_back(analyzer_->modular_set.EnterConstraint(constraint_));
-  recovery_functions_.push_back(
-      analyzer_->rewrite_simplify.EnterConstraint(constraint_, is_assume_));
+  recovery_functions_.push_back(analyzer_->rewrite_simplify.EnterConstraint(constraint_));
   recovery_functions_.push_back(analyzer_->int_set.EnterConstraint(constraint_));
   recovery_functions_.push_back(analyzer_->transitive_comparisons.EnterConstraint(constraint_));
-  recovery_functions_.push_back(analyzer_->z3_prover.EnterConstraint(constraint_, is_assume_));
+  recovery_functions_.push_back(analyzer_->z3_prover.EnterConstraint(constraint_));
 }
 
 void ConstraintContext::ExitWithScope() {
@@ -235,30 +234,15 @@ bool Analyzer::CanProve(const PrimExpr& expr, ProofStrength strength) {
     }
   }
 
-  if (z3_prover.CanProve(simplified)) {
+  // Z3 is an expensive best-effort fallback. Gate it behind the higher
+  // kSymbolicBound strength so the common kDefault path (including deeply
+  // recursive internal CanProve calls) never pays the prover cost.
+  if (strength >= ProofStrength::kSymbolicBound && z3_prover.CanProve(simplified)) {
     return true;
   }
   return false;
 }
 
-std::function<void()> Analyzer::EnterConstraint(const PrimExpr& constraint, bool is_assume) {
-  std::vector<std::function<void()>> recovery_functions;
-  recovery_functions.push_back(this->const_int_bound.EnterConstraint(constraint));
-  recovery_functions.push_back(this->modular_set.EnterConstraint(constraint));
-  recovery_functions.push_back(this->rewrite_simplify.EnterConstraint(constraint, is_assume));
-  recovery_functions.push_back(this->int_set.EnterConstraint(constraint));
-  recovery_functions.push_back(this->transitive_comparisons.EnterConstraint(constraint));
-  recovery_functions.push_back(this->z3_prover.EnterConstraint(constraint, is_assume));
-  return [recovery_functions]() {
-    for (auto it = recovery_functions.rbegin(); it != recovery_functions.rend(); ++it) {
-      auto& func = *it;
-      if (func) {
-        func();
-      }
-    }
-  };
-}
-
 PrimExpr Analyzer::Simplify(const PrimExpr& expr, int steps) {
   PrimExpr res = expr;
 
@@ -371,6 +355,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
           self->rewrite_simplify.SetEnabledExtensions(
               static_cast<RewriteSimplifier::Extension>(flags));
         });
+      } else if (name == "is_z3_enabled") {
+        return ffi::Function(
+            [self](ffi::PackedArgs args, ffi::Any* ret) { *ret = self->z3_prover.IsEnabled(); });
       } else if (name == "get_smtlib2") {
         return ffi::Function([self](ffi::PackedArgs args, ffi::Any* ret) {
           auto expr = args[0].cast<ffi::Optional<PrimExpr>>();

src/arith/rewrite_simplify.cc

@@ -526,14 +526,13 @@ PrimExpr RewriteSimplifier::Impl::VisitExpr_(const AddNode* op) {
   return ret;
 }
 
-std::function<void()> RewriteSimplifier::Impl::EnterConstraint(const PrimExpr& constraint,
-                                                               bool is_assume) {
+std::function<void()> RewriteSimplifier::Impl::EnterConstraint(const PrimExpr& constraint) {
   size_t old_literal_size = literal_constraints_.size();
   // we will compare the already simplified result with the constraint,
   // so simplify the constraint as well
   PrimExpr new_constraint = operator()(constraint);
   for (const PrimExpr& subconstraint : ExtractConstraints(new_constraint, false)) {
-    if (is_assume || SideEffect(subconstraint) <= CallEffectKind::kPure) {
+    if (SideEffect(subconstraint) <= CallEffectKind::kPure) {
       literal_constraints_.push_back(subconstraint);
       PrimExpr negation;
       if (subconstraint.dtype().is_bool()) {
@@ -2441,9 +2440,8 @@ void RewriteSimplifier::Update(const Var& var, const PrimExpr& info, bool allow_
   impl_->Update(var, info, allow_override);
 }
 
-std::function<void()> RewriteSimplifier::EnterConstraint(const PrimExpr& constraint,
-                                                         bool is_assume) {
-  return impl_->EnterConstraint(constraint, is_assume);
+std::function<void()> RewriteSimplifier::EnterConstraint(const PrimExpr& constraint) {
+  return impl_->EnterConstraint(constraint);
 }
 
 void RewriteSimplifier::SetEnabledExtensions(Extension flags) {