[POC] Metal via tvm-ffi

CMakeLists.txt

@@ -484,6 +484,15 @@ include(cmake/modules/Git.cmake)
 include(cmake/modules/LibInfo.cmake)
 include(cmake/modules/contrib/Mrvl.cmake)
 
+tvm_option(USE_Z3 "Build with Z3 SMT solver support" OFF)
+
+if (USE_Z3)
+  find_package(Z3 REQUIRED)
+  list(APPEND COMPILER_SRCS src/target/z3/z3_prover_on.cc)
+else()
+  list(APPEND COMPILER_SRCS src/target/z3/z3_prover_off.cc)
+endif()
+
 set(LIBINFO_FILE ${CMAKE_CURRENT_LIST_DIR}/src/support/libinfo.cc)
 add_lib_info(${LIBINFO_FILE})
 list(REMOVE_ITEM COMPILER_SRCS ${LIBINFO_FILE})
@@ -541,6 +550,42 @@ else()
   target_link_libraries(tvm_runtime PUBLIC tvm_ffi_shared)
 endif()
 
+if(USE_Z3)
+  target_include_directories(tvm_objs PRIVATE ${Z3_CXX_INCLUDE_DIRS})
+  target_include_directories(tvm_runtime_objs PRIVATE ${Z3_CXX_INCLUDE_DIRS})
+  target_include_directories(tvm_libinfo_objs PRIVATE ${Z3_CXX_INCLUDE_DIRS})
+  target_link_libraries(tvm PRIVATE z3::libz3)
+
+  if (APPLE)
+    # `libz3.dylib` from z3-solver on pypi have a "wrong" name `libz3.dylib`,
+    # so it won't be searched in rpath. We patch it to `@rpath/libz3.dylib` here.
+    # `POST_BUILD` command needs to be in same cmake file where the target's created.
+    add_custom_command(TARGET tvm POST_BUILD
+      COMMAND install_name_tool -change "libz3.dylib" "@rpath/libz3.dylib" $<TARGET_FILE:tvm>
+      COMMENT "Patching libz3 reference to use @rpath"
+    )
+  else()
+    find_program(PATCHELF_EXECUTABLE patchelf)
+    if (PATCHELF_EXECUTABLE)
+      execute_process(
+        COMMAND ${PATCHELF_EXECUTABLE} --print-soname ${Z3_LIBRARY}
+        OUTPUT_VARIABLE Z3_SONAME
+        OUTPUT_STRIP_TRAILING_WHITESPACE
+        RESULT_VARIABLE Z3_SONAME_RESULT
+      )
+      if(NOT Z3_SONAME_RESULT EQUAL "0")
+        message(FATAL_ERROR "Failed to get Z3 soname using patchelf")
+      endif()
+      message("-- Z3 SONAME: ${Z3_SONAME}")
+      add_custom_command(TARGET tvm POST_BUILD
+        COMMAND ${PATCHELF_EXECUTABLE} --replace-needed ${Z3_SONAME} libz3.so $<TARGET_FILE:tvm>
+        COMMENT "Patching libz3 reference to use soname ${Z3_SONAME}"
+      )
+    else()
+      message("patchelf not found, skip.")
+    endif()
+  endif()
+endif()
 
 target_include_directories(tvm_runtime PUBLIC "$<INSTALL_INTERFACE:${CMAKE_INSTALL_INCLUDEDIR}>")
 set_property(TARGET tvm_runtime APPEND PROPERTY LINK_OPTIONS "${TVM_VISIBILITY_FLAG}")
@@ -748,7 +793,7 @@ endif()
 
 # Change relative paths in backtrace to absolute ones
 if(TVM_IS_DEBUG_BUILD)
-  set(FILE_PREFIX_MAP_FLAG "-ffile-prefix-map=..=${CMAKE_CURRENT_SOURCE_DIR}")
+  # set(FILE_PREFIX_MAP_FLAG "-ffile-prefix-map=..=${CMAKE_CURRENT_SOURCE_DIR}")
   target_compile_options(tvm PRIVATE "${FILE_PREFIX_MAP_FLAG}")
   CHECK_CXX_COMPILER_FLAG("${FILE_PREFIX_MAP_FLAG}" FILE_PREFIX_MAP_SUPPORTED)
   if(FILE_PREFIX_MAP_SUPPORTED)

include/tvm/arith/analyzer.h

@@ -33,6 +33,7 @@
 #include <memory>
 #include <unordered_map>
 #include <vector>
