You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by GitBox <gi...@apache.org> on 2022/08/16 07:54:37 UTC

[GitHub] [tvm] yincs-intellif opened a new pull request, #12450: [TIR][Schedule] enhance compute_at primitive to choose proper position

yincs-intellif opened a new pull request, #12450:
URL: https://github.com/apache/tvm/pull/12450

   Current TIR "compute_at" primitive will compute at it's closest consumers. When a block has multiple producers, whoever 
   compute at later who is behind. But for some special hardware, we usually hope keep the a certain order whatever it's compute at early or late.
   eg: block A and block B are producers of block C. block A compute at block C first and block B compute at C late. we hope the result is block B->block A->block C under some loop var.
   
   cc @Hzfengsy @wrongtest-intellif 


-- 
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


[GitHub] [tvm] yincs-intellif commented on pull request #12450: [TIR][Schedule] enhance compute_at primitive to choose proper position

Posted by GitBox <gi...@apache.org>.
yincs-intellif commented on PR #12450:
URL: https://github.com/apache/tvm/pull/12450#issuecomment-1225046736

   > Hey sorry for the late reply! Happy to do another round of review tomorrow!
   
   @junrushao  Hi, I don't receive any feedback information. Did I miss something?


-- 
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


[GitHub] [tvm] Hzfengsy merged pull request #12450: [TIR][Schedule] enhance compute_at and reverse_compute_at primitive to choose possible position

Posted by GitBox <gi...@apache.org>.
Hzfengsy merged PR #12450:
URL: https://github.com/apache/tvm/pull/12450


-- 
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


[GitHub] [tvm] Hzfengsy commented on pull request #12450: [TIR][Schedule] enhance compute_at and reverse_compute_at primitive to choose possible position

Posted by GitBox <gi...@apache.org>.
Hzfengsy commented on PR #12450:
URL: https://github.com/apache/tvm/pull/12450#issuecomment-1228391601

   Thanks @yincs-intellif for this PR and continuous update.


-- 
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


[GitHub] [tvm] Hzfengsy commented on a diff in pull request #12450: [TIR][Schedule] enhance compute_at primitive to choose proper position

Posted by GitBox <gi...@apache.org>.
Hzfengsy commented on code in PR #12450:
URL: https://github.com/apache/tvm/pull/12450#discussion_r946457015


##########
tests/python/unittest/test_tir_schedule_compute_at.py:
##########
@@ -1353,5 +1353,154 @@ def _create_prim_func():
     verify_trace_roundtrip(sch=sch, mod=mod)
 
 
+def test_compute_at_to_early_stage():

Review Comment:
   Can we make the test "unit" enough? i.e. use a minimal case for better readability. 



-- 
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


[GitHub] [tvm] junrushao commented on a diff in pull request #12450: [TIR][Schedule] enhance compute_at and reverse_compute_at primitive to choose possible position

Posted by GitBox <gi...@apache.org>.
junrushao commented on code in PR #12450:
URL: https://github.com/apache/tvm/pull/12450#discussion_r955645217


##########
src/tir/schedule/primitive/compute_at.cc:
##########
@@ -585,10 +603,11 @@ void ComputeAtOrReverseComputeAtImpl(ScheduleState self, const StmtSRef& block_s
   std::unordered_map<const BlockNode*, const BlockRealizeNode*> block2realize;
   block2realize.reserve(self->block_info.size());
   int insert_position = FindInsertionPoint<!is_compute_at, is_compute_at>(
-      /*self=*/self,
+      /*self=*/self, /*scope=*/scope,
       /*subtrees=*/AsArray(loop->body),
       /*producer_srefs=*/producer_srefs,
-      /*consumer_srefs=*/consumer_srefs, /*block2realize=*/&block2realize);
+      /*consumer_srefs=*/consumer_srefs, /*block2realize=*/&block2realize,
+      /*index*/ index);

Review Comment:
   ```suggestion
         /*index=*/index);
   ```



-- 
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


[GitHub] [tvm] junrushao commented on a diff in pull request #12450: [TIR][Schedule] enhance compute_at primitive to choose proper position

Posted by GitBox <gi...@apache.org>.
junrushao commented on code in PR #12450:
URL: https://github.com/apache/tvm/pull/12450#discussion_r953853588


