Merge pull request #510 from facebookresearch/pr/unbind

makeIslAffFromExpr: only accept parametric expressions

Author: skimo-openhub

tc/core/halide2isl.cc

@@ -174,16 +174,9 @@ std::vector<isl::aff> makeIslAffBoundsFromExpr(
   const Max* maxOp = e.as<Max>();
 
   if (const Variable* op = e.as<Variable>()) {
-    isl::local_space ls = isl::local_space(space);
-    int pos = space.find_dim_by_name(isl::dim_type::param, op->name);
-    if (pos >= 0) {
-      return {isl::aff(ls, isl::dim_type::param, pos)};
-    } else {
-      // FIXME: thou shalt not rely upon set dimension names
-      pos = space.find_dim_by_name(isl::dim_type::set, op->name);
-      if (pos >= 0) {
-        return {isl::aff(ls, isl::dim_type::set, pos)};
-      }
+    isl::id id(space.get_ctx(), op->name);
+    if (space.has_param(id)) {
+      return {isl::aff::param_on_domain_space(space, id)};
     }
     LOG(FATAL) << "Variable not found in isl::space: " << space << ": " << op
                << ": " << op->name << '\n';
@@ -248,32 +241,28 @@ isl::set makeParamContext(isl::ctx ctx, const ParameterVector& params) {
   return context;
 }
 
+namespace {
+
 isl::map extractAccess(
-    isl::set domain,
+    const IterationDomain& domain,
     const IRNode* op,
     const std::string& tensor,
     const std::vector<Expr>& args,
     AccessMap* accesses) {
   // Make an isl::map representing this access. It maps from the iteration space
   // to the tensor's storage space, using the coordinates accessed.
+  // First construct a set describing the accessed element
+  // in terms of the parameters (including those corresponding
+  // to the outer loop iterators) and then convert this set
+  // into a map in terms of the iteration domain.
 
-  isl::space domainSpace = domain.get_space();
-  isl::space paramSpace = domainSpace.params();
+  isl::space paramSpace = domain.paramSpace;
   isl::id tensorID(paramSpace.get_ctx(), tensor);
-  auto rangeSpace = paramSpace.named_set_from_params_id(tensorID, args.size());
+  auto tensorSpace = paramSpace.named_set_from_params_id(tensorID, args.size());
 
-  // Add a tag to the domain space so that we can maintain a mapping
-  // between each access in the IR and the reads/writes maps.
-  std::string tag = "__tc_ref_" + std::to_string(accesses->size());
-  isl::id tagID(domain.get_ctx(), tag);
-  accesses->emplace(op, tagID);
-  isl::space tagSpace = paramSpace.named_set_from_params_id(tagID, 0);
-  domainSpace = domainSpace.product(tagSpace);
-
-  // Start with a totally unconstrained relation - every point in
-  // the iteration domain could write to every point in the allocation.
-  isl::map map =
-      isl::map::universe(domainSpace.map_from_domain_and_range(rangeSpace));
+  // Start with a totally unconstrained set - every point in
+  // the allocation could be accessed.
+  isl::set access = isl::set::universe(tensorSpace);
 
   for (size_t i = 0; i < args.size(); i++) {
     // Then add one equality constraint per dimension to encode the
@@ -283,19 +272,34 @@ isl::map extractAccess(
 
     // The coordinate written to in the range ...
     auto rangePoint =
-        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]);
+        isl::pw_aff(isl::local_space(tensorSpace), isl::dim_type::set, i);
+    // ... equals the coordinate accessed as a function of the parameters.
+    auto domainPoint = halide2isl::makeIslAffFromExpr(tensorSpace, args[i]);
     if (!domainPoint.is_null()) {
-      map = map.intersect(isl::pw_aff(domainPoint).eq_map(rangePoint));
+      access = access.intersect(isl::pw_aff(domainPoint).eq_set(rangePoint));
     }
   }
 
+  // Now convert the set into a relation with respect to the iteration domain.
+  auto map = access.unbind_params_insert_domain(domain.tuple);
+
+  // Add a tag to the domain space so that we can maintain a mapping
+  // between each access in the IR and the reads/writes maps.
+  std::string tag = "__tc_ref_" + std::to_string(accesses->size());
+  isl::id tagID(domain.paramSpace.get_ctx(), tag);
+  accesses->emplace(op, tagID);
+  isl::space domainSpace = map.get_space().domain();
+  isl::space tagSpace = domainSpace.params().named_set_from_params_id(tagID, 0);
+  domainSpace = domainSpace.product(tagSpace).unwrap();
+  map = map.preimage_domain(isl::multi_aff::domain_map(domainSpace));
+
   return map;
 }
 
-std::pair<isl::union_map, isl::union_map>
-extractAccesses(isl::set domain, const Stmt& s, AccessMap* accesses) {
+std::pair<isl::union_map, isl::union_map> extractAccesses(
+    const IterationDomain& domain,
+    const Stmt& s,
+    AccessMap* accesses) {
   class FindAccesses : public IRGraphVisitor {
     using IRGraphVisitor::visit;
 
@@ -313,28 +317,46 @@ extractAccesses(isl::set domain, const Stmt& s, AccessMap* accesses) {
           writes.unite(extractAccess(domain, op, op->name, op->args, accesses));
     }
 
-    const isl::set& domain;
+    const IterationDomain& domain;
     AccessMap* accesses;
 
    public:
     isl::union_map reads, writes;
 
-    FindAccesses(const isl::set& domain, AccessMap* accesses)
+    FindAccesses(const IterationDomain& domain, AccessMap* accesses)
         : domain(domain),
           accesses(accesses),
-          reads(isl::union_map::empty(domain.get_space())),
-          writes(isl::union_map::empty(domain.get_space())) {}
+          reads(isl::union_map::empty(domain.tuple.get_space())),
+          writes(isl::union_map::empty(domain.tuple.get_space())) {}
   } finder(domain, accesses);
   s.accept(&finder);
   return {finder.reads, finder.writes};
 }
 
+/*
+ * Take a parametric expression "f" and convert it into an expression
+ * on the iteration domains in "domain" by reinterpreting the parameters
+ * as set dimensions according to the corresponding tuples in "map".
+ */
+isl::union_pw_aff
+onDomains(isl::aff f, isl::union_set domain, const IterationDomainMap& map) {
+  auto upa = isl::union_pw_aff::empty(domain.get_space());
+  for (auto set : domain.get_set_list()) {
+    auto tuple = map.at(set.get_tuple_id()).tuple;
+    auto onSet = isl::union_pw_aff(f.unbind_params_insert_domain(tuple));
+    upa = upa.union_add(onSet);
+  }
+  return upa;
+}
+
+} // namespace
+
 /*
  * Helper function for extracting a schedule from a Halide Stmt,
  * recursively descending over the Stmt.
  * "s" is the current position in the recursive descent.
  * "set" describes the bounds on the outer loop iterators.
- * "outer" contains the names of the outer loop iterators
+ * "outer" contains the identifiers of the outer loop iterators
  * from outermost to innermost.
  * Return the schedule corresponding to the subtree at "s".
  *
@@ -343,81 +365,58 @@ 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,
     isl::set set,
-    std::vector<std::string>& outer,
+    isl::id_list outer,
     isl::union_map* reads,
     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
-    int thisLoopIdx = set.dim(isl::dim_type::set);
-    set = set.add_dims(isl::dim_type::set, 1);
-
-    // Make an id for this loop var. For set dimensions this is
-    // really just for pretty-printing.
+    // Make an id for this loop var.  It starts out as a parameter.
     isl::id id(set.get_ctx(), op->name);
-    set = set.set_dim_id(isl::dim_type::set, thisLoopIdx, id);
+    auto space = set.get_space().add_param(id);
 
-    // Construct a variable (affine function) that indexes the new dimension of
-    // this space.
-    isl::aff loopVar(
-        isl::local_space(set.get_space()), isl::dim_type::set, thisLoopIdx);
+    // Construct a variable (affine function) that references
+    // the new parameter.
+    auto loopVar = isl::aff::param_on_domain_space(space, id);
 
     // Then we add our new loop bound constraints.
-    auto lbs = halide2isl::makeIslAffBoundsFromExpr(
-        set.get_space(), op->min, false, true);
+    auto lbs =
+        halide2isl::makeIslAffBoundsFromExpr(space, op->min, false, true);
     TC_CHECK_GT(lbs.size(), 0u)
         << "could not obtain polyhedral lower bounds from " << op->min;
     for (auto lb : lbs) {
       set = set.intersect(loopVar.ge_set(lb));
     }
 
     Expr max = simplify(op->min + op->extent - 1);
-    auto ubs =
-        halide2isl::makeIslAffBoundsFromExpr(set.get_space(), max, true, false);
+    auto ubs = halide2isl::makeIslAffBoundsFromExpr(space, max, true, false);
     TC_CHECK_GT(ubs.size(), 0u)
         << "could not obtain polyhedral upper bounds from " << max;
     for (auto ub : ubs) {
       set = set.intersect(ub.ge_set(loopVar));
     }
 
     // Recursively descend.
-    auto outerNext = outer;
-    outerNext.push_back(op->name);
+    auto outerNext = outer.add(isl::id(set.get_ctx(), 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
-    // this loop. For each statement in the children we want the
-    // function that maps everything in its space to this
-    // dimension. The spaces may be different, but they'll all have
-    // this loop var at the same index.
-    isl::multi_union_pw_aff mupa;
-    body.get_domain().foreach_set([&](isl::set s) {
-      isl::aff newLoopVar(
-          isl::local_space(s.get_space()), isl::dim_type::set, thisLoopIdx);
-      if (mupa) {
-        mupa = mupa.union_add(isl::union_pw_aff(isl::pw_aff(newLoopVar)));
-      } else {
-        mupa = isl::union_pw_aff(isl::pw_aff(newLoopVar));
-      }
-    });
+    // 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 mupa = isl::multi_union_pw_aff(onDomains(aff, domain, *domains));
 
     schedule = body.insert_partial_schedule(mupa);
   } else if (auto op = s.as<Halide::Internal::Block>()) {
@@ -430,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, iterators));
+          stmt, set, outer, reads, writes, accesses, statements, domains));
     }
     schedule = schedules[0].sequence(schedules[1]);
 
@@ -441,13 +440,18 @@ 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);
-    isl::set domain = set.set_tuple_id(id);
+    auto tupleSpace = isl::space(set.get_ctx(), 0);
+    tupleSpace = tupleSpace.named_set_from_params_id(id, outer.n());
+    IterationDomain iterationDomain;
+    iterationDomain.paramSpace = set.get_space();
+    iterationDomain.tuple = isl::multi_id(tupleSpace, outer);
+    domains->emplace(id, iterationDomain);
+    auto domain = set.unbind_params(iterationDomain.tuple);
     schedule = isl::schedule::from_domain(domain);
 
     isl::union_map newReads, newWrites;
     std::tie(newReads, newWrites) =
-        halide2isl::extractAccesses(domain, op, accesses);
+        extractAccesses(iterationDomain, op, accesses);
 
     *reads = reads->unite(newReads);
     *writes = writes->unite(newWrites);
@@ -464,7 +468,7 @@ ScheduleTreeAndAccesses makeScheduleTree(isl::space paramSpace, const Stmt& s) {
   result.writes = result.reads = isl::union_map::empty(paramSpace);
 
   // Walk the IR building a schedule tree
-  std::vector<std::string> outer;
+  isl::id_list outer(paramSpace.get_ctx(), 0);
   auto schedule = makeScheduleTreeHelper(
       s,
       isl::set::universe(paramSpace),
@@ -473,7 +477,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

@@ -53,11 +53,25 @@ isl::aff makeIslAffFromInt(isl::space space, int64_t i);
 
 // Make an affine function over a space from a Halide Expr. Returns a
 // null isl::aff if the expression is not affine. Fails if Variable
-// does not correspond to a parameter or set dimension of the space.
+// does not correspond to a parameter of the space.
+// Note that the input space can be either a parameter space or
+// a set space, but the expression can only reference
+// the parameters in 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 {
+  // All parameters active at the point where the iteration domain
+  // was created, including those corresponding to outer loop iterators.
+  isl::space paramSpace;
+  // The identifier tuple corresponding to the iteration domain.
+  // The identifiers in the tuple are the outer loop iterators,
+  // from outermost to innermost.
+  isl::multi_id tuple;
+};
+
+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 +95,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,19 +637,18 @@ 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 tuple = scop.halide.domains.at(stmtId).tuple;
   auto& stmtIteratorMap = iteratorMaps[stmtId];
-  for (size_t i = 0; i < iterators.size(); ++i) {
+  for (int i = 0; i < tuple.size(); ++i) {
     auto expr = build.expr_from(iteratorMap.get_pw_aff(i));
-    stmtIteratorMap.emplace(iterators[i], expr);
+    stmtIteratorMap.emplace(tuple.get_id(i).get_name(), expr);
   }
   auto& subscripts = stmtSubscripts[stmtId];
   auto provide =
       scop.halide.statements.at(stmtId).as<Halide::Internal::Provide>();
   for (auto e : provide->args) {
     const auto& map = iteratorMap;
-    auto space = map.get_space().params();
-    auto aff = scop.makeIslAffFromStmtExpr(stmtId, space, e);
+    auto aff = scop.makeIslAffFromStmtExpr(stmtId, e);
     auto pulled = isl::pw_aff(aff).pullback(map);
     TC_CHECK_EQ(pulled.n_piece(), 1);
     subscripts.push_back(build.expr_from(pulled));

tc/core/polyhedral/cuda/codegen.h

@@ -136,8 +136,7 @@ struct CodegenStatementContext : CodegenContext {
   // of the variables does not correspond to a parameter or
   // an instance identifier of the statement.
   isl::aff makeIslAffFromExpr(const Halide::Expr& e) const {
-    auto space = iteratorMap().get_space().params();
-    return scop().makeIslAffFromStmtExpr(statementId(), space, e);
+    return scop().makeIslAffFromStmtExpr(statementId(), e);
   }
 
   isl::id astNodeId;