Merge pull request #542 from facebookresearch/separate-schedule-utils

Separate out schedule_utils.{cc,h} from schedule_transforms.{cc,h}

Author: ftynse

tc/core/CMakeLists.txt

@@ -29,6 +29,7 @@ add_library(
   polyhedral/schedule_tree.cc
   polyhedral/schedule_tree_elem.cc
   polyhedral/schedule_print.cc
+  polyhedral/schedule_utils.cc
   polyhedral/scop.cc
   polyhedral/separation.cc
   polyhedral/unroll.cc

tc/core/polyhedral/cuda/mapped_scop.cc

@@ -35,6 +35,7 @@
 #include "tc/core/polyhedral/functional.h"
 #include "tc/core/polyhedral/schedule_transforms.h"
 #include "tc/core/polyhedral/schedule_tree_matcher.h"
+#include "tc/core/polyhedral/schedule_utils.h"
 #include "tc/core/polyhedral/scop.h"
 #include "tc/core/polyhedral/separation.h"
 #include "tc/core/polyhedral/unroll.h"

tc/core/polyhedral/cuda/memory_promotion_heuristic.cc

@@ -23,6 +23,7 @@
 #include "tc/core/polyhedral/memory_promotion.h"
 #include "tc/core/polyhedral/schedule_tree.h"
 #include "tc/core/polyhedral/schedule_tree_matcher.h"
+#include "tc/core/polyhedral/schedule_utils.h"
 #include "tc/core/polyhedral/unroll.h"
 
 #include <algorithm>

tc/core/polyhedral/cuda/tighten_launch_bounds.cc

@@ -20,8 +20,8 @@
 #include "tc/core/polyhedral/cuda/mapping_types.h"
 #include "tc/core/polyhedral/exceptions.h"
 #include "tc/core/polyhedral/functional.h"
-#include "tc/core/polyhedral/schedule_transforms.h"
 #include "tc/core/polyhedral/schedule_tree.h"
+#include "tc/core/polyhedral/schedule_utils.h"
 
 namespace tc {
 namespace polyhedral {

tc/core/polyhedral/schedule_transforms.cc

@@ -34,6 +34,7 @@
 #include "tc/core/polyhedral/mapping_types.h"
 #include "tc/core/polyhedral/schedule_tree_elem.h"
 #include "tc/core/polyhedral/schedule_tree_matcher.h"
+#include "tc/core/polyhedral/schedule_utils.h"
 #include "tc/core/scope_guard.h"
 #include "tc/external/isl.h"
 
@@ -47,164 +48,6 @@ namespace tc {
 namespace polyhedral {
 using namespace detail;
 
-isl::union_map extendSchedule(
-    const ScheduleTree* node,
-    isl::union_map schedule) {
-  if (auto bandElem = node->elemAs<ScheduleTreeElemBand>()) {
-    if (bandElem->nMember() > 0) {
-      schedule =
-          schedule.flat_range_product(isl::union_map::from(bandElem->mupa_));
-    }
-  } else if (auto filterElem = node->elemAs<ScheduleTreeElemFilter>()) {
-    schedule = schedule.intersect_domain(filterElem->filter_);
-  } else if (auto extensionElem = node->elemAs<ScheduleTreeElemExtension>()) {
-    // FIXME: we may need to restrict the range of reversed extension map to
-    // schedule values that correspond to active domain elements at this
-    // point.
-    schedule = schedule.unite(
-        extensionElem->extension_.reverse().intersect_range(schedule.range()));
-  }
-
-  return schedule;
-}
-
-namespace {
-isl::union_map partialScheduleImpl(
-    const ScheduleTree* root,
-    const ScheduleTree* node,
-    bool useNode) {
-  auto nodes = node->ancestors(root);
-  if (useNode) {
-    nodes.push_back(node);
-  }
-  TC_CHECK_GT(nodes.size(), 0u) << "root node does not have a prefix schedule";
-  auto domain = root->elemAs<ScheduleTreeElemDomain>();
-  TC_CHECK(domain);
-  auto schedule = isl::union_map::from_domain(domain->domain_);
-  for (auto anc : nodes) {
-    if (anc->elemAs<ScheduleTreeElemDomain>()) {
-      TC_CHECK(anc == root);
-    } else {
-      schedule = extendSchedule(anc, schedule);
-    }
-  }
-  return schedule;
-}
-} // namespace
-
-isl::union_map prefixSchedule(
-    const ScheduleTree* root,
-    const ScheduleTree* node) {
-  return partialScheduleImpl(root, node, false);
-}
-
-isl::union_map partialSchedule(
-    const ScheduleTree* root,
-    const ScheduleTree* node) {
-  return partialScheduleImpl(root, node, true);
-}
-
-namespace {
-/*
- * If "node" is any filter, then intersect "domain" with that filter.
- */
-isl::union_set applyFilter(isl::union_set domain, const ScheduleTree* node) {
-  if (auto filterElem = node->elemAs<ScheduleTreeElemFilter>()) {
-    return domain.intersect(filterElem->filter_);
-  }
-  return domain;
-}
-
-/*
- * If "node" is a mapping, then intersect "domain" with its filter.
- */
-isl::union_set applyMapping(isl::union_set domain, const ScheduleTree* node) {
-  if (auto filterElem = node->elemAs<ScheduleTreeElemMapping>()) {
-    return domain.intersect(filterElem->filter_);
-  }
-  return domain;
-}
-
-// Get the set of domain elements that are active below
-// the given branch of nodes, filtered using "filter".
-//
-// Domain elements are introduced by the root domain node.  Some nodes
-// refine this set of elements based on "filter".  Extension nodes
-// are considered to introduce additional domain points.
-isl::union_set collectDomain(
-    const ScheduleTree* root,
-    const vector<const ScheduleTree*>& nodes,
-    isl::union_set (*filter)(isl::union_set domain, const ScheduleTree* node)) {
-  auto domainElem = root->elemAs<ScheduleTreeElemDomain>();
-  TC_CHECK(domainElem) << "root must be a Domain node" << *root;
-
-  auto domain = domainElem->domain_;
-
-  for (auto anc : nodes) {
-    domain = filter(domain, anc);
-    if (auto extensionElem = anc->elemAs<ScheduleTreeElemExtension>()) {
-      auto parentSchedule = prefixSchedule(root, anc);
-      auto extension = extensionElem->extension_;
-      TC_CHECK(parentSchedule) << "missing root domain node";
-      parentSchedule = parentSchedule.intersect_domain(domain);
-      domain = domain.unite(parentSchedule.range().apply(extension));
-    }
-  }
-  return domain;
-}
-
-// Get the set of domain elements that are active below
-// the given branch of nodes.
-isl::union_set activeDomainPointsHelper(
-    const ScheduleTree* root,
-    const vector<const ScheduleTree*>& nodes) {
-  return collectDomain(root, nodes, &applyFilter);
-}
-
-} // namespace
-
-isl::union_set prefixMappingFilter(
-    const ScheduleTree* root,
-    const ScheduleTree* node) {
-  return collectDomain(root, node->ancestors(root), &applyMapping);
-}
-
-isl::union_set activeDomainPoints(
-    const ScheduleTree* root,
-    const ScheduleTree* node) {
-  return activeDomainPointsHelper(root, node->ancestors(root));
-}
-
-isl::union_set activeDomainPointsBelow(
-    const ScheduleTree* root,
-    const ScheduleTree* node) {
-  auto ancestors = node->ancestors(root);
-  ancestors.emplace_back(node);
-  return activeDomainPointsHelper(root, ancestors);
-}
-
-vector<ScheduleTree*> collectScheduleTreesPath(
-    std::function<ScheduleTree*(ScheduleTree*)> next,
-    ScheduleTree* start) {
-  vector<ScheduleTree*> res{start};
-  auto n = start;
-  while ((n = next(n)) != nullptr) {
-    res.push_back(n);
-  }
-  return res;
-}
-
-vector<const ScheduleTree*> collectScheduleTreesPath(
-    std::function<const ScheduleTree*(const ScheduleTree*)> next,
-    const ScheduleTree* start) {
-  vector<const ScheduleTree*> res{start};
-  auto n = start;
-  while ((n = next(n)) != nullptr) {
-    res.push_back(n);
-  }
-  return res;
-}
-
 // Replace "tree" in the list of its parent's children with newTree.
 // Returns the pointer to newTree for call chaining purposes.
 ScheduleTree* swapSubtree(
@@ -432,85 +275,6 @@ ScheduleTree* bandScale(ScheduleTree* tree, const vector<size_t>& scales) {
   return tree;
 }
 
-namespace {
-
-template <typename T>
-vector<T> reversed(const vector<T>& vec) {
-  vector<T> result;
-  result.reserve(vec.size());
-  result.insert(result.begin(), vec.rbegin(), vec.rend());
-  return result;
-}
-
-template <typename T>
-vector<const ScheduleTree*> filterType(const vector<const ScheduleTree*>& vec) {
-  vector<const ScheduleTree*> result;
-  for (auto e : vec) {
-    if (e->elemAs<T>()) {
-      result.push_back(e);
-    }
-  }
-  return result;
-}
-
-template <typename T, typename Func>
-T foldl(const vector<const ScheduleTree*> vec, Func op, T init = T()) {
-  T value = init;
-  for (auto st : vec) {
-    value = op(st, value);
-  }
-  return value;
-}
-
-template <typename... Args>
-ostream& operator<<(ostream& os, const vector<Args...>& v) {
-  os << "[";
-  bool first = true;
-  for (auto const& ve : v) {
-    if (!first) {
-      os << ", ";
-    }
-    os << ve;
-    first = true;
-  }
-  os << "]";
-  return os;
-}
-} // namespace
-
-isl::multi_union_pw_aff infixScheduleMupa(
-    const ScheduleTree* root,
-    const ScheduleTree* relativeRoot,
-    const ScheduleTree* tree) {
-  auto domainElem = root->elemAs<ScheduleTreeElemDomain>();
-  TC_CHECK(domainElem);
-  auto domain = domainElem->domain_.universe();
-  auto zero = isl::multi_val::zero(domain.get_space().set_from_params());
-  auto prefix = isl::multi_union_pw_aff(domain, zero);
-  prefix = foldl(
-      filterType<ScheduleTreeElemBand>(tree->ancestors(relativeRoot)),
-      [](const ScheduleTree* st, isl::multi_union_pw_aff pref) {
-        auto mupa = st->elemAs<ScheduleTreeElemBand>()->mupa_;
-        return pref.flat_range_product(mupa);
-      },
-      prefix);
-  return prefix;
-}
-
-isl::multi_union_pw_aff prefixScheduleMupa(
-    const ScheduleTree* root,
-    const ScheduleTree* tree) {
-  return infixScheduleMupa(root, root, tree);
-}
-
-isl::multi_union_pw_aff partialScheduleMupa(
-    const detail::ScheduleTree* root,
-    const detail::ScheduleTree* tree) {
-  auto prefix = prefixScheduleMupa(root, tree);
-  auto band = tree->elemAs<ScheduleTreeElemBand>();
-  return band ? prefix.flat_range_product(band->mupa_) : prefix;
-}
-
 void updateTopLevelContext(detail::ScheduleTree* root, isl::set context) {
   if (!matchOne(tc::polyhedral::domain(tc::polyhedral::context(any())), root)) {
     root->appendChild(ScheduleTree::makeContext(
@@ -832,55 +596,5 @@ void orderAfter(ScheduleTree* root, ScheduleTree* tree, isl::union_set filter) {
   seq->insertChild(0, gistedFilter(other, parent->detachChild(childPos)));
   parent->insertChild(childPos, std::move(seq));
 }
-
-/*
- * Extract a mapping from the domain elements active at "tree"
- * to identifiers "ids", where all branches in "tree"
- * are assumed to have been mapped to these identifiers.
- * The result lives in a space of the form "tupleId"["ids"...].
- */
-isl::multi_union_pw_aff extractDomainToIds(
-    const detail::ScheduleTree* root,
-    const detail::ScheduleTree* tree,
-    const std::vector<mapping::MappingId>& ids,
-    isl::id tupleId) {
-  using namespace polyhedral::detail;
-
-  auto space = isl::space(tree->ctx_, 0);
-  auto empty = isl::union_set::empty(space);
-  space = space.named_set_from_params_id(tupleId, ids.size());
-  auto zero = isl::multi_val::zero(space);
-  auto domainToIds = isl::multi_union_pw_aff(empty, zero);
-
-  for (auto mapping : tree->collect(tree, ScheduleTreeType::Mapping)) {
-    auto mappingNode = mapping->elemAs<ScheduleTreeElemMapping>();
-    auto list = isl::union_pw_aff_list(tree->ctx_, ids.size());
-    for (auto id : ids) {
-      if (mappingNode->mapping.count(id) == 0) {
-        break;
-      }
-      auto idMap = mappingNode->mapping.at(id);
-      list = list.add(idMap);
-    }
-    // Ignore this node if it does not map to all required ids.
-    if (static_cast<size_t>(list.size()) != ids.size()) {
-      continue;
-    }
-    auto nodeToIds = isl::multi_union_pw_aff(space, list);
-    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);
-    domainToIds = domainToIds.union_add(nodeToIds);
-  }
-
-  auto active = activeDomainPoints(root, tree);
-  TC_CHECK(active.is_subset(domainToIds.domain()))
-      << "not all domain points of\n"
-      << active << "\nwere mapped to the required ids";
-
-  return domainToIds;
-}
-
 } // namespace polyhedral
 } // namespace tc