MappedScop::detectReductions: use dependences to check validity of ordering

The reduction matching moves out any statements other than
the reduction updates to be able to do a pure mapping
to library calls.
The original code took into account that the init statement
that corresponds to a reduction update statement can be moved
in front of that update statement, but it would punt
on any statements other than those explicitly marked
as init statements.
Use the dependences to check whether moving out the other statements
is valid instead.  This allows the matching to work in case
there are any other statements (that can be moved out) or
if the reduction init statements are not explicitly marked.

Author: Sven Verdoolaege

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

@@ -118,52 +118,19 @@ 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)) {
       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