[RFC] use templated isl types

Templated isl types require the user to specify the domain and range
universes of isl objects, allowing the compiler to check whether
it makes sense to combine pairs of objects.
This RFC only converts isPromotableToRegistersBelow and some related
functions to illustrate the effect.
The isPromotableToRegistersBelow was already applying operations
correctly, so the code itself did not require any changes.
However, one variable was reused to store different types
of intermediate result and this one had to be split up into
several variables because they now have different types.

Author: Sven Verdoolaege

tc/core/polyhedral/cuda/mapped_scop.cc

@@ -288,7 +288,7 @@ bool MappedScop::detectReductions(ScheduleTree* tree) {
   // a single reduction for now.
   // Support for multiple reductions would require a check
   // that these reductions do not interfere with each other.
-  auto domain = isl::UnionSet<Statement>(band->mupa_.domain());
+  auto domain = band->mupa_.domain();
   auto updates = reductionUpdates(domain, scop());
   if (updates.n_set() != 1) {
     return false;
@@ -560,7 +560,7 @@ Scop::SyncLevel MappedScop::findBestSync(
   auto activePoints2 = activeDomainPointsBelow(stRoot, st2);
 
   // The dependences between the two schedule trees
-  isl::union_map dependences = scop_->dependences;
+  auto dependences = scop_->dependences;
   dependences = dependences.intersect_domain(activePoints1);
   dependences = dependences.intersect_range(activePoints2);
   if (dependences.is_empty()) {

tc/core/polyhedral/cuda/memory_promotion_heuristic.cc

@@ -333,8 +333,8 @@ bool accessSubscriptsAreUnrolledLoops(
       }
     }
 
-    auto space =
-        subdomain.get_space().add_unnamed_tuple_ui<Unrolled>(unrolledDims.size());
+    auto space = subdomain.get_space().add_unnamed_tuple_ui<Unrolled>(
+        unrolledDims.size());
     auto unrolledDimsMupa = isl::multi_union_pw_aff(space, unrolledDims);
 
     // It is possible that no loops are unrolled, in which case
@@ -343,7 +343,7 @@ bool accessSubscriptsAreUnrolledLoops(
     unrolledDimsMupa =
         unrolledDimsMupa.intersect_domain(group.originalAccesses().domain());
 
-    auto accesses = group.originalAccesses();
+    isl::union_map accesses = group.originalAccesses();
     auto schedule = outerSchedule.flat_range_product(unrolledDimsMupa);
     accesses = accesses.apply_domain(isl::union_map::from(schedule));
 
@@ -656,16 +656,15 @@ void promoteToRegistersBelow(MappedScop& mscop, detail::ScheduleTree* scope) {
   auto blockSchedule = mscop.blockMappingSchedule(mscop.schedule());
 
   // Pure affine schedule without (mapping) filters.
-  isl::multi_union_pw_aff partialSchedMupa =
-      partialScheduleMupa<Scope>(root, scope);
+  auto partialSchedMupa = partialScheduleMupa<Scope>(root, scope);
   // Schedule with block mapping filter.
   auto partialSched =
       isl::union_map::from(partialSchedMupa).intersect_domain(blockMapping);
   // The following promotion validity and profitability checks need to be
   // performed with respect to the block mapping, so append the block schedule.
   // If the partial schedule contains it already, it will just end up with
   // identical dimensions without affecting the result of the checks.
-  partialSchedMupa = partialSchedMupa.flat_range_product(blockSchedule);
+  auto partialSchedBlockMupa = partialSchedMupa.range_product(blockSchedule);
 
   for (auto& tensorGroups : groupMap) {
     auto tensorId = tensorGroups.first;
@@ -679,11 +678,11 @@ void promoteToRegistersBelow(MappedScop& mscop, detail::ScheduleTree* scope) {
         continue;
       }
       if (!isPromotableToRegistersBelow(
-              *group, root, scope, partialSchedMupa, threadSchedule)) {
+              *group, root, scope, partialSchedBlockMupa, threadSchedule)) {
         continue;
       }
       // Check reuse within threads.
-      auto schedule = partialSchedMupa.flat_range_product(threadSchedule);
+      auto schedule = partialSchedBlockMupa.range_product(threadSchedule);
       if (!hasReuseWithin(*group, schedule)) {
         continue;
       }

tc/core/polyhedral/memory_promotion.h

@@ -137,8 +137,9 @@ class TensorReferenceGroup {
   // range spaces.
   isl::union_map originalWrites() const;
   isl::union_map originalReads() const;
-  isl::union_map originalAccesses() const {
-    return originalWrites().unite(originalReads());
+  isl::UnionMap<Statement, Tensor> originalAccesses() const {
+    auto accesses = originalWrites().unite(originalReads());
+    return isl::UnionMap<Statement, Tensor>(accesses);
   }
 
   // Rectangular overapproximation of the set of tensor elements accessed below

tc/core/polyhedral/schedule_utils-inl.h

@@ -129,10 +129,10 @@ inline isl::MultiUnionPwAff<Statement, Schedule> partialScheduleMupa(
     const detail::ScheduleTree* tree) {
   using namespace polyhedral::detail;
 
-  isl::multi_union_pw_aff prefix = prefixScheduleMupa<Schedule>(root, tree);
+  auto prefix = prefixScheduleMupa<Schedule>(root, tree);
   auto band = tree->as<detail::ScheduleTreeBand>();
-  auto partial = band ? prefix.flat_range_product(band->mupa_) : prefix;
-  return isl::MultiUnionPwAff<Statement, Schedule>(partial);
+  return band ? prefix.template flat_range_product<Schedule>(band->mupa_)
+              : prefix;
 }
 
 /*
@@ -171,7 +171,7 @@ isl::MultiUnionPwAff<Statement, MappingType> extractDomainToIds(
       continue;
     }
     auto nodeToIds = isl::MultiUnionPwAff<Statement, MappingType>(space, list);
-    auto active = isl::UnionSet<Statement>(activeDomainPoints(root, mapping));
+    auto active = activeDomainPoints(root, mapping);
     TC_CHECK(active.intersect(domainToIds.domain()).is_empty())
         << "conflicting mappings; are the filters in the tree disjoint?";
     nodeToIds = nodeToIds.intersect_domain(active);

tc/core/polyhedral/schedule_utils.cc

@@ -79,12 +79,12 @@ isl::UnionSet<Statement> collectDomain(
       domain = domain.unite(parentSchedule.range().apply(extension));
     }
   }
-  return domain;
+  return isl::UnionSet<Statement>(domain);
 }
 
 // Get the set of domain elements that are active below
 // the given branch of nodes.
-isl::union_set activeDomainPointsHelper(
+isl::UnionSet<Statement> activeDomainPointsHelper(
     const ScheduleTree* root,
     const vector<const ScheduleTree*>& nodes) {
   return collectDomain(root, nodes, &applyFilter);
@@ -101,16 +101,15 @@ isl::UnionSet<Statement> prefixMappingFilter(
 isl::UnionSet<Statement> activeDomainPoints(
     const ScheduleTree* root,
     const ScheduleTree* node) {
-  return isl::UnionSet<Statement>(
-      activeDomainPointsHelper(root, node->ancestors(root)));
+  return activeDomainPointsHelper(root, node->ancestors(root));
 }
 
 isl::UnionSet<Statement> activeDomainPointsBelow(
     const ScheduleTree* root,
     const ScheduleTree* node) {
   auto ancestors = node->ancestors(root);
   ancestors.emplace_back(node);
-  return isl::UnionSet<Statement>(activeDomainPointsHelper(root, ancestors));
+  return activeDomainPointsHelper(root, ancestors);
 }
 
 vector<ScheduleTree*> collectScheduleTreesPath(

tc/core/polyhedral/scop.cc

@@ -85,7 +85,7 @@ ScopUPtr Scop::makeScop(
   return makeScop(ctx, tc2halide::translate(ctx, treeRef, compilerOptions));
 }
 
-isl::union_set& Scop::domainRef() {
+isl::UnionSet<Statement>& Scop::domainRef() {
   auto dom = scheduleRoot()->as<ScheduleTreeDomain>();
   TC_CHECK(dom) << "root is not a domain in: " << *scheduleRoot();
   // TODO: activate this when the invariant has a chance of working (i.e. we
@@ -99,7 +99,7 @@ isl::union_set& Scop::domainRef() {
 }
 
 const isl::UnionSet<Statement> Scop::domain() const {
-  return isl::UnionSet<Statement>(const_cast<Scop*>(this)->domainRef());
+  return const_cast<Scop*>(this)->domainRef();
 }
 
 std::ostream& operator<<(std::ostream& os, const Scop& s) {

tc/core/polyhedral/scop.h

@@ -509,7 +509,7 @@ struct Scop {
   // By analogy with generalized functions, the domain is the "support" part
   // of the ScheduleTree "function".
  private:
-  isl::union_set& domainRef();
+  isl::UnionSet<Statement>& domainRef();
 
  public:
   const isl::UnionSet<Statement> domain() const;