halide2isl: store outer loop iterators in IterationDomain structure

This will make it easier to change the representation
of the outer loop iterators and to add extra information
related to the iteration domain.

Author: Sven Verdoolaege

tc/core/halide2isl.cc

@@ -347,8 +347,8 @@ extractAccesses(isl::set domain, const Stmt& s, AccessMap* accesses) {
  * (for the writes) to the corresponding tag in the access relations.
  * "statements" collects the mapping from instance set tuple identifiers
  * to the corresponding Provide node.
- * "iterators" collects the mapping from instance set tuple identifiers
- * to the corresponding outer loop iterator names, from outermost to innermost.
+ * "domains" collects the mapping from instance set tuple identifiers
+ * to the corresponding iteration domain information.
  */
 isl::schedule makeScheduleTreeHelper(
     const Stmt& s,
@@ -358,7 +358,7 @@ isl::schedule makeScheduleTreeHelper(
     isl::union_map* writes,
     AccessMap* accesses,
     StatementMap* statements,
-    IteratorMap* iterators) {
+    IterationDomainMap* domains) {
   isl::schedule schedule;
   if (auto op = s.as<For>()) {
     // Add one additional dimension to our set of loop variables
@@ -397,14 +397,7 @@ isl::schedule makeScheduleTreeHelper(
     auto outerNext = outer;
     outerNext.push_back(op->name);
     auto body = makeScheduleTreeHelper(
-        op->body,
-        set,
-        outerNext,
-        reads,
-        writes,
-        accesses,
-        statements,
-        iterators);
+        op->body, set, outerNext, reads, writes, 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
@@ -434,7 +427,7 @@ isl::schedule makeScheduleTreeHelper(
     std::vector<isl::schedule> schedules;
     for (Stmt stmt : stmts) {
       schedules.push_back(makeScheduleTreeHelper(
-          stmt, set, outer, reads, writes, accesses, statements, iterators));
+          stmt, set, outer, reads, writes, accesses, statements, domains));
     }
     schedule = schedules[0].sequence(schedules[1]);
 
@@ -445,7 +438,9 @@ isl::schedule makeScheduleTreeHelper(
     size_t stmtIndex = statements->size();
     isl::id id(set.get_ctx(), kStatementLabel + std::to_string(stmtIndex));
     statements->emplace(id, op);
-    iterators->emplace(id, outer);
+    IterationDomain iterationDomain;
+    iterationDomain.iterators = outer;
+    domains->emplace(id, iterationDomain);
     isl::set domain = set.set_tuple_id(id);
     schedule = isl::schedule::from_domain(domain);
 
@@ -476,7 +471,7 @@ ScheduleTreeAndAccesses makeScheduleTree(isl::space paramSpace, const Stmt& s) {
       &result.writes,
       &result.accesses,
       &result.statements,
-      &result.iterators);
+      &result.domains);
 
   result.tree = fromIslSchedule(schedule);
 

tc/core/halide2isl.h

@@ -56,8 +56,14 @@ isl::aff makeIslAffFromInt(isl::space space, int64_t i);
 // does not correspond to a parameter or set dimension of the space.
 isl::aff makeIslAffFromExpr(isl::space space, const Halide::Expr& e);
 
-typedef std::unordered_map<isl::id, std::vector<std::string>, isl::IslIdIslHash>
-    IteratorMap;
+// Iteration domain information associated to a statement identifier.
+struct IterationDomain {
+  // The outer loop iterators, from outermost to innermost.
+  std::vector<std::string> iterators;
+};
+
+typedef std::unordered_map<isl::id, IterationDomain, isl::IslIdIslHash>
+    IterationDomainMap;
 typedef std::unordered_map<isl::id, Halide::Internal::Stmt, isl::IslIdIslHash>
     StatementMap;
 typedef std::unordered_map<const Halide::Internal::IRNode*, isl::id> AccessMap;
@@ -81,9 +87,9 @@ struct ScheduleTreeAndAccesses {
   /// refered to above.
   StatementMap statements;
 
-  /// The correspondence between statement ids and the outer loop iterators
+  /// The correspondence between statement ids and the iteration domain
   /// of the corresponding leaf Stmt.
-  IteratorMap iterators;
+  IterationDomainMap domains;
 };
 
 /// Make a schedule tree from a Halide Stmt, along with auxiliary data

tc/core/polyhedral/codegen_llvm.cc

@@ -637,7 +637,7 @@ isl::ast_node collectIteratorMaps(
   auto stmtId = expr.get_arg(0).as<isl::ast_expr_id>().get_id();
   TC_CHECK_EQ(0u, iteratorMaps.count(stmtId)) << "entry exists: " << stmtId;
   auto iteratorMap = isl::pw_multi_aff(scheduleMap.reverse());
-  auto iterators = scop.halide.iterators.at(stmtId);
+  auto iterators = scop.halide.domains.at(stmtId).iterators;
   auto& stmtIteratorMap = iteratorMaps[stmtId];
   for (size_t i = 0; i < iterators.size(); ++i) {
     auto expr = build.expr_from(iteratorMap.get_pw_aff(i));

tc/core/polyhedral/scop.cc

@@ -64,7 +64,7 @@ ScopUPtr Scop::makeScop(
   scop->halide.statements = std::move(tree.statements);
   scop->halide.accesses = std::move(tree.accesses);
   scop->halide.reductions = halide2isl::findReductions(components.stmt);
-  scop->halide.iterators = std::move(tree.iterators);
+  scop->halide.domains = std::move(tree.domains);
 
   return scop;
 }
@@ -508,7 +508,7 @@ isl::aff Scop::makeIslAffFromStmtExpr(
     isl::space paramSpace,
     const Halide::Expr& e) const {
   auto ctx = stmtId.get_ctx();
-  auto iterators = halide.iterators.at(stmtId);
+  auto iterators = halide.domains.at(stmtId).iterators;
   auto space = paramSpace.named_set_from_params_id(stmtId, iterators.size());
   // Set the names of the set dimensions of "space" for use
   // by halide2isl::makeIslAffFromExpr.

tc/core/polyhedral/scop.h

@@ -490,7 +490,7 @@ struct Scop {
     std::unordered_map<isl::id, Halide::Internal::Stmt, isl::IslIdIslHash>
         statements;
     std::unordered_map<const Halide::Internal::IRNode*, isl::id> accesses;
-    halide2isl::IteratorMap iterators;
+    halide2isl::IterationDomainMap domains;
   } halide;
 
   // Polyhedral IR