extract may/must writes from Halide IR

Previous commits introduced may/must writes in Scop and dependence
analysis.  Extract those from Halide IR.  Change extractAccess to return
a flag indicating whether the affine access relation is constructed is
exact or not.  Exact relations correspond to must writes since we
statically know which tensor elements are written.  Inexact relations
overapproximate non-affine accesses and should be treated as may writes,
assuming the tensor elements are not necessarily written.

Author: ftynse

tc/core/halide2isl.cc

@@ -17,6 +17,7 @@
 
 #include <algorithm>
 #include <numeric>
+#include <tuple>
 #include <unordered_set>
 
 #include "tc/core/constants.h"
@@ -238,7 +239,20 @@ isl::set makeParamContext(isl::ctx ctx, const SymbolTable& symbolTable) {
   return context;
 }
 
-isl::map extractAccess(
+// Extract a tagged affine access relation from Halide IR.
+// The relation is tagged with a unique identifier, i.e. it lives in the space
+//   [D[...] -> __tc_ref_#[]] -> A[]
+// where # is a unique sequential number, D is the statement identifier
+// extracted from "domain" and A is the tensor identifier constructed from
+// "tensor".  "accesses" map is updated to keep track of the Halide IR nodes in
+// which a particular reference # appeared.
+// Returns the access relation and a flag indicating whether this relation is
+// exact or not.  The relation is overapproximated (that is, not exact) if it
+// represents a non-affine access, for example, an access with indirection such
+// as O(Index(i)) = 42.  In such overapproximated access relation, dimensions
+// that correspond to affine subscripts are still exact while those that
+// correspond to non-affine subscripts are not constrained.
+std::pair<isl::map, bool> extractAccess(
     isl::set domain,
     const IRNode* op,
     const std::string& tensor,
@@ -267,6 +281,7 @@ isl::map extractAccess(
   isl::map map =
       isl::map::universe(domainSpace.map_from_domain_and_range(rangeSpace));
 
+  bool exact = true;
   for (size_t i = 0; i < args.size(); i++) {
     // Then add one equality constraint per dimension to encode the
     // point in the allocation actually read/written for each point in
@@ -278,47 +293,64 @@ isl::map extractAccess(
         isl::pw_aff(isl::local_space(rangeSpace), isl::dim_type::set, i);
     // ... equals the coordinate accessed as a function of the domain.
     auto domainPoint = halide2isl::makeIslAffFromExpr(domainSpace, args[i]);
-    if (!domainPoint.is_null()) {
+    if (!domainPoint) {
+      exact = false;
+    } else {
       map = map.intersect(isl::pw_aff(domainPoint).eq_map(rangePoint));
     }
   }
 
-  return map;
+  return std::make_pair(map, exact);
 }
 
-std::pair<isl::union_map, isl::union_map>
+std::tuple<isl::union_map, isl::union_map, isl::union_map>
 extractAccesses(isl::set domain, const Stmt& s, AccessMap* accesses) {
   class FindAccesses : public IRGraphVisitor {
     using IRGraphVisitor::visit;
 
     void visit(const Call* op) override {
       IRGraphVisitor::visit(op);
       if (op->call_type == Call::Halide || op->call_type == Call::Image) {
-        reads = reads.unite(
-            extractAccess(domain, op, op->name, op->args, accesses));
+        // Read relations can be safely overapproximated.
+        isl::map read;
+        std::tie(read, std::ignore) =
+            extractAccess(domain, op, op->name, op->args, accesses);
+        reads = reads.unite(read);
       }
     }
 
     void visit(const Provide* op) override {
       IRGraphVisitor::visit(op);
-      writes =
-          writes.unite(extractAccess(domain, op, op->name, op->args, accesses));
+
+      // If the write access relation is not exact, we consider that any
+      // element _may_ be written by the statement.  If it is exact, then we
+      // can guarantee that all the elements specified by the relation _must_
+      // be written and any previously stored value will be killed.
+      isl::map write;
+      bool exact;
+      std::tie(write, exact) =
+          extractAccess(domain, op, op->name, op->args, accesses);
+      if (exact) {
+        mustWrites = mustWrites.unite(write);
+      }
+      mayWrites = mayWrites.unite(write);
     }
 
     const isl::set& domain;
     AccessMap* accesses;
 
    public:
-    isl::union_map reads, writes;
+    isl::union_map reads, mayWrites, mustWrites;
 
     FindAccesses(const isl::set& domain, AccessMap* accesses)
         : domain(domain),
           accesses(accesses),
           reads(isl::union_map::empty(domain.get_space())),
-          writes(isl::union_map::empty(domain.get_space())) {}
+          mayWrites(isl::union_map::empty(domain.get_space())),
+          mustWrites(isl::union_map::empty(domain.get_space())) {}
   } finder(domain, accesses);
   s.accept(&finder);
-  return {finder.reads, finder.writes};
+  return std::make_tuple(finder.reads, finder.mayWrites, finder.mustWrites);
 }
 
 /*
@@ -343,7 +375,8 @@ isl::schedule makeScheduleTreeHelper(
     isl::set set,
     std::vector<std::string>& outer,
     isl::union_map* reads,
-    isl::union_map* writes,
+    isl::union_map* mayWrites,
+    isl::union_map* mustWrites,
     AccessMap* accesses,
     StatementMap* statements,
     IteratorMap* iterators) {
@@ -389,7 +422,8 @@ isl::schedule makeScheduleTreeHelper(
         set,
         outerNext,
         reads,
-        writes,
+        mayWrites,
+        mustWrites,
         accesses,
         statements,
         iterators);
@@ -422,7 +456,15 @@ isl::schedule makeScheduleTreeHelper(
     std::vector<isl::schedule> schedules;
     for (Stmt s : stmts) {
       schedules.push_back(makeScheduleTreeHelper(
-          s, set, outer, reads, writes, accesses, statements, iterators));
+          s,
+          set,
+          outer,
+          reads,
+          mayWrites,
+          mustWrites,
+          accesses,
+          statements,
+          iterators));
     }
     schedule = schedules[0].sequence(schedules[1]);
 
@@ -437,23 +479,25 @@ isl::schedule makeScheduleTreeHelper(
     isl::set domain = set.set_tuple_id(id);
     schedule = isl::schedule::from_domain(domain);
 
-    isl::union_map newReads, newWrites;
-    std::tie(newReads, newWrites) =
+    isl::union_map newReads, newMayWrites, newMustWrites;
+    std::tie(newReads, newMayWrites, newMustWrites) =
         halide2isl::extractAccesses(domain, op, accesses);
 
     *reads = reads->unite(newReads);
-    *writes = writes->unite(newWrites);
+    *mayWrites = mayWrites->unite(newMayWrites);
+    *mustWrites = mustWrites->unite(newMustWrites);
 
   } else {
     LOG(FATAL) << "Unhandled Halide stmt: " << s;
   }
   return schedule;
-};
+}
 
 ScheduleTreeAndAccesses makeScheduleTree(isl::space paramSpace, const Stmt& s) {
   ScheduleTreeAndAccesses result;
 
-  result.writes = result.reads = isl::union_map::empty(paramSpace);
+  result.mayWrites = result.mustWrites = result.reads =
+      isl::union_map::empty(paramSpace);
 
   // Walk the IR building a schedule tree
   std::vector<std::string> outer;
@@ -462,7 +506,8 @@ ScheduleTreeAndAccesses makeScheduleTree(isl::space paramSpace, const Stmt& s) {
       isl::set::universe(paramSpace),
       outer,
       &result.reads,
-      &result.writes,
+      &result.mayWrites,
+      &result.mustWrites,
       &result.accesses,
       &result.statements,
       &result.iterators);

tc/core/halide2isl.h

@@ -70,7 +70,7 @@ struct ScheduleTreeAndAccesses {
   /// 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;
+  isl::union_map reads, mayWrites, mustWrites;
 
   /// The correspondence between from Call and Provide nodes and the
   /// reference ids in the reads and writes maps.

tc/core/polyhedral/scop.cc

@@ -61,8 +61,8 @@ ScopUPtr Scop::makeScop(
   auto tree = halide2isl::makeScheduleTree(paramSpace, components.stmt);
   scop->scheduleTreeUPtr = std::move(tree.tree);
   scop->reads = tree.reads;
-  scop->mayWrites = tree.writes;
-  scop->mustWrites = isl::union_map::empty(scop->mayWrites.get_space());
+  scop->mayWrites = tree.mayWrites;
+  scop->mustWrites = tree.mustWrites;
   scop->halide.statements = std::move(tree.statements);
   scop->halide.accesses = std::move(tree.accesses);
   scop->halide.reductions = halide2isl::findReductions(components.stmt);

test/test_core.cc

@@ -165,6 +165,10 @@ struct TC2Isl : public ::testing::Test {
     auto scheduleHalide = polyhedral::detail::fromIslSchedule(
         polyhedral::detail::toIslSchedule(scop->scheduleRoot()).reset_user());
   }
+
+  std::unique_ptr<polyhedral::Scop> MakeScop(const std::string& tc) {
+    return polyhedral::Scop::makeScop(isl::with_exceptions::globalIslCtx(), tc);
+  }
 };
 
 TEST_F(TC2Isl, Copy1D) {
@@ -313,6 +317,56 @@ def fun(float(M, N) I) -> (O1, O2, O3) {
   Check(tc, {123, 13});
 }
 
+// FIXME: range inference seems unaware of indirections on the LHS
+TEST_F(TC2Isl, DISABLED_MayWritesOnly) {
+  string tc = R"TC(
+def scatter(int32(N) A, int32(M) B) -> (O) {
+    O(A(i)) = B(i)
+}
+)TC";
+  auto scop = MakeScop(tc);
+  CHECK(scop->mustWrites.is_empty())
+      << "expected empty must-writes for scatter, got\n"
+      << scop->mustWrites;
+  CHECK(!scop->mayWrites.is_empty())
+      << "expected non-empty may-writes for scatter, got\n"
+      << scop->mayWrites;
+}
+
+TEST_F(TC2Isl, AllMustWrites) {
+  string tc = R"TC(
+def gather(int32(N) A, int32(N) B) -> (O) {
+    O(i) = A(B(i)) where i in 0:N
+}
+)TC";
+  auto scop = MakeScop(tc);
+  CHECK(!scop->mustWrites.is_empty())
+      << "expected non-empty must-writes for gather, got\n"
+      << scop->mustWrites;
+  CHECK_EQ(scop->mustWrites, scop->mayWrites);
+}
+
+// FIXME: range inference seems unaware of indirections on the LHS
+TEST_F(TC2Isl, DISABLED_MustWritesSubsetMayWrites) {
+  string tc = R"TC(
+def scatter_gather(int32(N) A, int32(N) B) -> (O1,O2) {
+    O1(i) = A(B(i)) where i in 0:N
+    O2(A(i)) = B(i)
+}
+)TC";
+  auto scop = MakeScop(tc);
+  CHECK(!scop->mustWrites.is_empty()) << "expected non-empty must-writes, got\n"
+                                      << scop->mustWrites;
+  CHECK(!scop->mayWrites.is_empty()) << "expected non-empty may-writes, got\n"
+                                     << scop->mustWrites;
+  CHECK(scop->mustWrites.is_subset(scop->mayWrites))
+      << scop->mustWrites << " is expected to be a subsetset of "
+      << scop->mayWrites;
+  CHECK(!scop->mayWrites.subtract(scop->mustWrites).is_empty())
+      << scop->mustWrites << "is expected to be a strict subset of "
+      << scop->mayWrites;
+}
+
 int main(int argc, char** argv) {
   ::testing::InitGoogleTest(&argc, argv);
   ::gflags::ParseCommandLineFlags(&argc, &argv, true);