[GitHub] [tvm] masahi commented on a diff in pull request #12171: [TIR] Asynchronous stage in software pipeline

[GitBox <gi...@apache.org>]

masahi commented on code in PR #12171:
URL: https://github.com/apache/tvm/pull/12171#discussion_r930522785


##########
src/tir/transforms/inject_software_pipeline.cc:
##########
@@ -530,18 +620,269 @@ class PipelineRewriter : public StmtExprMutator {
       Block new_block = Downcast<Block>(PipelineBodyRewriter(buffer_data_to_buffer_, buffer_remap_,
                                                              pipeline_loop_, max_stage_ != 1,
                                                              fragment_info_)(block));
-      Map<Var, PrimExpr> subst_map;
-      if (is_unit_loop) {
-        subst_map.Set(pipeline_loop_->loop_var, skewed_loop_var);
-      } else {
-        // normalize loop range
-        PrimExpr delta = start - pipeline_loop_->min;
-        subst_map.Set(pipeline_loop_->loop_var, skewed_loop_var + delta);
+
+      PrimExpr delta = start - pipeline_loop_->min;
+      // This variable corresponds to
+      // - "producer_head" if this stage is an async producer
+      // - "consumer_head" if this stage reads from asynchronously written buffers.
+      PrimExpr normalized_access_index = is_unit_loop ? skewed_loop_var : skewed_loop_var + delta;
+
+      // Adjust the block predicate and the body according to the final loop bound
+      //  [pipeline_loop_->min, extent).
+      if (!is_unit_loop) {
         Var loop_iter = Downcast<Var>(new_loop_var);
-        inbound = Substitute(inbound, Map<Var, PrimExpr>{{loop_iter, loop_iter + delta}});
+        inbound = Substitute(inbound, {{loop_iter, loop_iter + delta}});
+      }
+
+      new_block = Downcast<Block>(
+          Substitute(new_block, {{pipeline_loop_->loop_var, normalized_access_index}}));
+
+      if (pipeline_info_[block].async) {

Review Comment:
   ok moved the bulk of logic into two functions. Now `EmitImpl` itself is kept short.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org