[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/CMakeLists.txt

@@ -32,7 +32,6 @@ add_library(
   polyhedral/schedule_print.cc
   polyhedral/schedule_utils.cc
   polyhedral/scop.cc
-  polyhedral/separation.cc
   polyhedral/unroll.cc
   polyhedral/utils.cc
 )

tc/core/polyhedral/cuda/mapped_scop.cc

@@ -226,7 +226,10 @@ void fixThreadsBelow(MappedScop& mscop, ScheduleTree* tree, size_t begin) {
  * Anything that depends on an update statement is ordered after
  * the update statements.  Anything else is ordered before.
  */
-bool separatedOut(Scop& scop, ScheduleTree* tree, isl::union_set updates) {
+bool separatedOut(
+    Scop& scop,
+    ScheduleTree* tree,
+    isl::UnionSet<Statement> updates) {
   auto domain = activeDomainPoints(scop.scheduleRoot(), tree);
   auto other = domain.subtract(updates);
   if (other.is_empty()) {
@@ -285,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;
@@ -502,54 +505,54 @@ constexpr auto kWarp = "warp";
  * (of size "warpSize") to a warp identifier,
  * based on the thread sizes s_x, s_y up to s_z in "block".
  */
-isl::multi_aff constructThreadToWarp(
+isl::MultiAff<Thread, Warp> constructThreadToWarp(
     isl::ctx ctx,
     const unsigned warpSize,
     const Block& block) {
-  auto space = isl::space(ctx, 0);
+  auto space = isl::Space<>(isl::space(ctx, 0));
   auto id = isl::id(ctx, kBlock);
-  auto blockSpace = space.add_named_tuple_id_ui(id, block.view.size());
-  auto warpSpace = space.add_named_tuple_id_ui(isl::id(ctx, kWarp), 1);
-  auto aff = isl::aff::zero_on_domain(blockSpace);
+  auto blockSpace = space.add_named_tuple_id_ui<Thread>(id, block.view.size());
+  auto warpSpace = space.add_named_tuple_id_ui<Warp>(isl::id(ctx, kWarp), 1);
+  auto aff = isl::AffOn<Thread>::zero_on_domain(blockSpace);
 
   auto nThread = block.view.size();
-  auto identity = isl::multi_aff::identity(blockSpace.map_from_set());
+  auto identity = isl::MultiAff<Thread, Thread>::identity(blockSpace.map_from_set());
   for (int i = nThread - 1; i >= 0; --i) {
     aff = aff.scale(isl::val(ctx, block.view[i]));
     aff = aff.add(identity.get_aff(i));
   }
 
   aff = aff.scale_down(isl::val(ctx, warpSize)).floor();
-  auto mapSpace = blockSpace.product(warpSpace).unwrap();
-  return isl::multi_aff(mapSpace, isl::aff_list(aff));
+  auto mapSpace = blockSpace.product(warpSpace).unwrap<Thread, Warp>();
+  return isl::MultiAff<Thread, Warp>(mapSpace, aff.asAffList());
 }
 } // namespace
 
-isl::multi_union_pw_aff MappedScop::threadMappingSchedule(
+isl::MultiUnionPwAff<Statement, Thread> MappedScop::threadMappingSchedule(
     const ScheduleTree* tree) const {
   std::vector<mapping::MappingId> ids;
   for (size_t i = 0; i < numThreads.view.size(); ++i) {
     ids.emplace_back(mapping::ThreadId::makeId(i));
   }
   auto tupleId = isl::id(tree->ctx_, kBlock);
-  return extractDomainToIds(scop_->scheduleRoot(), tree, ids, tupleId);
+  return extractDomainToIds<Thread>(scop_->scheduleRoot(), tree, ids, tupleId);
 }
 
-isl::multi_union_pw_aff MappedScop::blockMappingSchedule(
+isl::MultiUnionPwAff<Statement, Block> MappedScop::blockMappingSchedule(
     const ScheduleTree* tree) const {
   std::vector<mapping::MappingId> ids;
   for (size_t i = 0; i < numBlocks.view.size(); ++i) {
     ids.emplace_back(mapping::BlockId::makeId(i));
   }
   auto tupleId = isl::id(tree->ctx_, kGrid);
-  return extractDomainToIds(scop_->scheduleRoot(), tree, ids, tupleId);
+  return extractDomainToIds<Block>(scop_->scheduleRoot(), tree, ids, tupleId);
 }
 
 Scop::SyncLevel MappedScop::findBestSync(
     ScheduleTree* st1,
     ScheduleTree* st2,
-    isl::multi_union_pw_aff domainToThread,
-    isl::multi_union_pw_aff domainToWarp) {
+    isl::MultiUnionPwAff<Statement, Thread> domainToThread,
+    isl::MultiUnionPwAff<Statement, Warp> domainToWarp) {
   // Active points in the two schedule trees
   auto stRoot = scop_->scheduleRoot();
   auto activePoints1 = activeDomainPointsBelow(stRoot, st1);

tc/core/polyhedral/cuda/mapped_scop.h

@@ -188,8 +188,8 @@ class MappedScop {
   Scop::SyncLevel findBestSync(
       detail::ScheduleTree* st1,
       detail::ScheduleTree* st2,
-      isl::multi_union_pw_aff domainToThread,
-      isl::multi_union_pw_aff domainToWarp);
+      isl::MultiUnionPwAff<Statement, Thread> domainToThread,
+      isl::MultiUnionPwAff<Statement, Warp> domainToWarp);
 
  public:
   // Find best configuration of synchronizations in a sequence, minimizing
@@ -210,14 +210,14 @@ class MappedScop {
   // to the thread identifiers, where all branches in "tree"
   // are assumed to have been mapped to thread identifiers.
   // The result lives in a space of the form block[x, ...].
-  isl::multi_union_pw_aff threadMappingSchedule(
+  isl::MultiUnionPwAff<Statement, Thread> threadMappingSchedule(
       const detail::ScheduleTree* tree) const;
 
   // Extract a mapping from the domain elements active at "tree"
   // to the block identifiers, where all branches in "tree"
   // are assumed to have been mapped to block identifiers.
   // The result lives in a space of the form grid[x, ...].
-  isl::multi_union_pw_aff blockMappingSchedule(
+  isl::MultiUnionPwAff<Statement, Block> blockMappingSchedule(
       const detail::ScheduleTree* tree) const;
 
  private:

tc/core/polyhedral/cuda/memory_promotion_heuristic.cc

@@ -225,7 +225,7 @@ bool promotionImprovesCoalescing(
     auto depth = marker->scheduleDepth(root);
     auto activePoints = activeDomainPoints(root, mapping);
     auto localAccesses = originalAccesses.intersect_domain(activePoints);
-    auto schedule = prefixSchedule(root, marker);
+    auto schedule = prefixSchedule<Prefix>(root, marker);
     auto scheduledAccesses = localAccesses.apply_domain(schedule);
     for (auto access : isl::UnionAsVector<isl::union_map>(scheduledAccesses)) {
       auto scheduleSpace = access.get_space().domain();
@@ -262,6 +262,8 @@ isl::union_set collectMappingsTo(const Scop& scop) {
   return mapping;
 }
 
+struct Unrolled;
+
 /*
  * Check that only unrolled loops may appear in access subscripts.
  * Because the scoping point can be above a branching tree, descend into each
@@ -292,11 +294,12 @@ isl::union_set collectMappingsTo(const Scop& scop) {
  * different references may have different values, but all of them remain
  * independent of non-unrolled loop iterators.
  */
+template <typename Outer>
 bool accessSubscriptsAreUnrolledLoops(
     const TensorReferenceGroup& group,
     const detail::ScheduleTree* root,
     const detail::ScheduleTree* scope,
-    isl::multi_union_pw_aff outerSchedule) {
+    isl::MultiUnionPwAff<Statement, Outer> outerSchedule) {
   using namespace detail;
 
   auto nodes = ScheduleTree::collect(scope);
@@ -315,7 +318,7 @@ bool accessSubscriptsAreUnrolledLoops(
 
     auto unrolledDims = isl::union_pw_aff_list(leaf->ctx_, 1);
     for (auto node : ancestors) {
-      auto band = node->as<detail::ScheduleTreeBand>();
+      auto band = node->template as<detail::ScheduleTreeBand>();
       if (!band) {
         continue;
       }
@@ -331,8 +334,9 @@ bool accessSubscriptsAreUnrolledLoops(
     }
 
     auto space =
-        subdomain.get_space().add_unnamed_tuple_ui(unrolledDims.size());
-    auto unrolledDimsMupa = isl::multi_union_pw_aff(space, unrolledDims);
+        subdomain.get_space().template add_unnamed_tuple_ui<Unrolled>(unrolledDims.size());
+    auto unrolledDimsMupa = isl::MultiUnionPwAff<Statement, Unrolled>(
+        space, isl::UnionPwAffListOn<Statement>(unrolledDims));
 
     // It is possible that no loops are unrolled, in which case
     // unrolledDimsMupa is zero-dimensional and needs an explicit domain
@@ -341,10 +345,11 @@ bool accessSubscriptsAreUnrolledLoops(
         unrolledDimsMupa.intersect_domain(group.originalAccesses().domain());
 
     auto accesses = group.originalAccesses();
-    auto schedule = outerSchedule.flat_range_product(unrolledDimsMupa);
-    accesses = accesses.apply_domain(isl::union_map::from(schedule));
+    auto schedule = outerSchedule.range_product(unrolledDimsMupa);
+    auto scheduleMap = schedule.toUnionMap();
+    auto scheduledAccesses = accesses.apply_domain(scheduleMap);
 
-    if (!accesses.is_single_valued()) {
+    if (!scheduledAccesses.is_single_valued()) {
       return false;
     }
   }
@@ -364,23 +369,25 @@ bool accessSubscriptsAreUnrolledLoops(
  * thread associated to a given pair of tensor element and outer schedule
  * iteration.
  */
+template <typename Outer>
 bool isPromotableToRegistersBelow(
     const TensorReferenceGroup& group,
     const detail::ScheduleTree* root,
     const detail::ScheduleTree* scope,
-    isl::multi_union_pw_aff outer,
-    isl::multi_union_pw_aff thread) {
+    isl::MultiUnionPwAff<Statement, Outer> outer,
+    isl::MultiUnionPwAff<Statement, Thread> thread) {
   if (!accessSubscriptsAreUnrolledLoops(
-          group, root, scope, outer.flat_range_product(thread))) {
+          group, root, scope, outer.range_product(thread))) {
     return false;
   }
 
   auto originalAccesses = group.originalAccesses();
-  auto map = isl::union_map::from(outer);
-  map = map.range_product(originalAccesses);
-  map = map.apply_domain(isl::union_map::from(thread));
+  auto outerMap = isl::UnionMap<Statement, Outer>::from(outer);
+  auto pair = outerMap.range_product(originalAccesses);
+  auto threadMap = isl::UnionMap<Statement, Thread>::from(thread);
+  auto threadToPair = pair.apply_domain(threadMap);
 
-  return map.is_injective();
+  return threadToPair.is_injective();
 }
 
 /*
@@ -653,15 +660,15 @@ void promoteToRegistersBelow(MappedScop& mscop, detail::ScheduleTree* scope) {
   auto blockSchedule = mscop.blockMappingSchedule(mscop.schedule());
 
   // Pure affine schedule without (mapping) filters.
-  auto partialSchedMupa = partialScheduleMupa(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;
@@ -675,11 +682,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.flat_range_product(threadSchedule);
       if (!hasReuseWithin(*group, schedule)) {
         continue;
       }

tc/core/polyhedral/memory_promotion.cc

@@ -408,25 +408,25 @@ namespace {
 // each dimension of the tensor is contrained by the min_aff on the left and
 // by the min_aff + extent_aff on the right.  Intersect this set with the
 // context of the scop.
-isl::set tensorElementsSet(const Scop& scop, isl::id tensorId) {
+isl::Set<Tensor> tensorElementsSet(const Scop& scop, isl::id tensorId) {
   auto halideParameter = scop.findArgument(tensorId).parameter();
-  auto space = isl::Space<>(scop.domain().get_space());
+  auto space = scop.domain().get_space();
   auto nDim = halideParameter.dimensions();
   auto tensorTuple = constructTensorTuple(space, tensorId, nDim);
   auto tensorSpace = tensorTuple.get_space();
 
-  auto tensorElements = isl::set::universe(tensorSpace);
-  auto identity = isl::multi_aff::identity(tensorSpace.map_from_set());
+  auto tensorElements = isl::Set<Tensor>::universe(tensorSpace);
+  auto identity = isl::MultiAff<Tensor,Tensor>::identity(tensorSpace.map_from_set());
   for (int i = 0; i < nDim; ++i) {
-    isl::aff minAff = halide2isl::makeIslAffFromExpr(
+    auto minAff = halide2isl::makeIslAffFromExpr(
         space, halideParameter.min_constraint(i));
-    isl::aff extentAff = halide2isl::makeIslAffFromExpr(
+    auto extentAff = halide2isl::makeIslAffFromExpr(
         space, halideParameter.extent_constraint(i));
-    minAff = minAff.unbind_params_insert_domain(tensorTuple);
-    extentAff = extentAff.unbind_params_insert_domain(tensorTuple);
+    auto minAff2 = minAff.unbind_params_insert_domain(tensorTuple);
+    auto extentAff2 = extentAff.unbind_params_insert_domain(tensorTuple);
     auto aff = identity.get_aff(i);
-    tensorElements = tensorElements & (minAff <= isl::aff_set(aff)) &
-        (isl::aff_set(aff) < (minAff + extentAff));
+    tensorElements = tensorElements & (minAff2.le_set(aff)) &
+        (aff.lt_set(minAff2 + extentAff2));
   }
 
   tensorElements = tensorElements.intersect_params(scop.context());
@@ -493,8 +493,8 @@ ScheduleTree* insertCopiesUnder(
   auto writeSchedule = isl::multi_union_pw_aff(identityCopySchedule.pullback(
       isl::multi_aff::wrapped_range_map(writeSpace)));
 
-  auto readBandNode = ScheduleTree::makeBand(readSchedule);
-  auto writeBandNode = ScheduleTree::makeBand(writeSchedule);
+  auto readBandNode = ScheduleTree::makeBand(isl::MultiUnionPwAff<Statement, Band>(readSchedule));
+  auto writeBandNode = ScheduleTree::makeBand(isl::MultiUnionPwAff<Statement, Band>(writeSchedule));
 
   if (unrollAllCopies) {
     unrollAllMembers(readBandNode->as<detail::ScheduleTreeBand>());
@@ -542,14 +542,18 @@ ScheduleTree* insertCopiesUnder(
 
   if (reads) {
     insertExtensionBefore(
-        root, tree, tree->child({0}), readExtension, std::move(readFilterNode));
+        root,
+        tree,
+        tree->child({0}),
+        isl::UnionMap<Prefix, Statement>(readExtension),
+        std::move(readFilterNode));
   }
   if (writes) {
     insertExtensionAfter(
         root,
         tree,
         tree->child({0}),
-        writeExtension,
+        isl::UnionMap<Prefix, Statement>(writeExtension),
         std::move(writeFilterNode));
   }
 

tc/core/polyhedral/memory_promotion.h

@@ -17,6 +17,7 @@
 
 #include <iostream>
 
+#include "tc/core/polyhedral/domain_types.h"
 #include "tc/core/polyhedral/schedule_tree.h"
 #include "tc/core/polyhedral/scop.h"
 #include "tc/external/isl.h"
@@ -136,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_isl_conversion.cc

@@ -23,6 +23,7 @@
 
 #include "tc/core/check.h"
 #include "tc/core/flags.h"
+#include "tc/core/polyhedral/domain_types.h"
 #include "tc/core/polyhedral/schedule_transforms.h"
 #include "tc/external/isl.h"
 
@@ -81,7 +82,7 @@ isl::schedule_node insertBranch(
  */
 std::vector<size_t> findCorePositions(
     const ScheduleTree* st,
-    isl::union_set domain) {
+    isl::UnionSet<Statement> domain) {
   std::vector<size_t> positions;
   TC_CHECK(st->as<ScheduleTreeSequence>());
   for (size_t i = 0; i < st->numChildren(); ++i) {
@@ -125,7 +126,7 @@ isl::schedule_node insertExtension(
     isl::schedule_node node,
     const ScheduleTree* st) {
   auto depth0 = node.get_tree_depth();
-  auto domain = node.get_universe_domain();
+  auto domain = isl::UnionSet<Statement>(node.get_universe_domain());
   auto child = st->child({0});
   auto corePos = findCorePositions(child, domain);
   TC_CHECK(!corePos.empty());
@@ -242,16 +243,17 @@ std::unique_ptr<ScheduleTreeBand> fromIslScheduleNodeBand(
   for (size_t i = 0; i < n; ++i) {
     coincident[i] = b.member_get_coincident(i);
   }
+  auto mupa = isl::MultiUnionPwAff<Statement, Band>(b.get_partial_schedule());
   return ScheduleTreeBand::make(
-      b.get_partial_schedule(), b.get_permutable(), coincident, unroll);
+      mupa, b.get_permutable(), coincident, unroll);
 }
 
 std::unique_ptr<ScheduleTree> elemFromIslScheduleNode(isl::schedule_node node) {
   auto ctx = node.get_ctx();
   if (auto band = node.as<isl::schedule_node_band>()) {
     return fromIslScheduleNodeBand(band);
   } else if (auto context = node.as<isl::schedule_node_context>()) {
-    auto c = context.get_context();
+    auto c = isl::Set<Prefix>(context.get_context());
     return ScheduleTreeContext::make(c);
   } else if (auto domain = node.as<isl::schedule_node_domain>()) {
     auto c = domain.get_domain();