use polyhedral reduction description to detect reductions

The exact notion of what constitutes a "reduction" was
not spelled out very clearly, but in the use case
of mapping the detected reduction to CUB, each thread
needs to write to a single tensor element.
Take this property as the definition of a "reduction".

A necessary and sufficient condition is that
the schedule dimensions outside the band members
that are considered to form the reduction,
determine a unique tensor element that is updated.

The condition used in the original detection mechanism
(that the reduction band members only involve a specific
reduction dimension of the reduction update statement)
is neither sufficient nor necessary.
It is not sufficient because it doesn't force
the schedule dimensions not involving the reduction dimensions
to appear outside of those band members.
It is not clear if such a case could ever arise in practice
in TC, because the coincident band members get placed outermost.
However, this situation could change in the future if other
criteria such as consecutivity get taken into account.

It is not necessary because it is perfectly fine
for such a band member to reference multiple
reduction dimensions or even to involve non-reduction
dimensions, as long as those non-reduction dimensions
also appear in outer schedule band members.
Since no skewing is currently allowed, this
is currently also unlikely to happen.

In summary, what the reduction band members look like
is largely irrelevant.  What is important is that
the outer band members make sure that only a single
tensor element gets accessed.  It just so happens
that the reduction band members are very likely
to have a particular form in that case.

Change the reduction detection mechanism to use
the polyhedral reduction representation.
Note that if a tensor element is involved in multiple
separate reductions (again, unlikely in the current
code base), then the detection will not only ensure
that a single tensor element is involved, but more
specifically, a single reduction on that tensor element.

This new mechanism is also much simpler.

Author: Sven Verdoolaege

tc/core/polyhedral/cuda/mapped_scop.cc

@@ -268,11 +268,11 @@ bool MappedScop::detectReductions(detail::ScheduleTree* tree) {
   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 = nCoincident;
-  auto member = band->mupa_.get_union_pw_aff(reductionDim);
-  if (!isReductionMember(member, updates, scop())) {
+  // The outer (coincident) members, together with the prefix schedule,
+  // need to determine a single reduction.
+  auto prefix = prefixScheduleMupa(schedule(), tree);
+  prefix = prefix.range_product(band->memberRange(0, nCoincident));
+  if (!isSingleReductionWithin(updates, prefix, scop())) {
     return false;
   }
   // Order the other statements (if any) before the update statements

tc/core/polyhedral/cuda/mapped_scop.h

@@ -154,6 +154,12 @@ class MappedScop {
   // Look for innermost reduction bands.
   // Store them in reductionBandUpdates_.
   // Return true if any were found.
+  // A band is considered to be a reduction band if it only involves
+  // instances of a single reduction update statement (modulo other
+  // statements that can be moved out of the way) and if the outer
+  // coincident members in the band, together with the prefix schedule,
+  // determine individual reductions.  In particular, each instance
+  // of this combined outer schedule only writes to a single tensor element.
   bool detectReductions(detail::ScheduleTree* band);
   // Does separateReduction need to be called on this node?
   bool needReductionSeparation(const detail::ScheduleTree* st);

tc/core/polyhedral/reduction_matcher.cc

@@ -69,52 +69,6 @@ bool isReductionUpdateId(
   return false;
 }
 
-/*
- * Does "aff" only involve the specified input dimension (and not
- * any other input dimensions).
- */
-bool pwAffInvolvesOnlyInputDim(isl::pw_aff pa, int redDimIdx) {
-  auto space = pa.get_space();
-
-  if (!pa.involves_dims(isl::dim_type::in, redDimIdx, 1)) {
-    return false;
-  }
-
-  if (pa.involves_dims(isl::dim_type::in, 0, redDimIdx) ||
-      pa.involves_dims(
-          isl::dim_type::in,
-          redDimIdx + 1,
-          space.dim(isl::dim_type::in) - redDimIdx - 1)) {
-    return false;
-  }
-
-  return true;
-}
-
-// Does pa have the form S(...) -> [(K*r)] where S is either a reduction init
-// or update statement and r is a known reduction loop in Scop?
-//
-// FIXME: now, K can be any value, including nested integer divisions, to
-// support detection after tiling; tighten this.
-bool isAlmostIdentityReduction(isl::pw_aff pa, const Scop& scop) {
-  auto space = pa.get_space();
-  if (!space.has_tuple_id(isl::dim_type::in)) {
-    return false;
-  }
-  auto stmtId = space.get_tuple_id(isl::dim_type::in);
-  std::vector<size_t> reductionDims;
-  if (!isReductionUpdateId(stmtId, scop, reductionDims)) {
-    return false;
-  }
-
-  for (auto redDimIdx : reductionDims) {
-    if (pwAffInvolvesOnlyInputDim(pa, redDimIdx)) {
-      return true;
-    }
-  }
-  return false;
-}
-
 } // namespace
 
 isl::union_set reductionUpdates(isl::union_set domain, const Scop& scop) {
@@ -129,14 +83,15 @@ isl::union_set reductionUpdates(isl::union_set domain, const Scop& scop) {
   return update;
 }
 
-bool isReductionMember(
-    isl::union_pw_aff member,
+bool isSingleReductionWithin(
     isl::union_set domain,
+    isl::multi_union_pw_aff prefix,
     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);
-  });
+  auto reductions = scop.body.reductions;
+  reductions = reductions.intersect_domain(domain);
+  auto prefixMap = isl::union_map::from(prefix);
+  auto prefixToReduction = reductions.apply_domain(prefixMap);
+  return prefixToReduction.is_single_valued();
 }
 
 } // namespace polyhedral

tc/core/polyhedral/schedule_tree_matcher.h

@@ -29,11 +29,12 @@ namespace polyhedral {
 // 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"?
-bool isReductionMember(
-    isl::union_pw_aff member,
+// Does "prefix" partition "domain" into individual reductions?
+// In particular, do the elements of "domain" access a single tensor
+// element within "prefix"?
+bool isSingleReductionWithin(
     isl::union_set domain,
+    isl::multi_union_pw_aff prefix,
     const Scop& scop);
 
 } // namespace polyhedral