[RFC] use templated isl types schedule tree

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/memory_promotion.cc

@@ -500,8 +500,10 @@ 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>());
@@ -551,14 +553,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/schedule_isl_conversion.cc

@@ -243,16 +243,16 @@ std::unique_ptr<ScheduleTreeBand> fromIslScheduleNodeBand(
   for (size_t i = 0; i < n; ++i) {
     coincident[i] = b.member_get_coincident(i);
   }
-  return ScheduleTreeBand::make(
-      b.get_partial_schedule(), b.get_permutable(), coincident, unroll);
+  auto mupa = isl::MultiUnionPwAff<Statement, Band>(b.get_partial_schedule());
+  return ScheduleTreeBand::make(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();

tc/core/polyhedral/schedule_transforms.cc

@@ -88,7 +88,8 @@ ScheduleTree* joinBandsHelper(ScheduleTree* st, bool& moveChildren) {
 
   auto& partialSchedule = eb->mupa_;
   auto& partialScheduleChild = ebChild->mupa_;
-  partialSchedule = partialSchedule.flat_range_product(partialScheduleChild);
+  partialSchedule =
+      partialSchedule.flat_range_product<Band>(partialScheduleChild);
   eb->coincident_.resize(
       eb->coincident_.size() + ebChild->coincident_.size(), false);
   eb->unroll_.insert(
@@ -284,7 +285,9 @@ ScheduleTree* insertTopLevelEmptyBand(ScheduleTree* root) {
   return insertNodeBelow(node, ScheduleTree::makeEmptyBand(root));
 }
 
-void updateTopLevelContext(detail::ScheduleTree* root, isl::set context) {
+void updateTopLevelContext(
+    detail::ScheduleTree* root,
+    isl::Set<Prefix> context) {
   if (!matchOne(tc::polyhedral::domain(tc::polyhedral::context(any())), root)) {
     root->appendChild(
         ScheduleTree::makeContext(context, root->detachChildren()));
@@ -393,7 +396,7 @@ void insertExtensionAt(
     ScheduleTree* relativeRoot,
     ScheduleTree* seqNode,
     size_t pos,
-    isl::union_map extension,
+    isl::UnionMap<Prefix, Statement> extension,
     ScheduleTreeUPtr&& filterNode) {
   auto extensionTree = seqNode->ancestor(relativeRoot, 1);
   auto extensionNode = extensionTree->as<detail::ScheduleTreeExtension>();
@@ -412,7 +415,7 @@ void insertExtensionBefore(
     const ScheduleTree* root,
     ScheduleTree* relativeRoot,
     ScheduleTree* tree,
-    isl::union_map extension,
+    isl::UnionMap<Prefix, Statement> extension,
     ScheduleTreeUPtr&& filterNode) {
   size_t pos;
   auto parent = tree->ancestor(relativeRoot, 1);
@@ -441,7 +444,7 @@ void insertExtensionAfter(
     const ScheduleTree* root,
     ScheduleTree* relativeRoot,
     ScheduleTree* tree,
-    isl::union_map extension,
+    isl::UnionMap<Prefix, Statement> extension,
     ScheduleTreeUPtr&& filterNode) {
   size_t pos;
   auto parent = tree->ancestor(relativeRoot, 1);
@@ -503,7 +506,7 @@ namespace {
  * of band node partial schedules.
  * Elements of a sequence that end up with an empty filter are removed.
  */
-void gist(ScheduleTree* tree, isl::union_set context) {
+void gist(ScheduleTree* tree, isl::UnionSet<Statement> context) {
   if (auto bandElem = tree->as<ScheduleTreeBand>()) {
     bandElem->mupa_ = bandElem->mupa_.gist(context);
   } else if (auto filterElem = tree->as<ScheduleTreeMapping>()) {
@@ -533,7 +536,9 @@ void gist(ScheduleTree* tree, isl::union_set context) {
  * Create a filter node with the given filter and single child node,
  * after simplifying the child node in the context of the filter.
  */
-ScheduleTreeUPtr gistedFilter(isl::union_set filter, ScheduleTreeUPtr child) {
+ScheduleTreeUPtr gistedFilter(
+    isl::UnionSet<Statement> filter,
+    ScheduleTreeUPtr child) {
   gist(child.get(), filter);
   return ScheduleTree::makeFilter(filter, std::move(child));
 }
@@ -556,7 +561,8 @@ bool canOrder(
   auto zero = isl::MultiVal<isl::Anonymous>::zero(space);
   auto one = zero.set_val(0, isl::val::one(ctx));
   auto order = isl::MultiUnionPwAff<Statement, isl::Anonymous>(first, zero);
-  order = order.union_add(isl::MultiUnionPwAff<Statement, isl::Anonymous>(second, one));
+  order = order.union_add(
+      isl::MultiUnionPwAff<Statement, isl::Anonymous>(second, one));
 
   // Check that this ordering preserves all dependences.
   auto preserved = dependences.lex_lt_at(order).unite(dependences.eq_at(order));

tc/core/polyhedral/schedule_transforms.h

@@ -111,7 +111,9 @@ detail::ScheduleTree* insertTopLevelEmptyBand(detail::ScheduleTree* root);
 // Update the top-level context node by intersecting it with "context".  The
 // top-level context node must be located directly under the root of the tree.
 // If there is no such node, insert one first.
-void updateTopLevelContext(detail::ScheduleTree* root, isl::set context);
+void updateTopLevelContext(
+    detail::ScheduleTree* root,
+    isl::Set<Prefix> context);
 
 // In a tree starting at "root", insert a sequence node with
 // as only child the node identified by "tree".
@@ -175,7 +177,7 @@ void insertExtensionBefore(
     const detail::ScheduleTree* root,
     detail::ScheduleTree* relativeRoot,
     detail::ScheduleTree* tree,
-    isl::union_map extension,
+    isl::UnionMap<Prefix, Statement> extension,
     ScheduleTreeUPtr&& filterNode);
 
 // Insert an extension with the given extension map and extension filter node
@@ -190,7 +192,7 @@ void insertExtensionAfter(
     const detail::ScheduleTree* root,
     detail::ScheduleTree* relativeRoot,
     detail::ScheduleTree* tree,
-    isl::union_map extension,
+    isl::UnionMap<Prefix, Statement> extension,
     ScheduleTreeUPtr&& filterNode);
 
 // Given a sequence node in the schedule tree, insert

tc/core/polyhedral/schedule_tree.cc

@@ -199,7 +199,7 @@ size_t ScheduleTree::scheduleDepth(const ScheduleTree* relativeRoot) const {
 }
 
 std::unique_ptr<ScheduleTree> ScheduleTree::makeBand(
-    isl::multi_union_pw_aff mupa,
+    isl::MultiUnionPwAff<Statement, Band> mupa,
     std::vector<ScheduleTreeUPtr>&& children) {
   std::vector<bool> coincident(mupa.size(), false);
   std::vector<bool> unroll(mupa.size(), false);
@@ -211,9 +211,9 @@ std::unique_ptr<ScheduleTree> ScheduleTree::makeBand(
 ScheduleTreeUPtr ScheduleTree::makeEmptyBand(const ScheduleTree* root) {
   auto domain = root->as<ScheduleTreeDomain>();
   TC_CHECK(domain);
-  auto space = domain->domain_.get_space().set_from_params();
-  auto mv = isl::multi_val::zero(space);
-  auto zero = isl::multi_union_pw_aff(domain->domain_, mv);
+  auto space = domain->domain_.get_space().add_unnamed_tuple_ui<Band>(0);
+  auto mv = isl::MultiVal<Band>::zero(space);
+  auto zero = isl::MultiUnionPwAff<Statement, Band>(domain->domain_, mv);
   return ScheduleTree::makeBand(zero);
 }
 
@@ -224,7 +224,7 @@ std::unique_ptr<ScheduleTree> ScheduleTree::makeDomain(
 }
 
 std::unique_ptr<ScheduleTree> ScheduleTree::makeContext(
-    isl::set context,
+    isl::Set<Prefix> context,
     std::vector<ScheduleTreeUPtr>&& children) {
   return ScheduleTreeContext::make(context, std::move(children));
 }

tc/core/polyhedral/schedule_tree.h

@@ -21,6 +21,7 @@
 #include <vector>
 
 #include "tc/core/check.h"
+#include "tc/core/polyhedral/domain_types.h"
 #include "tc/core/polyhedral/mapping_types.h"
 #include "tc/core/polyhedral/options.h"
 #include "tc/core/utils/vararg.h"
@@ -280,7 +281,7 @@ struct ScheduleTree {
   // Factory functions
   //
   static ScheduleTreeUPtr makeBand(
-      isl::multi_union_pw_aff mupa,
+      isl::MultiUnionPwAff<Statement, Band> mupa,
       std::vector<ScheduleTreeUPtr>&& children = {});
 
   // Return a zero-dimensional band for use in a tree with the given root.
@@ -291,7 +292,7 @@ struct ScheduleTree {
       std::vector<ScheduleTreeUPtr>&& children = {});
 
   static ScheduleTreeUPtr makeContext(
-      isl::set context,
+      isl::Set<Prefix> context,
       std::vector<ScheduleTreeUPtr>&& children = {});
 
   static ScheduleTreeUPtr makeFilter(
@@ -332,7 +333,7 @@ struct ScheduleTree {
 
   template <typename... Args>
   static ScheduleTreeUPtr makeBand(
-      isl::multi_union_pw_aff mupa,
+      isl::MultiUnionPwAff<Statement, Band> mupa,
       Args&&... args) {
     return makeBand(
         mupa, vectorFromArgs<ScheduleTreeUPtr>(std::forward<Args>(args)...));
@@ -345,7 +346,7 @@ struct ScheduleTree {
   }
 
   template <typename... Args>
-  static ScheduleTreeUPtr makeContext(isl::set context, Args&&... args) {
+  static ScheduleTreeUPtr makeContext(isl::Set<> context, Args&&... args) {
     return makeContext(
         context, vectorFromArgs<ScheduleTreeUPtr>(std::forward<Args>(args)...));
   }

tc/core/polyhedral/schedule_tree_elem.cc

@@ -26,6 +26,7 @@
 #include "tc/core/check.h"
 #include "tc/core/constants.h"
 #include "tc/core/flags.h"
+#include "tc/core/polyhedral/domain_types.h"
 #include "tc/core/polyhedral/schedule_isl_conversion.h"
 #include "tc/core/polyhedral/schedule_tree.h"
 #include "tc/core/scope_guard.h"
@@ -185,7 +186,7 @@ std::unique_ptr<ScheduleTreeSet> ScheduleTreeSet::make(
 }
 
 std::unique_ptr<ScheduleTreeBand> ScheduleTreeBand::make(
-    isl::multi_union_pw_aff mupa,
+    isl::MultiUnionPwAff<Statement, Band> mupa,
     bool permutable,
     std::vector<bool> coincident,
     std::vector<bool> unroll,
@@ -238,8 +239,8 @@ void ScheduleTreeBand::drop(size_t pos, size_t n) {
   auto list = mupa_.get_union_pw_aff_list();
   auto space = mupa_.get_space().params();
   list = list.drop(pos, n);
-  space = space.add_unnamed_tuple_ui(list.size());
-  mupa_ = isl::multi_union_pw_aff(space, list);
+  auto spaceBand = space.add_unnamed_tuple_ui<Band>(list.size());
+  mupa_ = isl::MultiUnionPwAff<Statement, Band>(spaceBand, list);
 
   std::copy(
       coincident_.begin() + pos + n,

tc/core/polyhedral/schedule_tree_elem.h

@@ -64,7 +64,7 @@ struct ScheduleTreeContext : public ScheduleTree {
   }
 
  public:
-  isl::set context_;
+  isl::Set<Prefix> context_;
 };
 
 struct ScheduleTreeDomain : public ScheduleTree {
@@ -100,7 +100,7 @@ struct ScheduleTreeDomain : public ScheduleTree {
   }
 
  public:
-  isl::union_set domain_;
+  isl::UnionSet<Statement> domain_;
 };
 
 struct ScheduleTreeExtension : public ScheduleTree {
@@ -136,7 +136,7 @@ struct ScheduleTreeExtension : public ScheduleTree {
   }
 
  public:
-  isl::union_map extension_;
+  isl::UnionMap<Prefix, Statement> extension_;
 };
 
 struct ScheduleTreeFilter : public ScheduleTree {
@@ -172,7 +172,7 @@ struct ScheduleTreeFilter : public ScheduleTree {
   }
 
  public:
-  isl::union_set filter_;
+  isl::UnionSet<Statement> filter_;
 };
 
 struct ScheduleTreeMapping : public ScheduleTree {
@@ -216,7 +216,7 @@ struct ScheduleTreeMapping : public ScheduleTree {
   // Mapping from identifiers to affine functions on domain elements.
   const Mapping mapping;
   // Assignment of the affine functions to the identifiers as parameters.
-  isl::union_set filter_;
+  isl::UnionSet<Statement> filter_;
 };
 
 struct ScheduleTreeSequence : public ScheduleTree {
@@ -310,7 +310,7 @@ struct ScheduleTreeBand : public ScheduleTree {
   // Replace "mupa" by its greatest integer part to ensure that the
   // schedule is always integral.
   static std::unique_ptr<ScheduleTreeBand> make(
-      isl::multi_union_pw_aff mupa,
+      isl::MultiUnionPwAff<Statement, Band> mupa,
       bool permutable,
       std::vector<bool> coincident,
       std::vector<bool> unroll,
@@ -336,18 +336,17 @@ struct ScheduleTreeBand : public ScheduleTree {
   isl::MultiUnionPwAff<Statement, Range> memberRange(size_t first, size_t n)
       const {
     auto list = mupa_.get_union_pw_aff_list();
-    auto space = mupa_.get_space().params().add_unnamed_tuple_ui(n);
+    auto space = mupa_.get_space().params().add_unnamed_tuple_ui<Range>(n);
     auto end = first + n;
     TC_CHECK_LE(end, nMember());
     list = list.drop(end, nMember() - end);
     list = list.drop(0, first);
-    return isl::MultiUnionPwAff<Statement, Range>(
-        isl::multi_union_pw_aff(space, list));
+    return isl::MultiUnionPwAff<Statement, Range>(space, list);
   }
 
  public:
   bool permutable_{false};
-  isl::multi_union_pw_aff mupa_;
+  isl::MultiUnionPwAff<Statement, Band> mupa_;
 
   std::vector<bool> coincident_;
   // For each member, should the corresponding loop in the generated code

tc/core/polyhedral/schedule_utils-inl.h

@@ -27,7 +27,8 @@ namespace tc {
 namespace polyhedral {
 
 namespace detail {
-inline isl::union_map partialScheduleImpl(
+template <typename Schedule>
+inline isl::UnionMap<Statement, Schedule> partialScheduleImpl(
     const ScheduleTree* root,
     const ScheduleTree* node,
     bool useNode) {
@@ -36,9 +37,10 @@ inline isl::union_map partialScheduleImpl(
     nodes.push_back(node);
   }
   TC_CHECK_GT(nodes.size(), 0u) << "root node does not have a prefix schedule";
-  auto domain = root->as<ScheduleTreeDomain>();
+  auto domain = root->as<detail::ScheduleTreeDomain>();
   TC_CHECK(domain);
-  auto schedule = isl::union_map::from_domain(domain->domain_);
+  auto schedule =
+      isl::UnionMap<Statement, Schedule>::from_domain(domain->domain_);
   for (auto anc : nodes) {
     if (anc->as<ScheduleTreeDomain>()) {
       TC_CHECK(anc == root);
@@ -54,15 +56,15 @@ template <typename Schedule>
 inline isl::UnionMap<Statement, Schedule> prefixSchedule(
     const detail::ScheduleTree* root,
     const detail::ScheduleTree* node) {
-  auto prefix = detail::partialScheduleImpl(root, node, false);
+  auto prefix = detail::partialScheduleImpl<Schedule>(root, node, false);
   return isl::UnionMap<Statement, Schedule>(prefix);
 }
 
 template <typename Schedule>
 inline isl::UnionMap<Statement, Schedule> partialSchedule(
     const detail::ScheduleTree* root,
     const detail::ScheduleTree* node) {
-  auto prefix = detail::partialScheduleImpl(root, node, true);
+  auto prefix = detail::partialScheduleImpl<Schedule>(root, node, true);
   return isl::UnionMap<Statement, Schedule>(prefix);
 }