[Dynamic] M2 for Pad-Einsum (S0) (#181)

mlc-ai · Jun 17, 2023 · 6ca9f26 · 6ca9f26
1 parent e473781
commit 6ca9f26
Show file tree

Hide file tree

Showing 7 changed files with 500 additions and 374 deletions.
diff --git a/include/tvm/tir/buffer.h b/include/tvm/tir/buffer.h
@@ -25,6 +25,7 @@
 #define TVM_TIR_BUFFER_H_
 
 #include <tvm/ir/expr.h>
+#include <tvm/node/script_printer.h>
 #include <tvm/runtime/container/array.h>
 #include <tvm/runtime/container/string.h>
 #include <tvm/tir/var.h>
@@ -162,6 +163,7 @@ class BufferNode : public Object {
   static constexpr const bool _type_has_method_sequal_reduce = true;
   static constexpr const bool _type_has_method_shash_reduce = true;
   TVM_DECLARE_FINAL_OBJECT_INFO(BufferNode, Object);
+  TVM_OBJECT_ENABLE_SCRIPT_PRINTER();
 };
 
 /*!

diff --git a/include/tvm/tir/op.h b/include/tvm/tir/op.h
@@ -493,6 +493,13 @@ TVM_DLL PrimExpr bitwise_neg(PrimExpr a, Span span = Span());
  */
 TVM_DLL PrimExpr if_then_else(PrimExpr cond, PrimExpr true_value, PrimExpr false_value,
                               Span span = Span());
+/*!
+ * \brief Protected write. This is only used on BufferStore's immediate RHS to indicate that
+ * out-of-bound access will not be performed.
+ * \param expr The expression to be protected.
+ * \return The result expression.
+ */
+TVM_DLL PrimExpr protected_write(PrimExpr expr);
 /*!
  * \brief Mark condition as likely.
  * \param cond The condition

diff --git a/src/tir/op/op.cc b/src/tir/op/op.cc
@@ -469,6 +469,13 @@ PrimExpr if_then_else(PrimExpr cond, PrimExpr true_value, PrimExpr false_value,
                    {cond, true_value, false_value}, span);
 }
 
+TVM_TIR_REGISTER_OP("protected_write");
+
+PrimExpr protected_write(PrimExpr expr) {
+  static const Op& op = Op::Get("tir.protected_write");
+  return tir::Call(expr.dtype(), op, {expr}, Span());
+}
+
 // likely
 PrimExpr likely(PrimExpr cond, Span span) {
   if (is_const_int(cond)) return cond;

diff --git a/src/tir/schedule/analysis.h b/src/tir/schedule/analysis.h
@@ -707,24 +707,6 @@ Array<arith::IntSet> AnalyzeRegionLowerBound(const BufferRegion& region, const P
                                              const StmtSRef& dom_high_exclusive,
                                              arith::Analyzer* analyzer);
 
-/*!
- * \brief Check if buffer indices are all Vars and extr
- * \param buffer_access The BufferLoad or BufferStore
- * \return The indices if the indices are all Vars, otherwise NullOpt
- */
-template <typename T>
-Optional<Array<Var>> CheckTrivialBufferIndices(const T& buffer_access) {
-  Array<Var> indices;
-  for (const PrimExpr& index : buffer_access->indices) {
-    const VarNode* var = index.as<VarNode>();
-    if (var == nullptr) {
-      return NullOpt;
-    }
-    indices.push_back(GetRef<Var>(var));
-  }
-  return indices;
-}
-
 /*!
  * \brief Simplify non-trivial expressions
  * \param expr The expression to be simplified

diff --git a/src/tir/schedule/primitive.h b/src/tir/schedule/primitive.h
@@ -632,8 +632,7 @@ TVM_DLL StmtSRef DecomposePadding(ScheduleState self, const StmtSRef& block_sref
  * \param block_sref The block sref that matches the Einsum pattern.
  * \param padding The padding for each block iter.
  */
-TVM_DLL void PadEinsum(ScheduleState self, const StmtSRef& block_sref,
-                       const Array<Integer>& padding);
+TVM_DLL void PadEinsum(ScheduleState self, StmtSRef block_sref, Array<Integer> padding);
 
 /******** Schedule: Buffer transformation ********/
 /*!