[ARITH] Add optional Z3-backed proving to Analyzer

CMakeLists.txt

@@ -90,6 +90,7 @@ tvm_option(COMPILER_RT_PATH "Path to COMPILER-RT" "3rdparty/compiler-rt")
 tvm_option(USE_BYODT_POSIT "Build with BYODT software emulated posit custom datatype" OFF)
 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)
@@ -439,9 +440,9 @@ include(cmake/utils/CCache.cmake)
 
 include(CheckCXXCompilerFlag)
 set(CMAKE_CXX_STANDARD_REQUIRED ON)
-set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD 20)
 set(CMAKE_CUDA_STANDARD_REQUIRED ON)
-set(CMAKE_CUDA_STANDARD 17)
+set(CMAKE_CUDA_STANDARD 20)
 
 # Module rules
 include(cmake/modules/CUDA.cmake)
@@ -460,6 +461,7 @@ include(cmake/modules/contrib/CUTLASS.cmake)
 include(cmake/modules/contrib/Random.cmake)
 include(cmake/modules/contrib/Posit.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)
@@ -546,6 +548,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,94 @@
+# 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}" -c
+        "import os, z3_static as m; f = getattr(m, 'get_cmake_dir', None); print(f() if f else os.path.join(os.path.dirname(m.__file__), 'static', 'lib', 'cmake', 'z3'))"
+      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`, 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")

cmake/utils/FindLLVM.cmake

@@ -254,9 +254,9 @@ macro(find_llvm use_llvm)
     # compiler-appropriate form so the probe works under MSVC as well.
     if(NOT CMAKE_CXX_STANDARD)
       if(MSVC)
-        set(CMAKE_REQUIRED_FLAGS "/std:c++17")
+        set(CMAKE_REQUIRED_FLAGS "/std:c++20")
       else()
-        set(CMAKE_REQUIRED_FLAGS "-std=c++17")
+        set(CMAKE_REQUIRED_FLAGS "-std=c++20")
       endif()
     endif()
     check_cxx_source_compiles("

docs/install/from_source.rst

@@ -35,11 +35,15 @@ Apache TVM requires the following dependencies:
 - CMake (>= 3.24.0)
 - LLVM (recommended >= 15)
 - Git
-- A recent C++ compiler supporting C++ 17, at the minimum
-    - GCC 7.1
-    - Clang 5.0
-    - Apple Clang 9.3
-    - Visual Studio 2019 (v16.7)
+- A recent C++ compiler supporting C++ 20, at the minimum
+    - GCC 10
+    - Clang 10
+    - Apple Clang 14
+    - Visual Studio 2022
+
+  Optional dependencies that use newer C++20 standard library facilities, such
+  as ``std::format``, may require a newer standard library (for example GCC 13
+  or newer on Linux).
 - Python (>= 3.10)
 - (Optional) Conda (Strongly Recommended)
 

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

include/tvm/tirx/op.h

@@ -825,7 +825,8 @@ inline bool IsPointerType(const Type& type, const DataType& element_type) {
  * \param span The location of this operation in the source.
  */
 template <typename ValueType,
-          typename = typename std::enable_if<std::is_pod<ValueType>::value>::type>
+          typename = typename std::enable_if<std::is_standard_layout<ValueType>::value &&
+                                             std::is_trivial<ValueType>::value>::type>
 inline PrimExpr make_const(DataType t, ValueType value, Span span = Span());
 /*!
  * \brief Make a const zero expr.

pyproject.toml

@@ -16,7 +16,8 @@
 # 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"]
 build-backend = "scikit_build_core.build"
 
 [project]
@@ -141,6 +142,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]