Merge pull request #448 from facebookresearch/pr/no_init

stop relying on reduction init statements

Author: ftynse

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

tc/core/polyhedral/cuda/mapped_scop.cc

@@ -215,9 +215,8 @@ bool MappedScop::detectReductions(detail::ScheduleTree* tree) {
   // a single reduction for now.
   // Support for multiple reductions would require a check
   // that these reductions do not interfere with each other.
-  auto initsUpdates = reductionInitsUpdates(band->mupa_.domain(), scop());
-  auto inits = initsUpdates.first;
-  auto updates = initsUpdates.second;
+  auto domain = band->mupa_.domain();
+  auto updates = reductionUpdates(domain, scop());
   if (updates.n_set() != 1) {
     return false;
   }
@@ -232,12 +231,19 @@ bool MappedScop::detectReductions(detail::ScheduleTree* tree) {
   if (!isReductionMember(member, updates, scop())) {
     return false;
   }
-  // Order the init statements (if any) before the update statements
+  // Order the other statements (if any) before the update statements
   // to ensure the band from which the reduction band has been split off
   // only contains update statements.
-  // Note that this relies on the outer members being coincident.
-  if (!inits.is_empty()) {
-    orderBefore(scop_->scheduleRoot(), tree, inits);
+  // Only do this if it doesn't violate any dependences.
+  // TODO (#454): order statements before or after the reduction based on
+  // dependences.
+  auto other = domain.subtract(updates);
+  if (!other.is_empty()) {
+    auto dependences = scop_->activeDependences(tree);
+    if (!canOrderBefore(scop_->scheduleRoot(), tree, other, dependences)) {
+      return false;
+    }
+    orderBefore(scop_->scheduleRoot(), tree, other);
   }
   reductionBandUpdates_.emplace(tree, Reduction(updateIds));
   return true;

tc/core/polyhedral/reduction_matcher.cc

@@ -48,12 +48,10 @@ bool isSupportedReduction(Halide::Internal::Stmt stmt) {
 // the reduction.  that is a kind of internal state dependence we want to avoid
 // If id is the statement identifier of an update statement
 // of a supported type of reduction,
-// then return the corresponding init statement in init and
-// the corresponding reduction dimensions in reductionDims.
+// then return the corresponding reduction dimensions in reductionDims.
 bool isReductionUpdateId(
     isl::id id,
     const Scop& scop,
-    Halide::Internal::Stmt& init,
     std::vector<size_t>& reductionDims) {
   CHECK_EQ(scop.halide.statements.count(id), 1u)
       << "id is not a statement in scop" << id;
@@ -63,7 +61,6 @@ bool isReductionUpdateId(
   }
   for (auto const& iup : scop.halide.reductions) {
     if (iup.update.same_as(provideNode)) {
-      init = iup.init;
       reductionDims = iup.dims;
       return true;
     }
@@ -104,9 +101,8 @@ bool isAlmostIdentityReduction(isl::pw_aff pa, const Scop& scop) {
     return false;
   }
   auto stmtId = space.get_tuple_id(isl::dim_type::in);
-  Halide::Internal::Stmt init;
   std::vector<size_t> reductionDims;
-  if (!isReductionUpdateId(stmtId, scop, init, reductionDims)) {
+  if (!isReductionUpdateId(stmtId, scop, reductionDims)) {
     return false;
   }
 
@@ -118,52 +114,18 @@ bool isAlmostIdentityReduction(isl::pw_aff pa, const Scop& scop) {
   return false;
 }
 
-/*
- * Return the identifier that maps to "stmt".
- */
-isl::id statementId(const Scop& scop, const Halide::Internal::Stmt& stmt) {
-  for (auto kvp : scop.halide.statements) {
-    if (kvp.second.same_as(stmt)) {
-      return kvp.first;
-    }
-  }
-  CHECK(false) << "no id recorded for statement" << stmt;
-  return isl::id();
-}
-
 } // namespace
 
-std::pair<isl::union_set, isl::union_set> reductionInitsUpdates(
-    isl::union_set domain,
-    const Scop& scop) {
-  auto initUnion = isl::union_set::empty(domain.get_space());
-  auto update = initUnion;
-  std::unordered_set<isl::id, isl::IslIdIslHash> init;
-  std::vector<isl::set> nonUpdate;
-  // First collect all the update statements,
-  // the corresponding init statement and all non-update statements.
-  domain.foreach_set([&init, &update, &nonUpdate, &scop](isl::set set) {
+isl::union_set reductionUpdates(isl::union_set domain, const Scop& scop) {
+  auto update = isl::union_set::empty(domain.get_space());
+  domain.foreach_set([&update, &scop](isl::set set) {
     auto setId = set.get_tuple_id();
-    Halide::Internal::Stmt initStmt;
     std::vector<size_t> reductionDims;
-    if (isReductionUpdateId(setId, scop, initStmt, reductionDims)) {
+    if (isReductionUpdateId(setId, scop, reductionDims)) {
       update = update.unite(set);
-      init.emplace(statementId(scop, initStmt));
-    } else {
-      nonUpdate.emplace_back(set);
     }
   });
-  // Then check if all the non-update statements are init statements
-  // that correspond to the update statements found.
-  // If not, return an empty list of update statements.
-  for (auto set : nonUpdate) {
-    if (init.count(set.get_tuple_id()) != 1) {
-      return std::pair<isl::union_set, isl::union_set>(
-          initUnion, isl::union_set::empty(domain.get_space()));
-    }
-    initUnion = initUnion.unite(set);
-  }
-  return std::pair<isl::union_set, isl::union_set>(initUnion, update);
+  return update;
 }
 
 bool isReductionMember(

tc/core/polyhedral/schedule_transforms.cc

@@ -731,6 +731,25 @@ ScheduleTreeUPtr gistedFilter(isl::union_set filter, ScheduleTreeUPtr child) {
 
 } // namespace
 
+bool canOrderBefore(
+    ScheduleTree* root,
+    ScheduleTree* tree,
+    isl::union_set filter,
+    isl::union_map dependences) {
+  // Create an ordering schedule function filter -> 0; other -> 1.
+  auto other = activeDomainPoints(root, tree).subtract(filter);
+  auto ctx = root->ctx_;
+  auto space = isl::space(ctx, 0).unnamed_set_from_params(1);
+  auto zero = isl::multi_val::zero(space);
+  auto one = zero.set_val(0, isl::val::one(ctx));
+  auto order = isl::multi_union_pw_aff(filter, zero);
+  order = order.union_add(isl::multi_union_pw_aff(other, one));
+
+  // Check that this ordering preserves all dependences.
+  auto preserved = dependences.lex_lt_at(order).unite(dependences.eq_at(order));
+  return dependences.is_subset(preserved);
+}
+
 void orderBefore(
     ScheduleTree* root,
     ScheduleTree* tree,

tc/core/polyhedral/schedule_transforms.h

@@ -233,6 +233,15 @@ void insertExtensionLabelAfter(
     detail::ScheduleTree* tree,
     isl::id id);
 
+// Is it possible to order the elements in the given filter
+// before the other active elements without violating
+// any of the given dependences?
+bool canOrderBefore(
+    detail::ScheduleTree* root,
+    detail::ScheduleTree* tree,
+    isl::union_set filter,
+    isl::union_map dependences);
+
 // Insert a sequence to ensure that the active domain elements
 // in the given filter are executed before the other active domain elements.
 void orderBefore(

tc/core/polyhedral/schedule_tree_matcher.h

@@ -25,15 +25,9 @@
 namespace tc {
 namespace polyhedral {
 
-// Return the union of the reduction init statements as well as
-// the union of the reduction update statements
-// that appear in "domain", assuming "domain" only contains
-// reduction init and update statements.
-// If "domain" contains any other statements, then return an empty set
-// of reduction update statements.
-std::pair<isl::union_set, isl::union_set> reductionInitsUpdates(
-    isl::union_set domain,
-    const Scop& scop);
+// Return the union of the reduction update statements
+// that appear in "domain".
+isl::union_set reductionUpdates(isl::union_set domain, const Scop& scop);
 
 // Does the band member with the given partial schedule correspond
 // to a reduction on all statements with a domain in "domain"?

tc/core/polyhedral/scop.cc

@@ -390,6 +390,16 @@ void Scop::computeAllDependences() {
   dependences = flowDeps.unite(falseDeps).coalesce();
 }
 
+isl::union_map Scop::activeDependences(detail::ScheduleTree* tree) {
+  auto prefix = prefixScheduleMupa(scheduleRoot(), tree);
+  auto domain = activeDomainPoints(scheduleRoot(), tree);
+  auto active = dependences;
+  active = active.intersect_domain(domain);
+  active = active.intersect_range(domain);
+  active = active.eq_at(prefix);
+  return active;
+}
+
 std::unique_ptr<detail::ScheduleTree> Scop::computeSchedule(
     isl::schedule_constraints constraints,
     const SchedulerOptionsView& schedulerOptions) {

tc/core/polyhedral/scop.h

@@ -471,6 +471,9 @@ struct Scop {
   // Do the simplest possible dependence analysis.
   // Compute all RAW, WAR, and WAW dependences, and save them in dependences.
   void computeAllDependences();
+  // Return the set of dependences that are active
+  // at the given position.
+  isl::union_map activeDependences(detail::ScheduleTree* tree);
 
  public:
   // Halide stuff

tc/core/tc2halide.cc

@@ -834,38 +834,6 @@ HalideComponents translateDef(const lang::Def& def, bool throwWarnings) {
     s = tagReduction.mutate(s);
   }
 
-  // Temporary hack: Fuse reduction initializations into the
-  // updates. Shouldn't matter because we're going to reschedule
-  // everything anyway, but starting at this point is more similar to
-  // the code this replaces. Only correct because we only currently
-  // create update definitions when there are reductions.
-  class FuseReductions : public IRMutator2 {
-    using IRMutator2::visit;
-    Stmt visit(const Block* op) override {
-      const For* first = op->first.as<For>();
-      const For* rest = op->rest.as<For>();
-      if (first && rest && equal(first->min, rest->min) &&
-          equal(first->extent, rest->extent) &&
-          first->for_type == rest->for_type &&
-          replace_all(first->name, ".s0.", ".s1.") == rest->name) {
-        Stmt body = rest->body;
-        body =
-            substitute(rest->name, Variable::make(Int(32), first->name), body);
-        body = mutate(Block::make(first->body, body));
-        return For::make(
-            first->name,
-            first->min,
-            first->extent,
-            first->for_type,
-            first->device_api,
-            body);
-      } else {
-        return IRMutator2::visit(op);
-      }
-    }
-  } fuser;
-  s = fuser.mutate(s);
-
   // Trim ProducerConsumer annotations. TC doesn't use them.
   class RemoveProducerConsumer : public IRMutator2 {
     using IRMutator2::visit;