replace findFirstReductionDim by isReductionMember

The original implementation would essentially take any band member
that looks like a reduction and simply take the outermost of those.
Later on, the choice of reduction schedule dimension was made more
stringent to appear at a specific position.
Instead of looking for the first reduction member and then checking
if it is in the right position, simply take the member in the right
position and check if it corresponds to a reduction.

Author: Sven Verdoolaege

tc/core/polyhedral/cuda/mapped_scop.cc

@@ -220,8 +220,9 @@ bool MappedScop::detectReductions(detail::ScheduleTree* tree) {
   });
   // 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);

tc/core/polyhedral/reduction_matcher.cc

@@ -167,23 +167,14 @@ std::pair<isl::union_set, isl::union_set> reductionInitsUpdates(
   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

@@ -35,10 +35,12 @@ 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