use templated isl types in schedule tree

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;

tc/core/polyhedral/memory_promotion.cc

@@ -501,8 +501,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>());
@@ -552,14 +554,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

@@ -89,7 +89,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(
@@ -285,7 +286,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()));
@@ -394,7 +397,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>();
@@ -413,7 +416,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);
@@ -442,7 +445,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);
@@ -504,7 +507,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>()) {
@@ -534,7 +537,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));
 }

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

@@ -200,7 +200,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);
@@ -212,9 +212,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);
 }
 
@@ -225,7 +225,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