[CI] Bump CI at the Ubuntu 24.04 images and re-enable USE_Z3

CMakeLists.txt

@@ -89,7 +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" 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)

ci/jenkins/docker-images.ini

@@ -17,7 +17,7 @@
 
 # This data file is read during when Jenkins runs job to determine docker images.
 [jenkins]
-ci_tag: 20260619-214849-4174cdf5
+ci_tag: 20260705-142349-cfb98e93
 ci_arm: tlcpack/ci-arm:%(ci_tag)s
 ci_cpu: tlcpack/ci-cpu:%(ci_tag)s
 ci_gpu: tlcpack/ci-gpu:%(ci_tag)s

pyproject.toml

@@ -16,7 +16,7 @@
 # under the License.
 
 [build-system]
-# z3-static ships the PIC static libz3 + headers for explicit USE_Z3=ON builds.
+# z3-static ships the PIC static libz3 + headers consumed by USE_Z3=ON.
 requires = [
   "scikit-build-core>=0.11",
   "setuptools-scm>=8",
@@ -146,8 +146,8 @@ logging.level = "INFO"
 [tool.scikit-build.cmake.define]
 TVM_BUILD_PYTHON_MODULE = "ON"
 USE_CUDA = "OFF"
-# Keep Z3 disabled by default until CI's C++ toolchain can link z3-static reliably.
-USE_Z3 = "OFF"
+# Statically link Z3 from the z3-static build dependency by default.
+USE_Z3 = "AUTO"
 BUILD_TESTING = "OFF"
 
 [tool.setuptools_scm]

src/arith/z3_prover.cc

@@ -267,7 +267,7 @@ class Z3Prover::Impl : ExprFunctor<z3::expr(const Expr&)> {
     scope_stack_.back().push_back(Scope{Scope::BindValue, var, value});
     // we add the binding whenever the value is pure,
     // because non-pure parts are handling by creating free variables in VisitExpr
-    memo_.emplace(var, ConvertInt(value));
+    memo_.emplace(var.as_or_throw<PrimExpr>(), ConvertInt(value));
   }
 
   /// @brief Bind a variable to a range
@@ -279,7 +279,7 @@ class Z3Prover::Impl : ExprFunctor<z3::expr(const Expr&)> {
     //    if the var is overrided later, we can just update the memo, and the old placeholder will
     //    be ignored
     auto var_expr = Create(var.get());
-    memo_.emplace(var, var_expr);
+    memo_.emplace(var.as_or_throw<PrimExpr>(), var_expr);
 
     // 2. Add constraint on the placeholder
     //    when min_expr >= max_expr, the range is empty, which is under undefined behavior
@@ -299,8 +299,9 @@ class Z3Prover::Impl : ExprFunctor<z3::expr(const Expr&)> {
         solver.add(var_expr < ctx->int_val(max_value));
       }
     } else {
+      PrimExpr prim_var = var.as_or_throw<PrimExpr>();
       solver.add(ConvertBool(range->extent <= 0 ||
-                             (range->min <= var && var < range->min + range->extent)));
+                             (range->min <= prim_var && prim_var < range->min + range->extent)));
     }
   }
 
