Merge pull request #336 from facebookresearch/pr/insert_node

extract out generic insertNode{Above,Below}

Author: ftynse

tc/core/polyhedral/cuda/mapped_scop.cc

@@ -109,7 +109,8 @@ void MappedScop::mapRemaining(
   auto root = scop_->scheduleRoot();
   auto domain = activeDomainPoints(root, tree);
   auto filter = makeFixRemainingZeroFilter(domain, ids);
-  insertMappingFilterAbove(root, tree, filter, ids);
+  auto mapping = detail::ScheduleTree::makeMappingFilter(filter, ids);
+  insertNodeAbove(root, tree, std::move(mapping));
 
   for (size_t i = nMapped; i < nToMap; ++i) {
     if (MappingTypeId::makeId(i) == mapping::ThreadId::x()) {
@@ -171,8 +172,9 @@ void fixThreadsBelowFilter(
   // invariant that leaf mapping filters have a single space.
   // So we intersect with the universe set of the filter to only keep the
   // space for the legitimate statement.
-  insertMappingFilterBelow(
-      filterTree, mappingFilter & filter->filter_.universe(), ids);
+  mappingFilter = mappingFilter & filter->filter_.universe();
+  auto mapping = detail::ScheduleTree::makeMappingFilter(mappingFilter, ids);
+  insertNodeBelow(filterTree, std::move(mapping));
 
   for (size_t i = begin; i < end; ++i) {
     if (mapping::ThreadId::makeId(i) == mapping::ThreadId::x()) {

tc/core/polyhedral/schedule_transforms-inl.h

@@ -17,32 +17,27 @@
 
 namespace tc {
 namespace polyhedral {
-template <typename MappingIdType>
-inline detail::ScheduleTree* insertMappingFilterAbove(
+inline detail::ScheduleTree* insertNodeAbove(
     detail::ScheduleTree* root,
     detail::ScheduleTree* tree,
-    isl::union_set filter,
-    const std::unordered_set<MappingIdType, typename MappingIdType::Hash>&
-        mappingIds) {
+    ScheduleTreeUPtr&& node) {
+  CHECK_EQ(node->numChildren(), 0u);
   auto parent = tree->ancestor(root, 1);
   auto childPos = tree->positionInParent(parent);
-  parent->insertChild(
-      childPos,
-      detail::ScheduleTree::makeMappingFilter(
-          filter, mappingIds, parent->detachChild(childPos)));
+  node->appendChild(parent->detachChild(childPos));
+  parent->insertChild(childPos, std::move(node));
   return parent->child({childPos});
 }
 
-template <typename MappingIdType>
-inline void insertMappingFilterBelow(
+inline detail::ScheduleTree* insertNodeBelow(
     detail::ScheduleTree* tree,
-    isl::union_set filter,
-    const std::unordered_set<MappingIdType, typename MappingIdType::Hash>&
-        mappingIds) {
+    ScheduleTreeUPtr&& node) {
+  CHECK_EQ(node->numChildren(), 0u);
   auto numChildren = tree->numChildren();
   CHECK_LE(numChildren, 1u);
-  tree->appendChild(detail::ScheduleTree::makeMappingFilter(
-      filter, mappingIds, tree->detachChildren()));
+  node->appendChildren(tree->detachChildren());
+  tree->appendChild(std::move(node));
+  return tree->child({0});
 }
 
 template <typename MappingIdType>
@@ -70,8 +65,9 @@ inline detail::ScheduleTree* mapToParameterWithExtent(
       band->mupa_.get_union_pw_aff(pos).mod_val(isl::val(tree->ctx_, extent));
   upa = upa.sub(isl::union_pw_aff::param_on_domain(domain, id));
   auto filter = upa.zero_union_set();
-  return insertMappingFilterAbove<MappingIdType>(root, tree, filter, {id})
-      ->child({0});
+  auto mapping =
+      detail::ScheduleTree::makeMappingFilter<MappingIdType>(filter, {id});
+  return insertNodeAbove(root, tree, std::move(mapping))->child({0});
 }
 } // namespace polyhedral
 } // namespace tc

tc/core/polyhedral/schedule_transforms.cc

@@ -462,26 +462,6 @@ isl::multi_union_pw_aff partialScheduleMupa(
   return prefixScheduleMupa(root, tree).flat_range_product(band->mupa_);
 }
 
-ScheduleTree* insertBandAbove(
-    ScheduleTree* root,
-    ScheduleTree* tree,
-    isl::multi_union_pw_aff mupa) {
-  auto parent = tree->ancestor(root, 1);
-  auto childPos = tree->positionInParent(parent);
-  auto child = parent->detachChild(childPos);
-  parent->insertChild(childPos, ScheduleTree::makeBand(mupa, std::move(child)));
-  return parent->child({childPos});
-}
-
-ScheduleTree* insertBandBelow(
-    detail::ScheduleTree* tree,
-    isl::multi_union_pw_aff mupa) {
-  auto numChildren = tree->numChildren();
-  CHECK_LE(numChildren, 1u);
-  tree->appendChild(ScheduleTree::makeBand(mupa, tree->detachChildren()));
-  return tree->child({0});
-}
-
 void updateTopLevelContext(detail::ScheduleTree* root, isl::set context) {
   if (!matchOne(tc::polyhedral::domain(tc::polyhedral::context(any())), root)) {
     root->appendChild(ScheduleTree::makeContext(

tc/core/polyhedral/schedule_transforms.h

@@ -131,27 +131,6 @@ detail::ScheduleTree* mapToParameterWithExtent(
     MappingIdType id,
     size_t extent);
 
-// In a tree starting at a (relative) "root", insert a band node with the
-// given partial schedule above the node identified by "tree".
-//
-// The tree is modified in place.
-// Return a non-owning pointer to the inserted band node
-// for call chaining purposes.
-detail::ScheduleTree* insertBandAbove(
-    detail::ScheduleTree* root,
-    detail::ScheduleTree* tree,
-    isl::multi_union_pw_aff mupa);
-
-// Insert a band node with the given partial schedule below node "tree",
-// which is assumed to have at most one child.
-//
-// The tree is modified in place.
-// Return a non-owning pointer to the inserted band node
-// for call chaining purposes.
-detail::ScheduleTree* insertBandBelow(
-    detail::ScheduleTree* tree,
-    isl::multi_union_pw_aff mupa);
-
 // Update the top-level conext 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 with universe context first.
@@ -178,31 +157,26 @@ detail::ScheduleTree* insertExtensionAbove(
     detail::ScheduleTree* tree,
     isl::union_map extension);
 
-// In a tree starting at a (relative) "root", insert a mapping filter node
-// with the given filter above the node identified by "tree".
+// In a tree starting at a (relative) "root", insert the given node
+// above the node identified by "tree".
 //
 // The tree is modified in place.
-// Return a non-owning pointer to the inserted filter node
+// Return a non-owning pointer to the inserted node
 // for call chaining purposes.
-template <typename MappingIdType>
-inline detail::ScheduleTree* insertMappingFilterAbove(
+inline detail::ScheduleTree* insertNodeAbove(
     detail::ScheduleTree* root,
     detail::ScheduleTree* tree,
-    isl::union_set filter,
-    const std::unordered_set<MappingIdType, typename MappingIdType::Hash>&
-        mappingIds);
+    ScheduleTreeUPtr&& node);
 
-// Insert a mapping filter node below node "tree", which is assumed to have at
-// most one child. The underlying isl::union_set filter is constructed from
-// the arguments.
+// Insert the given node below node "tree", which is assumed to have at
+// most one child.
 //
 // The tree is modified in place.
-template <typename MappingIdType>
-inline void insertMappingFilterBelow(
+// Return a non-owning pointer to the inserted node
+// for call chaining purposes.
+inline detail::ScheduleTree* insertNodeBelow(
     detail::ScheduleTree* tree,
-    isl::union_set filter,
-    const std::unordered_set<MappingIdType, typename MappingIdType::Hash>&
-        mappingIds);
+    ScheduleTreeUPtr&& node);
 
 // Given a sequence node in the schedule tree, insert
 // a zero-dimensional extension statement with the given identifier

tc/core/polyhedral/scop.cc

@@ -492,10 +492,11 @@ detail::ScheduleTree* obtainOuterBand(detail::ScheduleTree* root) {
     CHECK(domain);
     auto space = domain->domain_.get_space().set_from_params();
     auto zero = isl::multi_union_pw_aff::zero(space);
+    auto band = ScheduleTree::makeBand(zero);
     if (n == 0) {
-      return setPermutable(insertBandBelow(tree, zero));
+      return setPermutable(insertNodeBelow(tree, std::move(band)));
     } else {
-      return setPermutable(insertBandAbove(root, tree, zero));
+      return setPermutable(insertNodeAbove(root, tree, std::move(band)));
     }
   }
   return tree;