[TIR] Add software pipelining

include/tvm/tir/stmt.h

@@ -1361,6 +1361,12 @@ constexpr const char* script_parsing_detect_access = "tir.script_parsing_detect_
  */
 constexpr const char* pragma_loop_partition_hint = "pragma_loop_partition_hint";
 
+/*! \brief Mark the stage of a statement in the software pipeline */
+constexpr const char* software_pipeline_stage = "software_pipeline_stage";
+
+/*! \brief Mark the order of a statement in the software pipeline */
+constexpr const char* software_pipeline_order = "software_pipeline_order";
+
 /*! \brief Mark the tiling structure of blocks that are applied by rule Multi-Level-Tiling */
 constexpr const char* meta_schedule_tiling_structure = "meta_schedule.tiling_structure";
 

include/tvm/tir/transform.h

@@ -492,6 +492,39 @@ TVM_DLL Pass ConvertForLoopsToSerial();
  */
 TVM_DLL Pass UnifiedStaticMemoryPlanner();
 
+/*!
+ * \brief Transform annotated loops into pipelined one that ovarlaps producers and consumers.
+ *
+ * This pass detects loops with the software pipeline annotations and rewrite them to pipelined
+ * ones. The behavior of such rewriting depending on two annotations on the loop,
+ * attr::software_pipeline_stage, and attr::software_pipeline_order, which defines the stage and the
+ * order, respectively, of the components of the software pipeline. The components of the software
+ * pipeline is the direct children (ignoring BlockRealize / Block / SeqStmt) of the annotated loop.
+ * The value of the both annotations should be array of integers, with its size the same as the
+ * number of the components.
+ *
+ * The result of the rewriting is a block that has three blocks as its direct children which
+ * represents the prologue, the body, and the epilogue of the software pipeline. In the prologue,
+ * only components whose stage is less than max_stage will be executed. In the epilogue, only
+ * components whose stage is greater than 0 will be executed. In the body, all the components will
+ * be executed. Such rewriting enables behavior like prefetching, the components are not necessarily
+ * executed in the original order. attr::software_pipeline_order defines the order of the each
+ * component. Components belong to different stages can be reordered.
+ *
+ * Buffer allocated inside the software pipeline may be resized to accommodate multiple versions
+ * of the original buffer. Block annotation attr::double_buffer_scope can be used to indicate that
+ * the block need to write in the double-buffering style.
+ *
+ * Annotations:
+ * attr::software_pipeline_stage: Array of non-negative integers, each element should be in range
+ *                                [0, max_stage], where max_stage is the maximum (inclusive) stage.
+ * attr::software_pipeline_order: Array of non-negative integers, should be a permutation of
+ *                                [0, 1, ..., num_components - 1].
+ *
+ * \return The IR transform pass.
+ */
+TVM_DLL Pass InjectSoftwarePipeline();
+
 }  // namespace transform
 }  // namespace tir
 }  // namespace tvm

python/tvm/tir/transform/transform.py

@@ -749,3 +749,14 @@ def ConvertForLoopsToSerial():
         The result pass
     """
     return _ffi_api.ConvertForLoopsToSerial()  # type: ignore
+
+
+def InjectSoftwarePipeline():
+    """Transform annotated loops into pipelined one that parallelize producers and consumers
+
+    Returns
+    -------
+    fpass : tvm.transform.Pass
+        The result pass
+    """
+    return _ffi_api.InjectSoftwarePipeline()  # type: ignore

src/driver/driver_api.cc

@@ -247,6 +247,7 @@ Array<tvm::transform::Pass> CreatePassList(bool disable_loop_partition) {
   pass_list.push_back(tir::transform::UnifyThreadBinding());
   pass_list.push_back(tir::transform::CompactBufferAllocation());
   pass_list.push_back(tir::transform::LowerMatchBuffer());
+  pass_list.push_back(tir::transform::InjectSoftwarePipeline());
   pass_list.push_back(tir::transform::FlattenBuffer());
   pass_list.push_back(tir::transform::BF16Legalize());
   pass_list.push_back(tir::transform::NarrowDataType(32));