@@ -435,7 +436,7 @@ class Z3Prover::Impl : ExprFunctor<z3::expr(const Expr&)> {
     solver.push();
 
     // Convert the TVM variable to Z3 expression
-    z3::expr z3_var = VisitInt(var);
+    z3::expr z3_var = VisitInt(var.as_or_throw<PrimExpr>());
 
     int64_t count = 0;
     std::vector<int64_t> found_values;
@@ -596,7 +597,7 @@ class Z3Prover::Impl : ExprFunctor<z3::expr(const Expr&)> {
 
   z3::expr VisitExpr_(const LetNode* op) override {
     if (IsZ3SupportedExpr(op->var.get())) {
-      memo_.emplace(op->var, VisitInt(op->value));
+      memo_.emplace(op->var.as_or_throw<PrimExpr>(), VisitInt(op->value));
     }
     return VisitExpr(op->body);
   }
@@ -704,7 +705,9 @@ class Z3Prover::Impl : ExprFunctor<z3::expr(const Expr&)> {
     } else if (op->op.same_as(tirx::builtin::if_then_else()) && op->args.size() == 3 &&
                IsZ3SupportedExpr(op->args[1].get()) && IsZ3SupportedExpr(op->args[2].get())) {
       // tir.if_then_else(cond, a, b) is a select-like ternary.
-      return z3::ite(VisitBool(op->args[0]), VisitInt(op->args[1]), VisitInt(op->args[2]));
+      return z3::ite(VisitBool(op->args[0].as_or_throw<PrimExpr>()),
+                     VisitInt(op->args[1].as_or_throw<PrimExpr>()),
+                     VisitInt(op->args[2].as_or_throw<PrimExpr>()));
     } else {
       // For other call nodes, create a free variable
       return Create(op);
@@ -719,9 +722,9 @@ class Z3Prover::Impl : ExprFunctor<z3::expr(const Expr&)> {
       TVM_FFI_UNREACHABLE();
     }
 
-    const PrimExpr& a = op->args[0];
-    const PrimExpr& b = op->args[1];
-    unsigned bit_width = std::max(op->args[0].ty().bits(), op->args[1].ty().bits());
+    PrimExpr a = op->args[0].as_or_throw<PrimExpr>();
+    PrimExpr b = op->args[1].as_or_throw<PrimExpr>();
+    unsigned bit_width = std::max(a.ty().bits(), b.ty().bits());
 
     if (IsZ3SupportedExpr(a.get()) && IsZ3SupportedExpr(b.get())) {
       return z3::bv2int(
@@ -738,7 +741,7 @@ class Z3Prover::Impl : ExprFunctor<z3::expr(const Expr&)> {
       TVM_FFI_UNREACHABLE();
     }
 
-    const PrimExpr& a = op->args[0];
+    PrimExpr a = op->args[0].as_or_throw<PrimExpr>();
 
     if (IsZ3SupportedExpr(a.get())) {
       // Cast integer to bit-vector, apply bitwise not, then cast back.
@@ -758,8 +761,8 @@ class Z3Prover::Impl : ExprFunctor<z3::expr(const Expr&)> {
       TVM_FFI_UNREACHABLE();
     }
 
-    const PrimExpr& a = op->args[0];
-    const PrimExpr& b = op->args[1];
+    PrimExpr a = op->args[0].as_or_throw<PrimExpr>();
+    PrimExpr b = op->args[1].as_or_throw<PrimExpr>();
 
     // Shift operations require integer types for both operands
     if (IsZ3SupportedExpr(a.get()) && IsZ3SupportedExpr(b.get())) {

tests/python/arith/test_arith_z3.py

@@ -96,9 +96,10 @@ def worker():
 # ---------------------------------------------------------------------------
 # Examples the native analyzer cannot prove but Z3 can.
 #
-# Each case asserts both that the native analyzers (kDefault, Z3 gated off)
+# Most cases assert both that the native analyzers (kDefault, Z3 gated off)
 # fail and that Z3 (kSymbolicBound) succeeds. This demonstrates the added value
-# of the Z3 backend and that it is correctly gated behind kSymbolicBound.
+# of the Z3 backend and that it is correctly gated behind kSymbolicBound. Cases
+# the native analyzer learns to prove remain here as Z3 translation regressions.
 # ---------------------------------------------------------------------------
 
 
@@ -155,7 +156,6 @@ def test_z3_nested_floor_division_collapse():
         tirx.all(a >= 0, a < 128),
         a // 128 == (a // 64 * 32 + a % 32 // 16 * 8) // 64,
     )
-    assert not analyzer.can_prove(expr)
     assert analyzer.can_prove(expr, SB)