halide2isl: do not require init statement for reduction detection

Since the previous commit, the init statement is no longer used
by the TC mapper, so there is no longer any need to keep track of it.
This in turn means that there is also no longer any need to require
the presence of an init statement.
The new code does require that there is at least one reduction
dimension in order to prune out trivial reductions.

For example, in the TC below, both statements would
otherwise be treated as reduction updates.

    def fun(float(N,K) A, float(K,M) B, float(N,M) C) -> (O) {
      O(i,j) +=! A(i,k) * B(k,j)
      O(i,j) = O(i,j) + C(i,j)
    }

Note that the reduction is still identified by the tensor name,
meaning that multiple reductions on the same tensor are still
not allowed.  Identifying the reduction by the update statement
is left for future work.

Author: Sven Verdoolaege

tc/core/halide2isl.cc

@@ -16,7 +16,6 @@
 #include "tc/core/halide2isl.h"
 
 #include <algorithm>
-#include <numeric>
 #include <unordered_set>
 
 #include "tc/core/constants.h"
@@ -499,68 +498,42 @@ std::vector<Reduction> findReductions(const Stmt& s) {
       }
     }
 
-    // Keep track of any reduction variable name for use in isValidReduction
+    // Keep track of any reduction variable name for use in visit(Provide*)
     void visit(const Variable* op) {
       if (op->reduction_domain.defined()) {
         reductionVars.insert(op->name);
       }
     }
 
-    // Check that the given update node, together with the corresponding
-    // init node form a proper reduction pair.
-    // In particular, check that they share some outer For nodes and
-    // that the variables of the additional For nodes surrounding
-    // the update node are all reduction variables.
-    bool isValidReductionUpdate(const Provide* op) {
-      const auto& opInitVars = initVars[op->name];
-      auto n = opInitVars.size();
-      if (vars.size() <= n) {
-        return false;
-      }
-      if (!std::equal(opInitVars.begin(), opInitVars.end(), vars.begin())) {
-        return false;
-      }
-      for (auto i = vars.begin() + n; i != vars.end(); ++i) {
-        if (reductionVars.count(*i) == 0) {
-          return false;
-        }
-      }
-      return true;
-    }
-
     // Keep track of the names of the outer For nodes.
     void visit(const For* op) {
       vars.push_back(op->name);
       IRVisitor::visit(op);
       vars.pop_back();
     }
 
-    // Check if the node is an init node, keeping track of it,
-    // or an update node corresponding to init node that was found before,
-    // updating the information about the reduction.
-    // In particular, double-check that the pair are in the right
-    // relative positions and collect the positions of the reduction
+    // Check if the node is an update node with at least one reduction
+    // dimension, keeping track of the information about the reduction.
+    // In particular, collect the positions of the reduction
     // dimensions in the update statement domain.
     // Visit the children first to ensure that all relevant
     // reduction variables have been found first.
     void visit(const Provide* op) {
       IRVisitor::visit(op);
-      if (isReductionInit(op)) {
-        reductions[op->name].init = op;
-        initVars[op->name] = vars;
-      } else if (isReductionUpdate(op)) {
-        if (isValidReductionUpdate(op)) {
+      if (isReductionUpdate(op)) {
+        std::vector<size_t> dims;
+        auto n = vars.size();
+        for (size_t i = 0; i < n; ++i) {
+          if (reductionVars.count(vars[i]) != 0) {
+            dims.emplace_back(i);
+          }
+        }
+        if (dims.size() > 0) {
           auto& p = reductions[op->name];
-          CHECK(p.init.defined())
-              << "Missing reduction init or (unsupported) multiple updates";
           CHECK(!p.update.defined())
               << "Multiple reduction updates not yet implemented";
           p.update = op;
-          auto n = initVars[op->name].size();
-          p.dims.resize(vars.size() - n);
-          std::iota(p.dims.begin(), p.dims.end(), n);
-        } else {
-          reductions.erase(op->name);
+          p.dims = dims;
         }
       }
     }
@@ -570,8 +543,6 @@ std::vector<Reduction> findReductions(const Stmt& s) {
     std::unordered_set<std::string> reductionVars;
     // The names of the outer For nodes, outermost to innermost.
     std::vector<std::string> vars;
-    // For each init node, the names of its outer For nodes.
-    std::map<std::string, std::vector<std::string>> initVars;
     std::map<std::string, Reduction> reductions;
   } finder;
   s.accept(&finder);

tc/core/halide2isl.h

@@ -94,12 +94,11 @@ ScheduleTreeAndAccesses makeScheduleTree(
 /// Enumerate all reductions in a statement, by looking for the
 /// ReductionInit and ReductionUpdate markers inserted during lowering
 /// (see tc2halide.h).
-/// Each reduction object stores a reference to the init and
-/// the update statement, although the init statement is probably
-/// not strictly needed, and a list of reduction dimensions
+/// Each reduction object stores a reference to
+/// the update statement, and a list of reduction dimensions
 /// in the domain of the update statement.
 struct Reduction {
-  Halide::Internal::Stmt init, update;
+  Halide::Internal::Stmt update;
   std::vector<size_t> dims;
 };
 std::vector<Reduction> findReductions(const Halide::Internal::Stmt& s);