Merge pull request #266 from facebookresearch/pr/reduction

further simplify reduction handling

Author: ftynse

tc/core/polyhedral/cuda/mapped_scop.cc

@@ -211,13 +211,18 @@ bool MappedScop::detectReductions(detail::ScheduleTree* tree) {
   auto initsUpdates = reductionInitsUpdates(band->mupa_.domain(), scop());
   auto inits = initsUpdates.first;
   auto updates = initsUpdates.second;
-  if (updates.size() != 1) {
+  if (updates.n_set() != 1) {
     return false;
   }
+  std::vector<isl::id> updateIds;
+  updates.foreach_set([&updateIds](isl::set set) {
+    updateIds.emplace_back(set.get_tuple_id());
+  });
   // The reduction member needs to appear right underneath
   // the coincident members.
-  auto reductionDim = findFirstReductionDim(band->mupa_, scop());
-  if (reductionDim != nCoincident) {
+  auto reductionDim = nCoincident;
+  auto member = band->mupa_.get_union_pw_aff(reductionDim);
+  if (!isReductionMember(member, updates, scop())) {
     return false;
   }
   auto reductionTree = bandSplitOut(scop_->scheduleRoot(), tree, reductionDim);
@@ -229,7 +234,7 @@ bool MappedScop::detectReductions(detail::ScheduleTree* tree) {
     orderBefore(scop_->scheduleRoot(), tree, inits);
   }
   reductionFromParent_.emplace(tree, reductionTree);
-  reductionBandUpdates_.emplace(reductionTree, updates);
+  reductionBandUpdates_.emplace(reductionTree, updateIds);
   return true;
 }
 

tc/core/polyhedral/reduction_matcher.cc

@@ -134,56 +134,47 @@ isl::id statementId(const Scop& scop, const Halide::Internal::Stmt& stmt) {
 
 } // namespace
 
-std::pair<isl::union_set, std::vector<isl::id>> reductionInitsUpdates(
+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());
-  std::vector<isl::id> update;
+  auto update = initUnion;
   std::unordered_set<isl::id, isl::IslIdIslHash> init;
   std::vector<isl::set> nonUpdate;
-  // First collect all the update statement identifiers,
+  // 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) {
     auto setId = set.get_tuple_id();
     Halide::Internal::Stmt initStmt;
     std::vector<size_t> reductionDims;
     if (isReductionUpdateId(setId, scop, initStmt, reductionDims)) {
-      update.emplace_back(setId);
+      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 statement identifiers.
+  // 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, std::vector<isl::id>>(
-          initUnion, std::vector<isl::id>());
+      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, std::vector<isl::id>>(initUnion, update);
+  return std::pair<isl::union_set, isl::union_set>(initUnion, update);
 }
 
-int findFirstReductionDim(isl::multi_union_pw_aff islMupa, const Scop& scop) {
-  auto mupa = isl::MUPA(islMupa);
-  int reductionDim = -1;
-  int currentDim = 0;
-  for (auto const& upa : mupa) {
-    for (auto const& pa : upa) {
-      if (isAlmostIdentityReduction(pa.pa, scop)) {
-        reductionDim = currentDim;
-        break;
-      }
-    }
-    if (reductionDim != -1) {
-      break;
-    }
-    ++currentDim;
-  }
-  return reductionDim;
+bool isReductionMember(
+    isl::union_pw_aff member,
+    isl::union_set domain,
+    const Scop& scop) {
+  return domain.every_set([member, &scop](isl::set set) {
+    auto pa = member.extract_on_domain(set.get_space());
+    return isAlmostIdentityReduction(pa, scop);
+  });
 }
 
 } // namespace polyhedral

tc/core/polyhedral/schedule_tree_matcher.h

@@ -26,19 +26,21 @@ namespace tc {
 namespace polyhedral {
 
 // Return the union of the reduction init statements as well as
-// the identifiers of all reduction update statements
+// 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 vector
-// of identifiers.
-std::pair<isl::union_set, std::vector<isl::id>> reductionInitsUpdates(
+// 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);
 
-// Find the first band member that corresponds to a reduction.
-// TODO: heuristic to choose the "best" band member in presence of multiple
-// reductions.
-int findFirstReductionDim(isl::multi_union_pw_aff islMupa, const Scop& scop);
+// Does the band member with the given partial schedule correspond
+// to a reduction on all statements with a domain in "domain"?
+bool isReductionMember(
+    isl::union_pw_aff member,
+    isl::union_set domain,
+    const Scop& scop);
 
 } // namespace polyhedral
 } // namespace tc