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/06/09 04:42:09 UTC

[GitHub] [tvm] MasterJH5574 opened a new pull request, #11639: [TIR][Meta-Schedule] Tuple-reduction scheduling support

MasterJH5574 opened a new pull request, #11639:
URL: https://github.com/apache/tvm/pull/11639

   This PR improves our TIR scheduling primitives/transformations (rfactor & cross-thread reduction) designed for reduction operators, so that they can be applied to blocks of tuple-reduction.
   
   Most of the code changes are of unit tests.
   
   cc @junrushao1994 


-- 
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] junrushao1994 commented on pull request #11639: [TIR][Meta-Schedule] Tuple-reduction scheduling support

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

   Thanks a lot!! Having tuple reduction in TIR is definitely one of my top priority! Need some time to recharge this week and will review next week!


-- 
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] MasterJH5574 commented on a diff in pull request #11639: [TIR][Meta-Schedule] Tuple-reduction scheduling support

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


##########
tests/python/unittest/test_tir_transform_lower_cross_thread_reduction.py:
##########
@@ -44,7 +47,7 @@ def loop_split(a: T.handle, b: T.handle) -> None:
             with T.block("B"):
                 vi = T.axis.S(128, i)
                 vk = T.axis.R(128, ko * 32 + ki)
-                T.reads([B[vi], A[vi, vk]])
+                T.reads([A[vi, vk]])

Review Comment:
   I guess starting from some commit, we removed all reduction buffers from `T.reads` 🤔.
   
   One example in other test cases is https://github.com/apache/tvm/blob/9671aee942503815ad2a586406eef11391287ee5/tests/python/unittest/test_tvmscript_roundtrip.py#L2716-L2717
   
   But I can’t remember when it starts from.



-- 
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 merged pull request #11639: [TIR][Meta-Schedule] Tuple-reduction scheduling support

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


-- 
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 #11639: [TIR][Meta-Schedule] Tuple-reduction scheduling support

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


##########
src/meta_schedule/schedule_rule/cross_thread_reduction.cc:
##########
@@ -184,6 +184,13 @@ class CrossThreadReductionNode : public ScheduleRuleNode {
    */
   std::tuple<bool, tir::LoopRV, tir::BlockRV, tir::LoopRV> GetComputeTargetLoopAndBlock(
       const tir::Schedule& sch, const tir::BlockRV& block_rv) {
+    // Step 0. Due to technical reason of some primitives (e.g., compute-at), if the block is doing
+    // a tuple reduction, fusion is temporarily not supported.

Review Comment:
   nit: let's improve the doc a bit by being more specific, like, "fusion with the subsequent consumer is not allowed"



-- 
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 #11639: [TIR][Meta-Schedule] Tuple-reduction scheduling support

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


##########
tests/python/unittest/test_tir_transform_lower_cross_thread_reduction.py:
##########
@@ -44,7 +47,7 @@ def loop_split(a: T.handle, b: T.handle) -> None:
             with T.block("B"):
                 vi = T.axis.S(128, i)
                 vk = T.axis.R(128, ko * 32 + ki)
-                T.reads([B[vi], A[vi, vk]])
+                T.reads([A[vi, vk]])

Review Comment:
   Why are we removing `B[vi]` from `T.reads`? It's actually read, isn't it?



-- 
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] MasterJH5574 commented on a diff in pull request #11639: [TIR][Meta-Schedule] Tuple-reduction scheduling support

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


##########
tests/python/unittest/test_tir_transform_lower_cross_thread_reduction.py:
##########
@@ -44,7 +47,7 @@ def loop_split(a: T.handle, b: T.handle) -> None:
             with T.block("B"):
                 vi = T.axis.S(128, i)
                 vk = T.axis.R(128, ko * 32 + ki)
-                T.reads([B[vi], A[vi, vk]])
+                T.reads([A[vi, vk]])

Review Comment:
   My pleasure!



-- 
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 #11639: [TIR][Meta-Schedule] Tuple-reduction scheduling support

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


##########
tests/python/unittest/test_tir_transform_lower_cross_thread_reduction.py:
##########
@@ -44,7 +47,7 @@ def loop_split(a: T.handle, b: T.handle) -> None:
             with T.block("B"):
                 vi = T.axis.S(128, i)
                 vk = T.axis.R(128, ko * 32 + ki)
-                T.reads([B[vi], A[vi, vk]])
+                T.reads([A[vi, vk]])

Review Comment:
   I see! I just checked myself and found the same thing :-) Thanks for a heads-up!



##########
tests/python/unittest/test_tir_transform_lower_cross_thread_reduction.py:
##########
@@ -44,7 +47,7 @@ def loop_split(a: T.handle, b: T.handle) -> None:
             with T.block("B"):
                 vi = T.axis.S(128, i)
                 vk = T.axis.R(128, ko * 32 + ki)
-                T.reads([B[vi], A[vi, vk]])
+                T.reads([A[vi, vk]])

Review Comment:
   I see! I just checked myself and found the same thing :-) Thanks for the heads-up!



-- 
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 #11639: [TIR][Meta-Schedule] Tuple-reduction scheduling support

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


##########
src/tir/schedule/analysis/analysis.cc:
##########
@@ -1514,115 +1519,235 @@ class PatternMatcher : public ExprVisitor {
 
  private:
   bool match_success_{true};
-  PrimExpr pattern_, expr_to_match_;
+  Array<PrimExpr> pattern_;
+  PrimExpr expr_to_match_;
   std::unordered_map<const VarNode*, PrimExpr> filled_map_;
 };
 
 /******** Reduction Block Related ********/
 
-class InitBodyNotBufferStoreError : public ScheduleError {
- public:
-  explicit InitBodyNotBufferStoreError(IRModule mod, Block block, bool init_is_bufferstore,
-                                       bool body_is_bufferstore)
-      : mod_(std::move(mod)),
-        block_(std::move(block)),
-        init_is_bufferstore_(init_is_bufferstore),
-        body_is_bufferstore_(body_is_bufferstore) {}
-
-  String FastErrorString() const final {
-    return "ScheduleError: The `init` and `body` of reduction block are required to be both "
-           "BufferStore so that rfactor or cross-thread reduction can be applied";
-  }
+static const char* kRFactorCrossThreadReductionApplicableBlockDef =
+    R"(Definition of a reduction block that is applicable by RFactor and Cross-Thread Reduction:
+1) The block init should be a single BufferStore or a SeqStmt of BufferStores
+2) The buffers initialized in the block init should be all different
+3) The number of consecutive LetStmts in the block body (if any) should equal the number of BufferStores in the block init
+4) The variables of the LetStmts in the block body should be all different
+5) The body of the innermost LetStmt should be a single BufferStore or a SeqStmt of BufferStores
+6) The number of BufferStores under the block body should equal the number of BufferStores in the block init, and thereby equal the number of LetStmts above
+7) The variables bound by the LetStmts in the block body must all directly serve as values of the BufferStores inside, and the stored values of the BufferStores can only be those variables
+8) The variables stored by the BufferStores in the block body should be all different
+9) The buffers written by the BufferStores in the block body should be all different
+10) The buffers initialized in the block init and written in the block body should match
+11) The buffers written by the block should have same shape
+12) The indices of all BufferStores in the reduction block should be the same)";
+
+void ErrorRFactorCrossThreadReductionNotApplicable(const Optional<ScheduleState>& self, Block block,

Review Comment:
   nit: move the reducer-related logic into a separate file given they are relatively isolated and the logic is much longer



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