Merge pull request #462 from facebookresearch/pr/order

MappedScop::detectReductions: order away statements that depend on updates

Author: ftynse

tc/core/polyhedral/cuda/mapped_scop.cc

@@ -166,6 +166,8 @@ void MappedScop::mapToBlocksAndScaleBand(
   bandScale(band, tileSizes);
 }
 
+namespace {
+
 /*
  * Given a node in the schedule tree of a mapped scop,
  * insert a mapping filter underneath (if needed) that fixes
@@ -186,6 +188,41 @@ void fixThreadsBelow(
   mscop.mapThreadsBackward(bandTree);
 }
 
+/*
+ * Try and order the other statements in "domain" (if any)
+ * away from the "updates" statements, returning true is the operation succeeds.
+ * In particular, only do this if it doesn't violate any dependences.
+ * Anything that depends on an update statement is ordered after
+ * the update statements.  Anything else is ordered before.
+ */
+bool separatedOut(
+    Scop& scop,
+    detail::ScheduleTree* tree,
+    isl::union_set domain,
+    isl::union_set updates) {
+  auto other = domain.subtract(updates);
+  if (other.is_empty()) {
+    return true;
+  }
+  auto dependences = scop.activeDependences(tree);
+  auto after =
+      dependences.intersect_domain(updates).intersect_range(other).range();
+  auto before = other.subtract(after);
+  if (!canOrderBefore(scop.scheduleRoot(), tree, before, dependences) ||
+      !canOrderAfter(scop.scheduleRoot(), tree, after, dependences)) {
+    return false;
+  }
+  if (!before.is_empty()) {
+    orderBefore(scop.scheduleRoot(), tree, before);
+  }
+  if (!after.is_empty()) {
+    orderAfter(scop.scheduleRoot(), tree, after);
+  }
+  return true;
+}
+
+} // namespace
+
 bool MappedScop::detectReductions(detail::ScheduleTree* tree) {
   bool found = false;
   for (auto c : tree->children()) {
@@ -234,16 +271,8 @@ bool MappedScop::detectReductions(detail::ScheduleTree* tree) {
   // 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.
-  // 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);
+  if (!separatedOut(scop(), tree, domain, updates)) {
+    return false;
   }
   reductionBandUpdates_.emplace(tree, Reduction(updateIds));
   return true;

tc/core/polyhedral/schedule_transforms.cc

@@ -729,25 +729,49 @@ ScheduleTreeUPtr gistedFilter(isl::union_set filter, ScheduleTreeUPtr child) {
   return ScheduleTree::makeFilter(filter, std::move(child));
 }
 
+/*
+ * Given a partition of the (active) domain elements into "first" and "second",
+ * is it possible to order the "first" elements before the "second"
+ * without violating any of the (active) "dependences"?
+ */
+bool canOrder(
+    isl::union_set first,
+    isl::union_set second,
+    isl::union_map dependences) {
+  if (first.is_empty() || second.is_empty()) {
+    return true;
+  }
+  // Create an ordering schedule function first -> 0; second -> 1.
+  auto ctx = dependences.get_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(first, zero);
+  order = order.union_add(isl::multi_union_pw_aff(second, 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);
+}
+
 } // 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));
+  return canOrder(filter, other, dependences);
+}
 
-  // Check that this ordering preserves all dependences.
-  auto preserved = dependences.lex_lt_at(order).unite(dependences.eq_at(order));
-  return dependences.is_subset(preserved);
+bool canOrderAfter(
+    ScheduleTree* root,
+    ScheduleTree* tree,
+    isl::union_set filter,
+    isl::union_map dependences) {
+  auto other = activeDomainPoints(root, tree).subtract(filter);
+  return canOrder(other, filter, dependences);
 }
 
 void orderBefore(

tc/core/polyhedral/schedule_transforms.h

@@ -241,6 +241,14 @@ bool canOrderBefore(
     detail::ScheduleTree* tree,
     isl::union_set filter,
     isl::union_map dependences);
+// Is it possible to order the elements in the given filter
+// after the other active elements without violating
+// any of the given dependences?
+bool canOrderAfter(
+    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.

test/test_cuda_mapper.cc

@@ -867,7 +867,8 @@ struct ReductionTest : public PolyhedralMapperTest {
         .outerScheduleFusionStrategy(tc::FusionStrategy::Preserve3Coincident)
         .outerScheduleAllowSkewing(false)
         .outerSchedulePositiveOrthant(true)
-        .intraTileScheduleFusionStrategy(tc::FusionStrategy::Min)
+        .intraTileScheduleFusionStrategy(
+            tc::FusionStrategy::Preserve3Coincident)
         .intraTileScheduleAllowSkewing(false)
         .intraTileSchedulePositiveOrthant(true)
         .fixParametersBeforeScheduling(false)