makeScheduleTree: also extract polyhedral representation of reductions

Reductions are currently represented by lists of reduction dimensions
in the domains of the reduction update statements.
The detection of reductions in the schedule tree, based on this information,
may work out in practice, but it is difficult to reason about and
is technically incorrect.
This detection will be replaced by one that is based
on a polyhedral representation.  This commit introduces
this polyhedral representation.

In particular, each individual reduction (a set of updates
to a given tensor element) is represented by an element
in a reduction space, along with a mapping from the
update statement instances that contribute to that reduction.

The obvious choice for such a reduction space would be
the space of the tensor involved in the reduction.
However, there may in theory be multiple reductions
on the same tensor element separated by some other statement,
even though this is currently impossible inside TC.
Still, since the reductions are extracted from Halide,
it's best to only use information available in the Halide
representation and to not assume an identity matching
between reductions and tensor elements.

It is not immediately obvious what would be the best representation
for reductions.  This commit uses an isl::union_map representation,
but an isl::union_pw_multi_aff would also work.

This commit only extracts a polyhedral representation of reductions,
temporarily resulting in a duplicate representation of reductions.
After all uses of the AST representation of reductions have been
replaced by uses of the polyhedral representation, the AST representation
will be removed.

Author: Sven Verdoolaege

tc/core/constants.h

@@ -22,6 +22,7 @@ namespace polyhedral {
 // General constants to avoid hardcoding
 //
 constexpr auto kStatementLabel = "S_";
+constexpr auto kReductionLabel = "R_";
 
 constexpr auto kAstNodeIdPrefix = "__node_";
 

tc/core/halide2isl.cc

@@ -342,6 +342,82 @@ bool isReductionUpdate(const Provide* op) {
   }
 }
 
