[TensorIR] Cross-Thread Reduction

src/tir/schedule/analysis.h

@@ -19,6 +19,7 @@
 #ifndef TVM_TIR_SCHEDULE_ANALYSIS_H_
 #define TVM_TIR_SCHEDULE_ANALYSIS_H_
 
+#include <tvm/arith/analyzer.h>
 #include <tvm/tir/schedule/state.h>
 
 #include <tuple>
@@ -331,16 +332,14 @@ Buffer GetNthAccessBuffer(const ScheduleState& self, const Block& block, int n,
 
 /*!
  * \brief Convert the `init` and `body` of the input block to BufferStores
- * \tparam in_schedule Whether the function is called by schedule primitives
  * \param self The schedule state
  * \param block The block to be analyzed
  * \return The BufferStores of the `init` and `body` of the input block
  * \throw ScheduleError If the `init` or `body` is not BufferStore, or they don't write to the same
  * buffer
  */
-template <bool in_schedule>
-std::pair<BufferStore, BufferStore> GetBufferStoresFromReductionBlock(const ScheduleState& self,
-                                                                      const Block& block);
+std::pair<BufferStore, BufferStore> GetBufferStoresFromReductionBlock(
+    const Optional<ScheduleState>& self, const Block& block);
 
 /*!
  * \brief Check whether the input array of IterVars only contains data-parallel and reduction block
@@ -363,16 +362,14 @@ bool ReductionIterNotIndexOutputBuffer(const Block& block);
 /*!
  * \brief Given a reduction identity and a reduction combiner, detect the corresponding commutative
  * reducer, and extract the combiner lhs and combiner rhs
- * \tparam in_schedule Whether the function is called by schedule primitives
  * \param self The schedule state
  * \param identity The reduction identity to be analyzed
  * \param combiner The reduction combiner to be analyzed
  * \return The corresponding CommReducer, the combiner lhs and the combiner rhs
  * \throw ScheduleError If no corresponding commutative reducer can be matched
  */
-template <bool in_schedule>
 std::tuple<CommReducer, PrimExpr, PrimExpr> GetReducerAndCombinerLhsRhs(
-    const ScheduleState& self, const PrimExpr& identity, const BufferStore& combiner);
+    const Optional<ScheduleState>& self, const PrimExpr& identity, const BufferStore& combiner);
 
 /******** Commutative Reducer ********/
 
