reductionInitsUpdates: return updates as isl::union_set

Return inits and updates in the same way improves consistency.
The updates in isl::union_set will also be useful
for simplifying findFirstReductionDim.

Author: Sven Verdoolaege

tc/core/polyhedral/cuda/mapped_scop.cc

@@ -211,9 +211,13 @@ 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());
@@ -229,7 +233,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,37 +134,37 @@ 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) {

tc/core/polyhedral/schedule_tree_matcher.h

@@ -26,12 +26,12 @@ 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);