+/* Construct a multi-dimensional affine function mapping
+ * the given iteration domain
+ * to the outer loop iterators that do not appear in "skip".
+ * "id" is used as the identifier of the target space.
+ * For each of these outer loop iterators, an affine function
+ * is first constructed in terms of the parameter space
+ * active at the point where the iteration domain was created and
+ * then converted into an expression on that iteration domain
+ * by reinterpreting the parameters as input dimensions.
+ */
+static isl::multi_aff mapToOther(
+    const IterationDomain& iterationDomain,
+    std::unordered_set<std::string> skip,
+    isl::id id) {
+  auto ctx = iterationDomain.tuple.get_ctx();
+  auto list = isl::aff_list(ctx, 0);
+  for (auto id : iterationDomain.tuple.get_id_list()) {
+    if (skip.count(id.get_name()) == 1) {
+      continue;
+    }
+    auto aff = isl::aff::param_on_domain_space(iterationDomain.paramSpace, id);
+    aff = aff.unbind_params_insert_domain(iterationDomain.tuple);
+    list = list.add(aff);
+  }
+  auto domainSpace = iterationDomain.tuple.get_space();
+  auto space = domainSpace.params().named_set_from_params_id(id, list.size());
+  space = domainSpace.product(space).unwrap();
+  return isl::multi_aff(space, list);
+}
+
+/*
+ * If "op" performs a reduction, then return a mapping from
+ * the statement instances to the individual reductions.
+ * Otherwise, return an empty isl::union_map.
+ *
+ * "op" is considered to be a reduction if it has been marked
+ * as performing a reduction and if more than one statement instance
+ * is involved in the individual reductions.
+ *
+ * The space of the reduction has a name of the form R_<op->name>_<index>.
+ * Each reduction is indexed by the outer loop variables
+ * that are not marked as reduction variables.
+ * Since the loop variables that iterate over output tensor elements
+ * are never marked as reduction variables, this means in particular
+ * that all statement instances that belong to the same reduction
+ * write to the same tensor element.
+ */
+isl::union_map extractReduction(
+    const IterationDomain& iterationDomain,
+    const Provide* op,
+    size_t index) {
+  class FindReductionVars : public IRVisitor {
+    void visit(const Variable* op) {
+      if (op->reduction_domain.defined()) {
+        reductionVars.insert(op->name);
+      }
+    }
+
+   public:
+    // The variables that are known to be reduction variables.
+    std::unordered_set<std::string> reductionVars;
+  } finder;
+
+  if (!isReductionUpdate(op)) {
+    return isl::union_map::empty(iterationDomain.tuple.get_space().params());
+  }
+  op->accept(&finder);
+  if (finder.reductionVars.size() == 0) {
+    return isl::union_map::empty(iterationDomain.tuple.get_space().params());
+  }
+  auto ctx = iterationDomain.tuple.get_ctx();
+  isl::id id(ctx, kReductionLabel + op->name + "_" + std::to_string(index));
+  auto reduction = mapToOther(iterationDomain, finder.reductionVars, id);
+  return isl::union_map(isl::map(reduction));
+}
+
 /*
  * Take a parametric expression "f" and convert it into an expression
  * on the iteration domains in "domain" by reinterpreting the parameters
@@ -369,7 +445,7 @@ onDomains(isl::aff f, isl::union_set domain, const IterationDomainMap& map) {
  * from outermost to innermost.
  * Return the schedule corresponding to the subtree at "s".
  *
- * "body" collects the accesses found along the way.
+ * "body" collects the accesses and reductions found along the way.
  * "accesses" collects the mapping from Call (for the reads) and Provide nodes
  * (for the writes) to the corresponding tag in the access relations.
  * "statements" collects the mapping from instance set tuple identifiers
@@ -460,9 +536,13 @@ isl::schedule makeScheduleTreeHelper(
     isl::union_map newReads, newWrites;
     std::tie(newReads, newWrites) =
         extractAccesses(iterationDomain, op, accesses);
+    // A tensor may be involved in multiple reductions.
+    // Use the statement index to differentiate between them.
+    auto newReduction = extractReduction(iterationDomain, op, stmtIndex);
 
     body->reads = body->reads.unite(newReads);
     body->writes = body->writes.unite(newWrites);
+    body->reductions = body->reductions.unite(newReduction);
 
   } else {
     LOG(FATAL) << "Unhandled Halide stmt: " << s;

tc/core/polyhedral/body.cc

@@ -26,6 +26,7 @@ namespace polyhedral {
 std::ostream& operator<<(std::ostream& os, const Body& body) {
   os << "reads: " << body.reads << "\n";
   os << "writes: " << body.writes << "\n";
+  os << "reductions: " << body.reductions << "\n";
 
   return os;
 }

tc/core/polyhedral/body.h

@@ -26,19 +26,54 @@ namespace polyhedral {
 struct Body {
   Body() = default;
   Body(isl::space paramSpace) {
-    writes = reads = isl::union_map::empty(paramSpace);
+    reductions = writes = reads = isl::union_map::empty(paramSpace);
   }
 
   // Specialize to the given context.
   void specialize(isl::set context) {
     reads = reads.intersect_params(context);
     writes = writes.intersect_params(context);
+    reductions = reductions.intersect_params(context);
   }
 
   // Union maps describing the reads and writes done. Uses the ids in
   // the schedule tree to denote the containing Stmt, and tags each
   // access with a unique reference id of the form __tc_ref_N.
   isl::union_map reads, writes;
+
+  // A function on reduction update statement instances that partitions them
+  // into individual reductions, where each reduction consists of
+  // associative updates to the same tensor element.
+  // Since each reduction involves a single tensor element,
+  // the partition of statement instances based
+  // on the reductions forms a refinement of the partition based
+  // on the element modified by the statement.
+  // That is, if W is the write access relation and R is the reduction function,
+  // then, in iscc notation, (W.W^-1) >= (R.R^-1).
+  // In theory, it is possible for the inclusion to be strict, i.e.,
+  // for (W.W^-1) > (R.R^-1) to hold.  For example, for a statement
+  //
+  //	A += T(i)
+  //
+  // with T of size 4, the write access relation is
+  //
+  //	{ S[i] -> A[] : 0 <= i < 4 }
+  //
+  // and the reduction relation could in theory be something like
+  //
+  //	{ S[i] -> R[i] : 0 <= i < 4 }
+  //
+  // or even
+  //
+  //	{ S[i] -> R1[i] : 0 <= i < 2; S[i] -> R2[i - 2] : 3 <= i < 4 }
+  //
+  // In practice, the reduction map is usually equal to
+  // the write access relation on reduction update statements,
+  // with different target spaces.
+  // That is, in the example above, it would just be
+  //
+  //	{ S[i] -> R[] : 0 <= i < 4 }
+  isl::union_map reductions;
 };
 
 std::ostream& operator<<(std::ostream& os, const Body& body);