@@ -381,7 +378,7 @@ std::tuple<CommReducer, PrimExpr, PrimExpr> GetReducerAndCombinerLhsRhs(
  * \return The list of the registered reducer-getter functions
  * \sa ReducerRegistry
  */
-std::vector<TypedPackedFunc<CommReducer(DataType)>> GetReducerGetters();
+std::vector<runtime::TypedPackedFunc<CommReducer(DataType)>> GetReducerGetters();
 
 /*!
  * \brief Given the input identity and the combiner BufferStore of a reduction, extract the

src/tir/schedule/analysis/analysis.cc

@@ -523,7 +523,6 @@ bool GetVarsTouchedByBlockIters(const BlockRealize& block_realize,
     if (set == nullptr) {
       continue;
     }
-
     Array<Var> vars_in_binding = UndefinedVars(iter_value);
     for (const Var& var : vars_in_binding) {
       set->insert(var.get());
@@ -1166,39 +1165,38 @@ class InitBodyNotSameBufferAccessError : public ScheduleError {
   Block block_;
 };
 
-template <bool in_schedule>
-std::pair<BufferStore, BufferStore> GetBufferStoresFromReductionBlock(const ScheduleState& self,
-                                                                      const Block& block) {
-  const char* error_str1 =
+std::pair<BufferStore, BufferStore> GetBufferStoresFromReductionBlock(
+    const Optional<ScheduleState>& self, const Block& block) {
+  static constexpr const char* error_str1 =
       "ValueError: The `init` and `body` of the reduction block are required to be both "
       "BufferStore so that rfactor or cross-thread reduction can be applied. However, a reduction "
       "block that doesn't meet this requirement is ";
-  const char* error_str2 =
+  static constexpr const char* error_str2 =
       "ValueError: The `init` and `body` of the reduction block are required to have the same "
       "buffer access pattern so that rfactor or cross-thread reduction can be applied. However, a "
       "reduction block that doesn't meet this requirement is ";
 
   const auto* init = block->init.as<BufferStoreNode>();
   const auto* body = block->body.as<BufferStoreNode>();
   if (!(init && body)) {
-    if (in_schedule) {
-      throw InitBodyNotBufferStoreError(self->mod, block, init != nullptr, body != nullptr);
+    if (self.defined()) {
+      throw InitBodyNotBufferStoreError(self.value()->mod, block, init != nullptr, body != nullptr);
     } else {
       LOG(FATAL) << error_str1 << block;
     }
   }
   if (!init->buffer.same_as(body->buffer)) {
-    if (in_schedule) {
-      throw InitBodyNotSameBufferAccessError(self->mod, block);
+    if (self.defined()) {
+      throw InitBodyNotSameBufferAccessError(self.value()->mod, block);
     } else {
       LOG(FATAL) << error_str2 << block;
     }
   }
   int ndim = static_cast<int>(init->buffer->shape.size());
   for (int i = 0; i < ndim; ++i) {
     if (!ExprDeepEqual()(init->indices[i], body->indices[i])) {
-      if (in_schedule) {
-        throw InitBodyNotSameBufferAccessError(self->mod, block);
+      if (self.defined()) {
+        throw InitBodyNotSameBufferAccessError(self.value()->mod, block);
       } else {
         LOG(FATAL) << error_str2 << block;
       }
@@ -1231,26 +1229,30 @@ bool ReductionIterNotIndexOutputBuffer(const Block& block) {
   for (const BufferRegion& write_region : block->writes) {
     buffer_written.insert(write_region->buffer.get());
   }
+  auto f_uses_reduction_block_var = [&](const PrimExpr& expr) -> bool {
+    return UsesVar(expr, [&](const VarNode* var) {  //
+      return reduction_block_iters.count(var);
+    });
+  };
   bool affected = false;
   PreOrderVisit(block->body, [&](const ObjectRef& obj) {
     if (affected) {
       return false;
     }
-    if (const auto* store = obj.as<BufferStoreNode>()) {
-      ICHECK(buffer_written.count(store->buffer.get()))
-          << "ValueError: The buffer \"" << store->buffer
-          << "\" is written in the block but is not in the block's signature";
-      for (const PrimExpr& index : store->indices) {
-        if (UsesVar(index, [&reduction_block_iters](const VarNode* var) {
-              return reduction_block_iters.count(var);
-            })) {
-          affected = true;
-          return false;
-        }
+    const auto* store = obj.as<BufferStoreNode>();
+    if (!store) {
+      return true;
+    }
+    ICHECK(buffer_written.count(store->buffer.get()))
+        << "ValueError: The buffer \"" << store->buffer
+        << "\" is written in the block but is not in the block's signature";
+    for (const PrimExpr& index : store->indices) {
+      if (f_uses_reduction_block_var(index)) {
+        affected = true;
+        return false;
       }
-      return false;
     }
-    return true;
+    return false;
   });
   return !affected;
 }
@@ -1281,15 +1283,14 @@ class NoMatchedReducerError : public ScheduleError {
   BufferStore combiner_;
 };
 
-template <bool in_schedule>
 std::tuple<CommReducer, PrimExpr, PrimExpr> GetReducerAndCombinerLhsRhs(
-    const ScheduleState& self, const PrimExpr& identity, const BufferStore& combiner) {
+    const Optional<ScheduleState>& self, const PrimExpr& identity, const BufferStore& combiner) {
   CommReducer reducer{nullptr};
   PrimExpr combiner_lhs{nullptr}, combiner_rhs{nullptr};
   bool matched = FromIdentityCombiner(identity, combiner, &reducer, &combiner_lhs, &combiner_rhs);
   if (!matched) {
-    if (in_schedule) {
-      throw NoMatchedReducerError(self->mod, identity, combiner);
+    if (self.defined()) {
+      throw NoMatchedReducerError(self.value()->mod, identity, combiner);
     } else {
       LOG(FATAL) << "ValueError: No matched reducer for the identity and the combiner of the "
                     "reduction block. So rfactor and cross-thread reduction cannot be applied.";
@@ -1298,15 +1299,6 @@ std::tuple<CommReducer, PrimExpr, PrimExpr> GetReducerAndCombinerLhsRhs(
   return std::make_tuple(std::move(reducer), std::move(combiner_lhs), std::move(combiner_rhs));
 }
 
-template std::pair<BufferStore, BufferStore> GetBufferStoresFromReductionBlock<true>(
-    const ScheduleState& self, const Block& block);
-template std::pair<BufferStore, BufferStore> GetBufferStoresFromReductionBlock<false>(
-    const ScheduleState& self, const Block& block);
-template std::tuple<CommReducer, PrimExpr, PrimExpr> GetReducerAndCombinerLhsRhs<true>(
-    const ScheduleState& self, const PrimExpr& identity, const BufferStore& combiner);
-template std::tuple<CommReducer, PrimExpr, PrimExpr> GetReducerAndCombinerLhsRhs<false>(
-    const ScheduleState& self, const PrimExpr& identity, const BufferStore& combiner);
-
 /******** Commutative Reducer ********/
 
 bool MatchReducer(const CommReducer& reducer, const PrimExpr& identity, const PrimExpr& combiner,

src/tir/schedule/primitive/reduction.cc

@@ -1041,9 +1041,9 @@ StmtSRef RFactor(ScheduleState self, const StmtSRef& rf_loop_sref, int factor_ax
   BufferStore update;
   CommReducer reducer;
   PrimExpr combiner_lhs, combiner_rhs;
-  std::tie(init, update) = GetBufferStoresFromReductionBlock<true>(self, block);
+  std::tie(init, update) = GetBufferStoresFromReductionBlock(self, block);
   std::tie(reducer, combiner_lhs, combiner_rhs) =
-      GetReducerAndCombinerLhsRhs<true>(self, init->value, update);
+      GetReducerAndCombinerLhsRhs(self, init->value, update);
 
   // Step 6. Check whether `factor_axis` is in a correct range, and convert it to non-negative if it
   // is negative.