Support nested software pipelining (#463)

* Support nested software pipelining * Update test_schedule_software_pipeline.py
tlc-pack · Sep 1, 2021 · 8d398d0 · 8d398d0
1 parent ca726d7
commit 8d398d0
Show file tree

Hide file tree

Showing 5 changed files with 417 additions and 479 deletions.
diff --git a/include/tvm/tir/transform.h b/include/tvm/tir/transform.h
@@ -539,6 +539,12 @@ TVM_DLL Pass UnifyThreadBinding();
  */
 TVM_DLL Pass LowerLogicalLayout();
 
+/*!
+ * \brief Transform annotated loops into pipelined one that parallelize producers and consumers.
+ * \return The IR transform pass.
+ */
+TVM_DLL Pass InjectSoftwarePipeline();
+
 /*!
  * \brief Lower logical intrinsics into physical intrinsics.
  */

diff --git a/src/tir/transforms/flatten_buffer.cc b/src/tir/transforms/flatten_buffer.cc
@@ -99,7 +99,7 @@ class BufferFlattener : public StmtExprMutator {
     for (const auto& annotation : op->annotations) {
       const String& ann_key = annotation.first;
       const ObjectRef& ann_value = annotation.second;
-      if (attr::IsPragmaKey(ann_key)) {
+      if (attr::IsPragmaKey(ann_key) || ann_key == attr::pipeline_scope ) {
         body = AttrStmt(op->loop_var, ann_key, Downcast<PrimExpr>(ann_value), std::move(body));
       }
     }