store data extracted from statement bodies in Body structure

This data is collected by makeScheduleTree and then copied
into and between Scops.  It is easier to do this collection
on a single object rather than on multiple objects.
For now, the store information consists of two objects,
the read and write accesses, but this will be extended
with a polyhedral representation of reductions.
The fields of Body also support some common operations,
supporting the fact that they belong together.

Author: Sven Verdoolaege

tc/core/CMakeLists.txt

@@ -21,6 +21,7 @@ add_library(
   halide2isl.cc
   halide_utils.cc
 
+  polyhedral/body.cc
   polyhedral/codegen.cc
   polyhedral/memory_promotion.cc
   polyhedral/reduction_matcher.cc

tc/core/halide2isl.cc

@@ -20,6 +20,7 @@
 
 #include "tc/core/check.h"
 #include "tc/core/constants.h"
+#include "tc/core/polyhedral/body.h"
 #include "tc/core/polyhedral/schedule_isl_conversion.h"
 #include "tc/core/polyhedral/schedule_transforms.h"
 #include "tc/core/polyhedral/schedule_tree.h"
@@ -360,7 +361,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".
  *
- * "reads" and "writes" collect the accesses found along the way.
+ * "body" collects the accesses 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
@@ -372,8 +373,7 @@ isl::schedule makeScheduleTreeHelper(
     const Stmt& s,
     isl::set set,
     isl::id_list outer,
-    isl::union_map* reads,
-    isl::union_map* writes,
+    Body* body,
     AccessMap* accesses,
     StatementMap* statements,
     IterationDomainMap* domains) {
@@ -406,19 +406,19 @@ isl::schedule makeScheduleTreeHelper(
 
     // Recursively descend.
     auto outerNext = outer.add(isl::id(set.get_ctx(), op->name));
-    auto body = makeScheduleTreeHelper(
-        op->body, set, outerNext, reads, writes, accesses, statements, domains);
+    auto bodySchedule = makeScheduleTreeHelper(
+        op->body, set, outerNext, body, accesses, statements, domains);
 
     // Create an affine function that defines an ordering for all
     // the statements in the body of this loop over the values of
     // this loop.  Start from a parametric expression equal
     // to the current loop iterator and then convert it to
     // a function on the statements in the domain of the body schedule.
     auto aff = isl::aff::param_on_domain_space(space, id);
-    auto domain = body.get_domain();
+    auto domain = bodySchedule.get_domain();
     auto mupa = isl::multi_union_pw_aff(onDomains(aff, domain, *domains));
 
-    schedule = body.insert_partial_schedule(mupa);
+    schedule = bodySchedule.insert_partial_schedule(mupa);
   } else if (auto op = s.as<Halide::Internal::Block>()) {
     std::vector<Stmt> stmts;
     stmts.push_back(op->first);
@@ -429,7 +429,7 @@ isl::schedule makeScheduleTreeHelper(
     std::vector<isl::schedule> schedules;
     for (Stmt stmt : stmts) {
       schedules.push_back(makeScheduleTreeHelper(
-          stmt, set, outer, reads, writes, accesses, statements, domains));
+          stmt, set, outer, body, accesses, statements, domains));
     }
     schedule = schedules[0].sequence(schedules[1]);
 
@@ -453,8 +453,8 @@ isl::schedule makeScheduleTreeHelper(
     std::tie(newReads, newWrites) =
         extractAccesses(iterationDomain, op, accesses);
 
-    *reads = reads->unite(newReads);
-    *writes = writes->unite(newWrites);
+    body->reads = body->reads.unite(newReads);
+    body->writes = body->writes.unite(newWrites);
 
   } else {
     LOG(FATAL) << "Unhandled Halide stmt: " << s;
@@ -465,20 +465,20 @@ isl::schedule makeScheduleTreeHelper(
 ScheduleTreeAndAccesses makeScheduleTree(isl::space paramSpace, const Stmt& s) {
   ScheduleTreeAndAccesses result;
 
-  result.writes = result.reads = isl::union_map::empty(paramSpace);
+  Body body(paramSpace);
 
   // Walk the IR building a schedule tree
   isl::id_list outer(paramSpace.get_ctx(), 0);
   auto schedule = makeScheduleTreeHelper(
       s,
       isl::set::universe(paramSpace),
       outer,
-      &result.reads,
-      &result.writes,
+      &body,
       &result.accesses,
       &result.statements,
       &result.domains);
 
+  result.body = body;
   result.tree = fromIslSchedule(schedule);
 
   return result;

tc/core/halide2isl.h

@@ -22,6 +22,7 @@
 
 #include <Halide.h>
 
+#include "tc/core/polyhedral/body.h"
 #include "tc/core/polyhedral/schedule_tree.h"
 #include "tc/core/tc2halide.h"
 #include "tc/external/isl.h"
@@ -82,10 +83,8 @@ struct ScheduleTreeAndAccesses {
   /// for each leaf node is captured below.
   tc::polyhedral::ScheduleTreeUPtr tree;
 
-  /// 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;
+  /// Information extracted from the bodies of the statements.
+  tc::polyhedral::Body body;
 
   /// The correspondence between from Call and Provide nodes and the
   /// reference ids in the reads and writes maps.

tc/core/polyhedral/body.cc

@@ -0,0 +1,34 @@
+/**
+ * Copyright (c) 2018, Facebook, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "tc/core/polyhedral/body.h"
+
+#include <iostream>
+
+#include "tc/external/isl.h"
+
+namespace tc {
+namespace polyhedral {
+
+std::ostream& operator<<(std::ostream& os, const Body& body) {
+  os << "reads: " << body.reads << "\n";
+  os << "writes: " << body.writes << "\n";
+
+  return os;
+}
+
+} // namespace polyhedral
+} // namespace tc

tc/core/polyhedral/body.h

@@ -0,0 +1,47 @@
+/**
+ * Copyright (c) 2018, Facebook, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+#include <iostream>
+
+#include "tc/external/isl.h"
+
+namespace tc {
+namespace polyhedral {
+
+// Information about the bodies of the polyhedral statements.
+struct Body {
+  Body() = default;
+  Body(isl::space paramSpace) {
+    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);
+  }
+
+  // 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;
+};
+
+std::ostream& operator<<(std::ostream& os, const Body& body);
+
+} // namespace polyhedral
+} // namespace tc

tc/core/polyhedral/cuda/codegen.cc

@@ -23,6 +23,7 @@
 #include "tc/core/check.h"
 #include "tc/core/cuda/cuda_libraries.h"
 #include "tc/core/flags.h"
+#include "tc/core/polyhedral/body.h"
 #include "tc/core/polyhedral/codegen.h"
 #include "tc/core/polyhedral/cuda/codegen.h"
 #include "tc/core/polyhedral/cuda/mapping_types.h"
@@ -769,8 +770,8 @@ std::unordered_set<isl::id, isl::IslIdIslHash> gatherReadOnlySet(
 
   const auto& scop = mscop.scop();
 
-  auto read = scop.reads.universe().range();
-  auto written = scop.writes.universe().range();
+  auto read = scop.body.reads.universe().range();
+  auto written = scop.body.writes.universe().range();
   auto readOnly = read.subtract(written);
   for (auto s : readOnly.get_set_list()) {
     readOnlySet.emplace(s.get_tuple_id());

tc/core/polyhedral/cuda/memory_promotion_heuristic.cc

@@ -523,7 +523,7 @@ void promoteToSharedGreedy(
     auto partialSched = partialSchedule(root, bandNode);
 
     auto groupMap = TensorReferenceGroup::accessedWithin(
-        partialSched.intersect_domain(activePoints), scop.reads, scop.writes);
+        partialSched.intersect_domain(activePoints), scop.body);
     // Pure affine schedule without (mapping) filters.
     auto partialSchedMupa = partialScheduleMupa(root, bandNode);
 
@@ -650,7 +650,7 @@ void promoteToRegistersBelow(MappedScop& mscop, detail::ScheduleTree* scope) {
       collectMappingsTo<mapping::ThreadId>(scop).intersect(blockMapping);
   auto schedule = partialSchedule(scop.scheduleRoot(), scope);
   auto groupMap = TensorReferenceGroup::accessedWithin(
-      schedule.intersect_domain(mapping), scop.reads, scop.writes);
+      schedule.intersect_domain(mapping), scop.body);
 
   auto threadSchedule = mscop.threadMappingSchedule(mscop.schedule());
   auto blockSchedule = mscop.blockMappingSchedule(mscop.schedule());

tc/core/polyhedral/memory_promotion.cc

@@ -21,6 +21,7 @@
 #include <unordered_map>
 
 #include "tc/core/check.h"
+#include "tc/core/polyhedral/body.h"
 #include "tc/core/polyhedral/exceptions.h"
 #include "tc/core/polyhedral/schedule_tree.h"
 #include "tc/core/polyhedral/scop.h"
@@ -341,15 +342,14 @@ void addSingletonReferenceGroups(
 // references.
 TensorGroups TensorReferenceGroup::accessedWithin(
     isl::union_map outerSchedule,
-    isl::union_map reads,
-    isl::union_map writes) {
+    const Body& body) {
   TensorGroups tensorGroups;
   auto domain = outerSchedule.domain();
 
   addSingletonReferenceGroups(
-      tensorGroups, writes, domain, outerSchedule, AccessType::Write);
+      tensorGroups, body.writes, domain, outerSchedule, AccessType::Write);
   addSingletonReferenceGroups(
-      tensorGroups, reads, domain, outerSchedule, AccessType::Read);
+      tensorGroups, body.reads, domain, outerSchedule, AccessType::Read);
 
   // For each tensor, join groups whose footprints overlap and at least one
   // access is a write.  Do not join between tensors because no aliasing.

tc/core/polyhedral/memory_promotion.h

@@ -113,8 +113,7 @@ class TensorReferenceGroup {
  public:
   static TensorGroups accessedWithin(
       isl::union_map outerSchedule,
-      isl::union_map reads,
-      isl::union_map writes);
+      const Body& body);
 
   bool isReadOnly() const;
 

tc/core/polyhedral/scop.cc

@@ -25,6 +25,7 @@
 
 #include "tc/core/check.h"
 #include "tc/core/halide2isl.h"
+#include "tc/core/polyhedral/body.h"
 #include "tc/core/polyhedral/functional.h"
 #include "tc/core/polyhedral/memory_promotion.h"
 #include "tc/core/polyhedral/schedule_isl_conversion.h"
@@ -60,8 +61,7 @@ ScopUPtr Scop::makeScop(
 
   auto tree = halide2isl::makeScheduleTree(paramSpace, components.stmt);
   scop->scheduleTreeUPtr = std::move(tree.tree);
-  scop->reads = tree.reads;
-  scop->writes = tree.writes;
+  scop->body = tree.body;
   scop->halide.statements = std::move(tree.statements);
   scop->halide.accesses = std::move(tree.accesses);
   scop->halide.reductions = halide2isl::findReductions(components.stmt);
@@ -97,8 +97,7 @@ const isl::union_set Scop::domain() const {
 
 std::ostream& operator<<(std::ostream& os, const Scop& s) {
   os << "domain: " << s.domain() << "\n";
-  os << "reads: " << s.reads << "\n";
-  os << "writes: " << s.writes << "\n";
+  os << s.body;
   os << "schedule: " << *s.scheduleRoot() << "\n";
   os << "idx: { ";
   for (auto i : s.halide.idx) {
@@ -256,7 +255,7 @@ void Scop::promoteEverythingAt(std::vector<size_t> pos) {
   checkFiltersDisjointStatements(scheduleRoot());
   auto schedule = partialSchedule(root, tree);
 
-  auto groupMap = TensorReferenceGroup::accessedWithin(schedule, reads, writes);
+  auto groupMap = TensorReferenceGroup::accessedWithin(schedule, body);
   for (auto& p : groupMap) {
     for (auto& gr : p.second) {
       promoteGroup(
@@ -353,8 +352,8 @@ isl::schedule_constraints makeScheduleConstraints(
 
 void Scop::computeAllDependences() {
   auto schedule = toIslSchedule(scheduleRoot());
-  auto allReads = reads.domain_factor_domain();
-  auto allWrites = writes.domain_factor_domain();
+  auto allReads = body.reads.domain_factor_domain();
+  auto allWrites = body.writes.domain_factor_domain();
   // RAW
   auto flowDeps = computeDependences(allWrites, allReads, schedule);
   // WAR and WAW