simplify MultiLevelTilingHexagon

masahi · Nov 29, 2022 · 44ad353 · 44ad353
1 parent a64f0d3
commit 44ad353
Show file tree

Hide file tree

Showing 3 changed files with 58 additions and 331 deletions.
diff --git a/src/meta_schedule/schedule_rule/multi_level_tiling.h b/src/meta_schedule/schedule_rule/multi_level_tiling.h
@@ -21,6 +21,7 @@
 
 #include <tvm/meta_schedule/schedule_rule.h>
 #include <tvm/tir/schedule/schedule.h>
+#include "../../tir/schedule/transform.h"
 
 #include <unordered_map>
 #include <utility>
@@ -135,6 +136,23 @@ std::vector<State> SubRule(std::vector<State> states, FLambda sub_rule) {
   return results;
 }
 
+/*!
+ * \brief Tile a subset of loops in the block according to the given tensor intrinsic, and annotate
+ * the tiled block for tensorization by postproc rewrite.
+ */
+inline Optional<tir::BlockRV> TileForIntrin(tir::Schedule sch, tir::BlockRV block,
+                                     const std::string& intrin_name) {
+  Optional<tir::LoopRV> tiled_loop_rv = TileWithTensorIntrin(sch, block, intrin_name);
+  if (!tiled_loop_rv) {
+    return NullOpt;
+  }
+  ICHECK(tiled_loop_rv.defined());
+  tir::BlockRV outer_block = sch->Blockize(tiled_loop_rv.value());
+  sch->Annotate(outer_block, tir::attr::meta_schedule_auto_tensorize, String(intrin_name));
+  return outer_block;
+}
+
+
 /*!
  * \brief The mega rule: multi-level tiling with data reuse
  */