+#include "tvm/ffi/object.h"
 
 namespace tvm {
 /*! \brief namespace of arithmetic analysis. */
@@ -175,6 +176,8 @@ class ConstIntBoundAnalyzer {
   friend class ConstraintContext;
   explicit ConstIntBoundAnalyzer(Analyzer* parent);
   TVM_DLL ~ConstIntBoundAnalyzer();
+  // Deep-copy internal state from another instance (for Analyzer::Clone)
+  void CopyFrom(const ConstIntBoundAnalyzer& other);
   /*!
    * \brief Update the internal state to enter constraint.
    * \param constraint A constraint expression.
@@ -254,6 +257,8 @@ class ModularSetAnalyzer {
   friend class ConstraintContext;
   explicit ModularSetAnalyzer(Analyzer* parent);
   TVM_DLL ~ModularSetAnalyzer();
+  // Deep-copy internal state from another instance (for Analyzer::Clone)
+  void CopyFrom(const ModularSetAnalyzer& other);
   /*!
    * \brief Update the internal state to enter constraint.
    * \param constraint A constraint expression.
@@ -294,7 +299,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);
+  TVM_DLL std::function<void()> EnterConstraint(const PrimExpr& constraint, bool is_assume=false);
 
   /*! \brief Flags to enable more computationally-intensive simplifications
    *
@@ -407,6 +412,8 @@ class RewriteSimplifier {
   friend class CanonicalSimplifier;
   explicit RewriteSimplifier(Analyzer* parent);
   TVM_DLL ~RewriteSimplifier();
+  // Deep-copy internal state from another instance (for Analyzer::Clone)
+  void CopyFrom(const RewriteSimplifier& other);
   class Impl;
   /*! \brief Internal impl */
   Impl* impl_;
@@ -438,6 +445,8 @@ class CanonicalSimplifier {
   friend class ConstraintContext;
   explicit CanonicalSimplifier(Analyzer* parent);
   TVM_DLL ~CanonicalSimplifier();
+  // Deep-copy internal state from another instance (for Analyzer::Clone)
+  void CopyFrom(const CanonicalSimplifier& other);
   class Impl;
   /*! \brief Internal impl */
   Impl* impl_;
@@ -523,6 +532,8 @@ class TransitiveComparisonAnalyzer {
   friend class ConstraintContext;
   TransitiveComparisonAnalyzer();
   TVM_DLL ~TransitiveComparisonAnalyzer();
+  // Deep-copy internal state from another instance (for Analyzer::Clone)
+  void CopyFrom(const TransitiveComparisonAnalyzer& other);
   class Impl;
   /*! \brief Internal impl */
   std::unique_ptr<Impl> impl_;
@@ -553,8 +564,8 @@ class ConstraintContext {
    * \param analyzer The analyzer.
    * \param constraint The constraint to be applied.
    */
-  ConstraintContext(Analyzer* analyzer, PrimExpr constraint)
-      : analyzer_(analyzer), constraint_(constraint) {}
+  ConstraintContext(Analyzer* analyzer, PrimExpr constraint, bool is_assume=false)
+      : analyzer_(analyzer), constraint_(constraint), is_assume_(is_assume) {}
   // enter the scope.
   void EnterWithScope();
   // exit the scope.
@@ -565,6 +576,7 @@ class ConstraintContext {
   PrimExpr constraint_;
   /*! \brief function to be called in recovery */
   std::vector<std::function<void()>> recovery_functions_;
+  bool is_assume_;
 };
 
 /*!
@@ -616,11 +628,85 @@ class IntSetAnalyzer {
   friend class Analyzer;
   explicit IntSetAnalyzer(Analyzer* parent);
   TVM_DLL ~IntSetAnalyzer();
+  // Deep-copy internal state from another instance (for Analyzer::Clone)
+  void CopyFrom(const IntSetAnalyzer& other);
   class Impl;
   /*! \brief Internal impl */
   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 can we prove expr is always true.
+   *
+   * \param expr The expression.
+   * \return Whether we can prove it.
+   *
+   * \note Analyzer will call into sub-analyzers to get the result.
+   */
+  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, bool is_assume=false);
+
+  /*!
+   * \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 (like maxinum step or sth.)
+   */
+  void SetRLimit(unsigned rlimit);
+
+ private:
+  friend class Analyzer;
+  explicit Z3Prover(Analyzer* parent);
+  TVM_DLL ~Z3Prover();
+  void CopyFrom(const Z3Prover & other);
+  class Impl;
+  Impl* impl_;
+};
+
 /*!
  * \brief Analyzer that contains bunch of sub-analyzers.
  *
@@ -650,8 +736,15 @@ class TVM_DLL Analyzer {
   IntSetAnalyzer int_set;
   /*! \brief sub-analyzer transitive comparisons */
   TransitiveComparisonAnalyzer transitive_comparisons;
+  /*! \brief analyzer using z3 */
+  Z3Prover z3_prover;
   /*! \brief constructor */
   Analyzer();
+  /*!
+   * \brief Create a deep copy of this Analyzer, including all sub-analyzer states.
+   * \return A new Analyzer with copied internal state.
+   */
+  std::unique_ptr<Analyzer> Clone() const;
   /*!
    * \brief Mark the value as non-negative value globally in analyzer.
    *

include/tvm/ir/type.h

@@ -309,6 +309,5 @@ class TensorMapType : public Type {
   TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE_WITHOUT_DEFAULT_CONSTRUCTOR(TensorMapType, Type,
                                                                          TensorMapTypeNode);
 };
-
 }  // namespace tvm
 #endif  // TVM_IR_TYPE_H_

include/tvm/script/ir_builder/tir/frame.h

@@ -350,11 +350,12 @@ class LetFrameNode : public TIRFrameNode {
   /*! \brief The value we bind var to */
   PrimExpr value;
 
+
   static void RegisterReflection() {
     namespace refl = tvm::ffi::reflection;
     refl::ObjectDef<LetFrameNode>()
         .def_ro("var", &LetFrameNode::var)
-        .def_ro("value", &LetFrameNode::value);
+        .def_rw("value", &LetFrameNode::value);
   }
   TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tir.LetFrame", LetFrameNode, TIRFrameNode);
 

include/tvm/script/ir_builder/tir/ir.h

@@ -488,6 +488,7 @@ TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_SIZES(UInt, DataType::UInt);
 TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_SIZES(Int, DataType::Int);
 
 #define TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_LANES(FuncName, FDType, Size) \
+  TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(FuncName##x2, FDType(Size, 2));     \
   TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(FuncName##x4, FDType(Size, 4));     \
   TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(FuncName##x8, FDType(Size, 8));     \
   TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(FuncName##x16, FDType(Size, 16));   \
@@ -507,6 +508,7 @@ TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_SIZES_LANES(Int, DataType::Int);
 
 #define TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(DType, FDType) \
   TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(DType, FDType(1));                    \
+  TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(DType##x2, FDType(2));                \
   TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(DType##x4, FDType(4));                \
   TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(DType##x8, FDType(8));                \
   TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(DType##x16, FDType(16));              \

include/tvm/tir/expr.h

@@ -731,9 +731,21 @@ class CallNode : public PrimExprNode {
   /*! \brief The arguments. */
   ffi::Array<PrimExpr> args;
 
+  /*!
+   * \brief Additional annotations about the call.
+   *
+   *  These annotations can be used to pass additional metadata
+   *  to lowering passes. For tile operators, this can include
+   *  coalesced_width, disable_tma, eviction_policy, etc.
+   */
+  ffi::Map<ffi::String, ObjectRef> annotations;
+
   static void RegisterReflection() {
     namespace refl = tvm::ffi::reflection;
-    refl::ObjectDef<CallNode>().def_ro("op", &CallNode::op).def_ro("args", &CallNode::args);
+    refl::ObjectDef<CallNode>()
+        .def_ro("op", &CallNode::op)
+        .def_ro("args", &CallNode::args)
+        .def_ro("annotations", &CallNode::annotations);
   }
   TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tir.Call", CallNode, PrimExprNode);
 };
@@ -744,7 +756,9 @@ class CallNode : public PrimExprNode {
  */
 class Call : public PrimExpr {
  public:
-  TVM_DLL Call(DataType dtype, RelaxExpr op, ffi::Array<PrimExpr> args, Span span = Span());
+  TVM_DLL Call(DataType dtype, RelaxExpr op, ffi::Array<PrimExpr> args,
+               ffi::Map<ffi::String, ObjectRef> annotations = {},
+               Span span = Span());
   TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Call, PrimExpr, CallNode);
   TVM_DEFINE_OBJECT_REF_COW_METHOD(CallNode);
 };

include/tvm/tir/op.h

@@ -722,17 +722,17 @@ TVM_DLL PrimExpr fast_erf_float_expr(PrimExpr arg, int bits);
 
 // Intrinsic operators
 #define TVM_DECLARE_INTRIN_UNARY(OpName)                                \
-  inline PrimExpr OpName(PrimExpr x, Span span = Span()) {              \
-    static const Op& op = Op::Get("tir." #OpName);                      \
-    if (x.dtype().is_bfloat16()) {                                      \
-      DataType bf16_dtype = x.dtype();                                  \
-      DataType fp32_dtype(kDLFloat, 32, bf16_dtype.lanes());            \
-      PrimExpr x_fp32 = tir::Cast(fp32_dtype, {x}, span);               \
-      PrimExpr result_fp32 = tir::Call(fp32_dtype, op, {x_fp32}, span); \
-      return tir::Cast(bf16_dtype, {result_fp32}, span);                \
-    } else {                                                            \
-      return tir::Call(x.dtype(), op, {x}, span);                       \
-    }                                                                   \
+  inline PrimExpr OpName(PrimExpr x, Span span = Span()) {                  \
+    static const Op& op = Op::Get("tir." #OpName);                          \
+    if (x.dtype().is_bfloat16()) {                                          \
+      DataType bf16_dtype = x.dtype();                                      \
+      DataType fp32_dtype(kDLFloat, 32, bf16_dtype.lanes());                \
+      PrimExpr x_fp32 = tir::Cast(fp32_dtype, {x}, span);                   \
+      PrimExpr result_fp32 = tir::Call(fp32_dtype, op, {x_fp32}, {}, span); \
+      return tir::Cast(bf16_dtype, {result_fp32}, span);                    \
+    } else {                                                                \
+      return tir::Call(x.dtype(), op, {x}, {}, span);                       \
+    }                                                                       \
   }
 
 TVM_DECLARE_INTRIN_UNARY(exp);
@@ -764,7 +764,7 @@ TVM_DECLARE_INTRIN_UNARY(clz);
 #define TVM_DECLARE_INTRIN_BINARY(OpName)                              \
   inline PrimExpr OpName(PrimExpr x, PrimExpr y, Span span = Span()) { \
     static const Op& op = Op::Get("tir." #OpName);                     \
-    return tir::Call(x.dtype(), op, {x, y}, span);                     \
+    return tir::Call(x.dtype(), op, {x, y}, {}, span);                 \
   }
 
 TVM_DECLARE_INTRIN_BINARY(atan2);

include/tvm/tir/schedule/schedule.h

@@ -535,7 +535,7 @@ class ScheduleNode : public runtime::Object {
    * \return The reindex stage block.
    */
   virtual BlockRV ReIndex(const BlockRV& block_rv, int buffer_index,
-                          BufferIndexType buffer_index_type) = 0;
+                          BufferIndexType buffer_index_type, bool skip_simplify = false) = 0;
   /******** Schedule: Data movement ********/
   virtual BlockRV ReadAt(const LoopRV& loop_rv, const BlockRV& block_rv, int read_buffer_index,
                          const ffi::String& storage_scope) = 0;

include/tvm/tir/stmt.h

@@ -1319,6 +1319,8 @@ constexpr const char* explicit_read_region = "explicit_read_region";
  */
 constexpr const char* explicit_write_region = "explicit_write_region";
 
+constexpr const char* tilelang_assume = "tl.assume";
+
 /*! \brief ,ark a ForNode represent an irregular loop of non-structural control flow edges. */
 constexpr const char* irregular_loop_mark = "irregular_loop_mark";
 

include/tvm/topi/transform.h

@@ -1285,12 +1285,15 @@ inline Tensor take(const Tensor& a, ffi::Variant<Tensor, PrimExpr> indices, int
           for (size_t j = 0; j < static_cast<size_t>(axis); ++j) {
             real_indices.push_back(out_index[j]);
           }
-          auto idx = truncmod(truncmod(get_index(indices_position), axis_dim) + axis_dim, axis_dim);
+          PrimExpr idx = get_index(indices_position);
           real_indices.push_back(idx);
           for (size_t j = axis + indices_len; j < out_index.size(); ++j) {
             real_indices.push_back(out_index[j]);
           }
-          return a(real_indices);
+          PrimExpr in_bounds = idx >= 0 && idx < axis_dim;
+          return tvm::if_then_else(
+              in_bounds, a(real_indices),
+              tvm::tir::make_const(a->dtype, std::numeric_limits<float>::quiet_NaN()));
         },
         name, tag);
   }