##########
src/tir/schedule/primitive/compute_at.cc:
##########
@@ -160,7 +174,57 @@ int FindInsertionPoint(
   // The valid indices are: (last_producer_position, first_consumer_position]
   ICHECK(split.last_producer_position < split.first_consumer_position);
   // Step 4. Return the last valid insertion point
-  return split.first_consumer_position;
+  int insert_position = split.first_consumer_position;
+  if (index == -1) {
+    return insert_position;
+  } else if (index == -2) {
+    if (require_all_producers_visited) {
+      if (split.last_producer_position >= 0) {
+        insert_position = split.last_producer_position + 1;
+      } else {
+        insert_position = split.first_consumer_position;
+      }
+    } else if (require_all_consumers_visited) {
+      class Finder : public StmtVisitor {
+       public:
+        void VisitStmt_(const BlockRealizeNode* realize) final {
+          const BlockNode* block = realize->block.get();
+          if (producer_blocks_.count(block)) {
+            ++this->n_producers_visited_;
+          }
+        }
+
+        std::unordered_set<const StmtNode*> producer_blocks_;
+        int n_producers_visited_ = 0;
+      };
+      // adjust the inserted position by compute at order
+      for (int i = split.first_consumer_position; i - 1 > split.last_producer_position; --i) {
+        auto blk = GetBlock(subtrees[i]);
+        if (!blk.defined()) break;

Review Comment:
   Quick question: do we have the assumption that each subtree has only one block inside?



-- 
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


[GitHub] [tvm] Hzfengsy commented on a diff in pull request #12450: [TIR][Schedule] enhance compute_at primitive to choose proper position

Posted by GitBox <gi...@apache.org>.
Hzfengsy commented on code in PR #12450:
URL: https://github.com/apache/tvm/pull/12450#discussion_r953882199


##########
src/tir/schedule/primitive/compute_at.cc:
##########
@@ -160,7 +174,57 @@ int FindInsertionPoint(
   // The valid indices are: (last_producer_position, first_consumer_position]
   ICHECK(split.last_producer_position < split.first_consumer_position);
   // Step 4. Return the last valid insertion point
-  return split.first_consumer_position;
+  int insert_position = split.first_consumer_position;
+  if (index == -1) {
+    return insert_position;
+  } else if (index == -2) {
+    if (require_all_producers_visited) {
+      if (split.last_producer_position >= 0) {
+        insert_position = split.last_producer_position + 1;
+      } else {
+        insert_position = split.first_consumer_position;
+      }

Review Comment:
   This line does not make sense to me. `split.last_producer_position` should always be `>=-1`, so `split.last_producer_position + 1`, which is `>=0` is always a valid position for `index == -2`
   
   lmk if I make some mistakes.



##########
src/tir/schedule/primitive/compute_at.cc:
##########
@@ -160,7 +174,57 @@ int FindInsertionPoint(
   // The valid indices are: (last_producer_position, first_consumer_position]
   ICHECK(split.last_producer_position < split.first_consumer_position);
   // Step 4. Return the last valid insertion point
-  return split.first_consumer_position;
+  int insert_position = split.first_consumer_position;
+  if (index == -1) {
+    return insert_position;
+  } else if (index == -2) {

Review Comment:
   The logic for `index == -2` is so confused. Based on the branch `index > 0`, all possible location are always `(last_producer_position, first_consumer_position]`, so can we just return `last_producer_position + 1` in this case?



##########
src/tir/schedule/primitive/compute_at.cc:
##########
@@ -160,7 +174,57 @@ int FindInsertionPoint(
   // The valid indices are: (last_producer_position, first_consumer_position]
   ICHECK(split.last_producer_position < split.first_consumer_position);
   // Step 4. Return the last valid insertion point

Review Comment:
   Please update the comment here.



-- 
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


[GitHub] [tvm] junrushao commented on a diff in pull request #12450: [TIR][Schedule] enhance compute_at and reverse_compute_at primitive to choose possible position

Posted by GitBox <gi...@apache.org>.
junrushao commented on code in PR #12450:
URL: https://github.com/apache/tvm/pull/12450#discussion_r955645029


##########
src/tir/schedule/primitive/compute_at.cc:
##########
@@ -159,8 +162,22 @@ int FindInsertionPoint(
   // Step 3. Check if there is at least one index of the position can be inserted into
   // The valid indices are: (last_producer_position, first_consumer_position]
   ICHECK(split.last_producer_position < split.first_consumer_position);
-  // Step 4. Return the last valid insertion point
-  return split.first_consumer_position;
+  // Step 4. Return the possible insertion point according to index
+  int insert_position;
+  if (index == -1) {
+    insert_position = split.first_consumer_position;
+  } else if (index == -2) {
+    insert_position = split.last_producer_position + 1;
+  } else if (index >= 0 && index >= split.last_producer_position + 1 &&
+             index <= split.first_consumer_position) {
+    insert_position = index;
+  } else {
+    LOG(FATAL) << "Valid index:(-1, -2, [" << split.last_producer_position + 1 << ", "
+               << split.first_consumer_position << "]), "
+               << "current index=" << index;
+    throw;
+  }
+  return insert_position;

Review Comment:
   Thanks for updating the logic here! It's much clearer now!



-- 
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


[GitHub] [tvm] yincs-intellif commented on pull request #12450: [TIR][Schedule] enhance compute_at primitive to choose proper position

Posted by GitBox <gi...@apache.org>.
yincs-intellif commented on PR #12450:
URL: https://github.com/apache/tvm/pull/12450#issuecomment-1221748923

   @junrushao Hi, I have changed as you mentioned, could you please help review again?
   
   Thanks in advance!


-- 
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


[GitHub] [tvm] junrushao commented on a diff in pull request #12450: [TIR][Schedule] enhance compute_at primitive to choose proper position

Posted by GitBox <gi...@apache.org>.
junrushao commented on code in PR #12450:
URL: https://github.com/apache/tvm/pull/12450#discussion_r953851809


##########
include/tvm/tir/schedule/schedule.h:
##########
@@ -430,9 +430,11 @@ class ScheduleNode : public runtime::Object {
    * \param block_rv The block to be moved
    * \param loop_rv The loop where the block to be moved under
    * \param preserve_unit_loops Whether to keep the trivial loops whose extents are 1
+   * \param index The block index of the loop body subtree blocks
+   * -1, -2 means keep closed to or away from it's consumer

Review Comment:
   I polished the doc a little bit. let me know if makes sense
   
   ```suggestion
      * \param index The block index of the loop body subtree blocks.
      * A nonnegative number means the index of a possible insertion point
      * -1 means inserted into the last possible insertion point
      * -2 means inserted into the first possible insertion point
   ```



-- 
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


[GitHub] [tvm] yincs-intellif commented on a diff in pull request #12450: [TIR][Schedule] enhance compute_at primitive to choose proper position

Posted by GitBox <gi...@apache.org>.
yincs-intellif commented on code in PR #12450:
URL: https://github.com/apache/tvm/pull/12450#discussion_r954427886


##########
src/tir/schedule/primitive/compute_at.cc:
##########
@@ -160,7 +174,57 @@ int FindInsertionPoint(
   // The valid indices are: (last_producer_position, first_consumer_position]
   ICHECK(split.last_producer_position < split.first_consumer_position);
   // Step 4. Return the last valid insertion point
-  return split.first_consumer_position;
+  int insert_position = split.first_consumer_position;
+  if (index == -1) {
+    return insert_position;
+  } else if (index == -2) {
+    if (require_all_producers_visited) {
+      if (split.last_producer_position >= 0) {
+        insert_position = split.last_producer_position + 1;
+      } else {
+        insert_position = split.first_consumer_position;
+      }

Review Comment:
   when  split.last_producer_position==-1, we still hope the insert_position is split.first_consumer_position(the last block) instead of 0.



-- 
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


[GitHub] [tvm] junrushao commented on pull request #12450: [TIR][Schedule] enhance compute_at primitive to choose proper position

Posted by GitBox <gi...@apache.org>.
junrushao commented on PR #12450:
URL: https://github.com/apache/tvm/pull/12450#issuecomment-1221824915

   Hey sorry for the late reply! Happy to do another round of review tomorrow!


-- 
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


[GitHub] [tvm] Hzfengsy commented on pull request #12450: [TIR][Schedule] enhance compute_at primitive to choose proper position

Posted by GitBox <gi...@apache.org>.
Hzfengsy commented on PR #12450:
URL: https://github.com/apache/tvm/pull/12450#issuecomment-1216283589

   Do you have ideas about the naming of `to_early_stage`? cc @junrushao1994 @spectrometerHBH 


-- 
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


[GitHub] [tvm] junrushao commented on a diff in pull request #12450: [TIR][Schedule] enhance compute_at and reverse_compute_at primitive to choose possible position

Posted by GitBox <gi...@apache.org>.
junrushao commented on code in PR #12450:
URL: https://github.com/apache/tvm/pull/12450#discussion_r955644837


##########
src/tir/schedule/primitive.h:
##########
@@ -299,10 +299,12 @@ TVM_DLL StmtSRef ReIndex(ScheduleState self, const StmtSRef& block_sref, int buf
  * \param self The schedule state
  * \param block_sref The block to be moved
  * \param loop_sref The loop where the block to be moved to
- * \param preserve_unit_loops Whether to keep the trivial loops whose extents are 1
+ * \param index The block index of the loop body subtree blocks
+ * -1 means inserted into the last possible insertion point
+ * -2 means inserted into the first possible insertion point

Review Comment:
   Let's update the doc accordingly



-- 
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


[GitHub] [tvm] yincs-intellif commented on pull request #12450: [TIR][Schedule] enhance compute_at and reverse_compute_at primitive to choose possible position

Posted by GitBox <gi...@apache.org>.
yincs-intellif commented on PR #12450:
URL: https://github.com/apache/tvm/pull/12450#issuecomment-1228364837

   Hi @junrushao @Hzfengsy  The final modification has been updated. 
   Thank you for the high quality review!


-- 
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


[GitHub] [tvm] junrushao commented on a diff in pull request #12450: [TIR][Schedule] enhance compute_at and reverse_compute_at primitive to choose possible position

Posted by GitBox <gi...@apache.org>.
junrushao commented on code in PR #12450:
URL: https://github.com/apache/tvm/pull/12450#discussion_r955644743


##########
python/tvm/tir/schedule/schedule.py:
##########
@@ -1303,6 +1304,11 @@ def compute_at(
         preserve_unit_loops: bool
             Whether to keep the trivial loops whose extents are 1
 
+        index: int
+            The block index of the loop body subtree blocks
+            -1 means inserted into the last possible insertion point
+            -2 means inserted into the first possible insertion point

Review Comment:
   ```suggestion
               The block index of the loop body subtree blocks:
               - `index = -1` means inserted into the last possible insertion point;
               - `index = -2` means inserted into the first possible insertion point;
               - Otherwise, `index` is a nonnegative number that indicates the insertion point
   ```



-- 
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


[GitHub] [tvm] yincs-intellif commented on a diff in pull request #12450: [TIR][Schedule] enhance compute_at primitive to choose proper position

Posted by GitBox <gi...@apache.org>.
yincs-intellif commented on code in PR #12450:
URL: https://github.com/apache/tvm/pull/12450#discussion_r946495534


##########
tests/python/unittest/test_tir_schedule_compute_at.py:
##########
@@ -1353,5 +1353,154 @@ def _create_prim_func():
     verify_trace_roundtrip(sch=sch, mod=mod)
 
 
+def test_compute_at_to_early_stage():

Review Comment:
   Thanks @Hzfengsy  I have update test case by a minimal case.



-- 
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


[GitHub] [tvm] yincs-intellif commented on pull request #12450: [TIR][Schedule] enhance compute_at primitive to choose proper position

Posted by GitBox <gi...@apache.org>.
yincs-intellif commented on PR #12450:
URL: https://github.com/apache/tvm/pull/12450#issuecomment-1216334866

   Thanks @junrushao1994 I quite agree with your suggestion. Currently I modify it according to scenario we have encountered. Do you have any suggestion about the better name?


-- 
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


[GitHub] [tvm] yincs-intellif commented on pull request #12450: [TIR][Schedule] enhance compute_at primitive to choose proper position

Posted by GitBox <gi...@apache.org>.
yincs-intellif commented on PR #12450:
URL: https://github.com/apache/tvm/pull/12450#issuecomment-1216398026

   Hi @junrushao1994 , what do you think of my follow suggestion?
   1. we give a integer parameter instead of to_early_stage, which means block number, masked as N.
   2. Default we will get a compute at position, which is after the last producer  and before the first consumer, masked as [P, C].
   3. when N in [P,C], we use N for the position of compute at, otherwise use C.
   
   The above rules apply to "reverse_compute_at" at same time
   
   Thanks in advance!


-- 
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


[GitHub] [tvm] Hzfengsy commented on a diff in pull request #12450: [TIR][Schedule] enhance compute_at primitive to choose proper position

Posted by GitBox <gi...@apache.org>.
Hzfengsy commented on code in PR #12450:
URL: https://github.com/apache/tvm/pull/12450#discussion_r954445989


##########
src/tir/schedule/primitive/compute_at.cc:
##########
@@ -160,7 +174,57 @@ int FindInsertionPoint(
   // The valid indices are: (last_producer_position, first_consumer_position]
   ICHECK(split.last_producer_position < split.first_consumer_position);
   // Step 4. Return the last valid insertion point
-  return split.first_consumer_position;
+  int insert_position = split.first_consumer_position;
+  if (index == -1) {
+    return insert_position;
+  } else if (index == -2) {

Review Comment:
   Why do we insert it to the last position in case 2?



-- 
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


[GitHub] [tvm] yincs-intellif commented on a diff in pull request #12450: [TIR][Schedule] enhance compute_at primitive to choose proper position

Posted by GitBox <gi...@apache.org>.
yincs-intellif commented on code in PR #12450:
URL: https://github.com/apache/tvm/pull/12450#discussion_r954440662


##########
src/tir/schedule/primitive/compute_at.cc:
##########
@@ -160,7 +174,57 @@ int FindInsertionPoint(
   // The valid indices are: (last_producer_position, first_consumer_position]
   ICHECK(split.last_producer_position < split.first_consumer_position);
   // Step 4. Return the last valid insertion point
-  return split.first_consumer_position;
+  int insert_position = split.first_consumer_position;
+  if (index == -1) {
+    return insert_position;
+  } else if (index == -2) {

Review Comment:
   1. For compute at primitive, most of time, we can just return last_producer_position + 1. When other blocks insert ahead of block B, block B maybe far way block C, this is why I add some code to adjust the position. But the adjustment scheme  is not good either. I agree return last_producer_position + 1.
   2. For reverse_compute_at, when last_producer_position=-1 we hope it insert the last position, which is mostly caused by parameters.



-- 
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


[GitHub] [tvm] yincs-intellif commented on a diff in pull request #12450: [TIR][Schedule] enhance compute_at primitive to choose proper position

Posted by GitBox <gi...@apache.org>.
yincs-intellif commented on code in PR #12450:
URL: https://github.com/apache/tvm/pull/12450#discussion_r954441541


##########
include/tvm/tir/schedule/schedule.h:
##########
@@ -430,9 +430,11 @@ class ScheduleNode : public runtime::Object {
    * \param block_rv The block to be moved
    * \param loop_rv The loop where the block to be moved under
    * \param preserve_unit_loops Whether to keep the trivial loops whose extents are 1
+   * \param index The block index of the loop body subtree blocks
+   * -1, -2 means keep closed to or away from it's consumer

Review Comment:
   Agree. Thanks.



-- 
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


[GitHub] [tvm] yincs-intellif commented on a diff in pull request #12450: [TIR][Schedule] enhance compute_at and reverse_compute_at primitive to choose possible position

Posted by GitBox <gi...@apache.org>.
yincs-intellif commented on code in PR #12450:
URL: https://github.com/apache/tvm/pull/12450#discussion_r954464119


##########
src/tir/schedule/primitive/compute_at.cc:
##########
@@ -160,7 +174,57 @@ int FindInsertionPoint(
   // The valid indices are: (last_producer_position, first_consumer_position]
   ICHECK(split.last_producer_position < split.first_consumer_position);
   // Step 4. Return the last valid insertion point
-  return split.first_consumer_position;
+  int insert_position = split.first_consumer_position;
+  if (index == -1) {
+    return insert_position;
+  } else if (index == -2) {
+    if (require_all_producers_visited) {
+      if (split.last_producer_position >= 0) {
+        insert_position = split.last_producer_position + 1;
+      } else {
+        insert_position = split.first_consumer_position;
+      }
+    } else if (require_all_consumers_visited) {
+      class Finder : public StmtVisitor {
+       public:
+        void VisitStmt_(const BlockRealizeNode* realize) final {
+          const BlockNode* block = realize->block.get();
+          if (producer_blocks_.count(block)) {
+            ++this->n_producers_visited_;
+          }
+        }
+
+        std::unordered_set<const StmtNode*> producer_blocks_;
+        int n_producers_visited_ = 0;
+      };
+      // adjust the inserted position by compute at order
+      for (int i = split.first_consumer_position; i - 1 > split.last_producer_position; --i) {
+        auto blk = GetBlock(subtrees[i]);
+        if (!blk.defined()) break;

Review Comment:
   I have deleted this part of the code



-- 
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