Merge pull request facebookresearch#543 from facebookresearch/schedule-xyz-subclasses

ScheduleTree subclasses

Author: ftynse

tc/core/polyhedral/codegen.cc

@@ -19,6 +19,7 @@
 #include <sstream>
 
 #include "tc/core/polyhedral/schedule_tree.h"
+#include "tc/core/polyhedral/schedule_tree_elem.h"
 
 namespace tc {
 namespace polyhedral {
@@ -30,7 +31,7 @@ isl::id_list Codegen::makeLoopIterators(
       detail::ScheduleTree::collect(root, detail::ScheduleTreeType::Band);
   size_t n = 0;
   for (auto const& node : bands) {
-    auto bandElem = node->elemAs<detail::ScheduleTreeElemBand>();
+    auto bandElem = node->as<detail::ScheduleTreeBand>();
     auto depth = node->scheduleDepth(root) + bandElem->nMember();
     if (depth > n) {
       n = depth;

tc/core/polyhedral/cuda/codegen.cc

@@ -785,9 +785,8 @@ string emitCudaKernel(
     const std::string& specializedName,
     const MappedScop& mscop) {
   // Expecting a schedule with domain root and context first child.
-  TC_CHECK(mscop.schedule()->elemAs<detail::ScheduleTreeElemDomain>());
-  TC_CHECK(
-      mscop.schedule()->child({0})->elemAs<detail::ScheduleTreeElemContext>());
+  TC_CHECK(mscop.schedule()->as<detail::ScheduleTreeDomain>());
+  TC_CHECK(mscop.schedule()->child({0})->as<detail::ScheduleTreeContext>());
   const auto& scop = mscop.scop();
 
   // Make a map of the specialized scalar parameter values

tc/core/polyhedral/cuda/mapped_scop.cc

@@ -72,7 +72,7 @@ void validate(const detail::ScheduleTree* root) {
       root);
 }
 
-bool anyNonCoincidentMember(const detail::ScheduleTreeElemBand* band) {
+bool anyNonCoincidentMember(const detail::ScheduleTreeBand* band) {
   return band->nOuterCoincident() < band->nMember();
 }
 
@@ -140,7 +140,7 @@ detail::ScheduleTree* MappedScop::map(
 detail::ScheduleTree* MappedScop::mapBlocksForward(
     detail::ScheduleTree* band,
     size_t nToMap) {
-  auto bandNode = band->elemAs<detail::ScheduleTreeElemBand>();
+  auto bandNode = band->as<detail::ScheduleTreeBand>();
   TC_CHECK(bandNode) << "expected a band, got " << *band;
 
   auto list = bandNode->mupa_.get_union_pw_aff_list();
@@ -155,7 +155,7 @@ void MappedScop::mapToBlocksAndScaleBand(
     std::vector<size_t> tileSizes) {
   using namespace tc::polyhedral::detail;
 
-  auto bandNode = band->elemAs<ScheduleTreeElemBand>();
+  auto bandNode = band->as<ScheduleTreeBand>();
   TC_CHECK(bandNode->permutable_) << "cannot map non-permutable band to blocks";
 
   auto nBlocksToMap = bandNode->nOuterCoincident();
@@ -235,7 +235,7 @@ bool MappedScop::detectReductions(detail::ScheduleTree* tree) {
   for (auto c : tree->children()) {
     found |= detectReductions(c);
   }
-  auto band = tree->elemAs<detail::ScheduleTreeElemBand>();
+  auto band = tree->as<detail::ScheduleTreeBand>();
   // Nested reductions are not currently supported.
   if (!band || found) {
     return found;
@@ -296,7 +296,7 @@ bool MappedScop::needReductionSeparation(const detail::ScheduleTree* st) {
 isl::multi_union_pw_aff MappedScop::reductionMapSchedule(
     const detail::ScheduleTree* st) {
   TC_CHECK(reductionBandUpdates_.count(st) == 1);
-  auto reductionBand = st->elemAs<detail::ScheduleTreeElemBand>();
+  auto reductionBand = st->as<detail::ScheduleTreeBand>();
   TC_CHECK(reductionBand);
 
   auto nMember = reductionBand->nMember();
@@ -359,7 +359,7 @@ detail::ScheduleTree* MappedScop::separateReduction(detail::ScheduleTree* st) {
 
 detail::ScheduleTree* MappedScop::mapThreadsBackward(
     detail::ScheduleTree* band) {
-  auto bandNode = band->elemAs<detail::ScheduleTreeElemBand>();
+  auto bandNode = band->as<detail::ScheduleTreeBand>();
   TC_CHECK(bandNode);
   auto nMember = bandNode->nMember();
   auto nToMap = std::min(nMember, numThreads.view.size());
@@ -376,7 +376,7 @@ detail::ScheduleTree* MappedScop::mapThreadsBackward(
 size_t MappedScop::mapToThreads(detail::ScheduleTree* band) {
   using namespace tc::polyhedral::detail;
 
-  auto bandNode = band->elemAs<ScheduleTreeElemBand>();
+  auto bandNode = band->as<ScheduleTreeBand>();
   // Cannot map non-permutable bands.
   if (!bandNode->permutable_) {
     return 0;
@@ -417,7 +417,7 @@ size_t MappedScop::mapToThreads(detail::ScheduleTree* band) {
       reductionBandUpdates_.erase(band);
     }
     band = child;
-    bandNode = band->elemAs<ScheduleTreeElemBand>();
+    bandNode = band->as<ScheduleTreeBand>();
   }
 
   if (nMappedThreads < bandNode->nMember()) {
@@ -450,11 +450,11 @@ bool hasOuterSequentialMember(
   auto ancestors = st->ancestors(root);
   std::reverse(ancestors.begin(), ancestors.end());
   for (auto a : ancestors) {
-    auto band = a->elemAs<detail::ScheduleTreeElemBand>();
+    auto band = a->as<detail::ScheduleTreeBand>();
     if (band && band->nMember() > band->nOuterCoincident()) {
       return true;
     }
-    if (a->elemAs<detail::ScheduleTreeElemSequence>()) {
+    if (a->as<detail::ScheduleTreeSequence>()) {
       return false;
     }
   }
@@ -542,7 +542,7 @@ Scop::SyncLevel MappedScop::findBestSync(
 
   TC_CHECK_LE(1u, scop_->scheduleRoot()->children().size());
   auto contextSt = scop_->scheduleRoot()->children()[0];
-  auto contextElem = contextSt->elemAs<detail::ScheduleTreeElemContext>();
+  auto contextElem = contextSt->as<detail::ScheduleTreeContext>();
   TC_CHECK(nullptr != contextElem);
   dependences = dependences.intersect_params(contextElem->context_);
 
@@ -705,7 +705,7 @@ std::vector<std::pair<int, int>> MappedScop::findBestSyncConfigInSeq(
 }
 
 void MappedScop::insertBestSyncInSeq(detail::ScheduleTree* seq) {
-  TC_CHECK(seq->elemAs<detail::ScheduleTreeElemSequence>());
+  TC_CHECK(seq->as<detail::ScheduleTreeSequence>());
 
   auto children = seq->children();
   auto nChildren = children.size();
@@ -779,7 +779,7 @@ size_t MappedScop::mapInnermostBandsToThreads(detail::ScheduleTree* st) {
   }
   auto n = nChildren > 0 ? *std::max_element(nInner.begin(), nInner.end()) : 0;
   if (nChildren > 1) {
-    auto needSync = st->elemAs<detail::ScheduleTreeElemSequence>() && n > 0;
+    auto needSync = st->as<detail::ScheduleTreeSequence>() && n > 0;
     if (n > 0) {
       for (size_t i = 0; i < nChildren; ++i) {
         fixThreadsBelow(*this, children[i], nInner[i]);
@@ -790,7 +790,7 @@ size_t MappedScop::mapInnermostBandsToThreads(detail::ScheduleTree* st) {
     }
   }
 
-  if (auto band = st->elemAs<detail::ScheduleTreeElemBand>()) {
+  if (auto band = st->as<detail::ScheduleTreeBand>()) {
     if (n == 0) {
       // If children were not mapped to threads, the current band can be mapped.
       // First, map the coincidence and reduction dimension to threads.
@@ -1040,7 +1040,7 @@ std::unique_ptr<MappedScop> MappedScop::makeWithOuterBlockInnerThreadStrategy(
       sharedMemorySize -= reductionMemoryRequirement;
     }
 
-    auto band = outerBand->elemAs<ScheduleTreeElemBand>();
+    auto band = outerBand->as<ScheduleTreeBand>();
     LOG_IF(WARNING, FLAGS_debug_tc_mapper && band->nMember() == 0)
         << "Aborting memory promotion because outer band has 0 members (NYI)";
     if (band->nMember() > 0 && sharedMemorySize > 0) {

tc/core/polyhedral/cuda/memory_promotion_heuristic.cc

@@ -64,7 +64,7 @@ void mapCopiesToThreads(MappedScop& mscop, bool unroll) {
     }
 
     auto bandNode = node->child({0});
-    auto band = bandNode->elemAs<ScheduleTreeElemBand>();
+    auto band = bandNode->as<ScheduleTreeBand>();
     if (!band) {
       throw promotion::PromotionLogicError("no copy band");
     }
@@ -143,7 +143,7 @@ std::vector<T> collectBranchMarkers(T root, T node) {
 isl::union_map fullSchedule(const detail::ScheduleTree* root) {
   using namespace tc::polyhedral::detail;
 
-  if (!root->elemAs<ScheduleTreeElemDomain>()) {
+  if (!root->as<ScheduleTreeDomain>()) {
     throw promotion::PromotionLogicError("expected root to be a domain node");
   }
 
@@ -157,18 +157,18 @@ isl::union_map fullSchedule(const detail::ScheduleTree* root) {
   // are innermost, the partial schedule can no longer be affected by deeper
   // nodes and hence is full.
   auto schedule = isl::union_map::empty(
-      root->elemAs<ScheduleTreeElemDomain>()->domain_.get_space());
+      root->as<ScheduleTreeDomain>()->domain_.get_space());
   for (auto node : leaves) {
-    auto domain = root->elemAs<ScheduleTreeElemDomain>()->domain_;
+    auto domain = root->as<ScheduleTreeDomain>()->domain_;
     auto prefixMupa = prefixScheduleMupa(root, node);
-    if (auto band = node->elemAs<ScheduleTreeElemBand>()) {
+    if (auto band = node->as<ScheduleTreeBand>()) {
       prefixMupa = prefixMupa.flat_range_product(band->mupa_);
     }
 
     auto pathToRoot = node->ancestors(root);
     pathToRoot.push_back(node);
     for (auto n : pathToRoot) {
-      if (auto filterNode = n->elemAs<ScheduleTreeElemFilter>()) {
+      if (auto filterNode = n->as<ScheduleTreeFilter>()) {
         domain = domain.intersect(filterNode->filter_);
       }
     }
@@ -308,7 +308,7 @@ isl::union_set collectMappingsTo(const Scop& scop) {
   mappingFilters = functional::Filter(isMappingTo<MappingType>, mappingFilters);
   auto mapping = isl::union_set::empty(domain.get_space());
   for (auto mf : mappingFilters) {
-    auto filterNode = mf->elemAs<detail::ScheduleTreeElemMapping>();
+    auto filterNode = mf->as<detail::ScheduleTreeMapping>();
     auto filter = filterNode->filter_.intersect(activeDomainPoints(root, mf));
     mapping = mapping.unite(filterNode->filter_);
   }
@@ -356,7 +356,7 @@ bool accessSubscriptsAreUnrolledLoops(
   auto leaves = functional::Filter(
       [](const ScheduleTree* tree) { return tree->numChildren() == 0; }, nodes);
 
-  auto domainNode = root->elemAs<detail::ScheduleTreeElemDomain>();
+  auto domainNode = root->as<detail::ScheduleTreeDomain>();
   TC_CHECK(domainNode);
   auto domain = domainNode->domain_;
 
@@ -368,7 +368,7 @@ bool accessSubscriptsAreUnrolledLoops(
 
     auto unrolledDims = isl::union_pw_aff_list(leaf->ctx_, 1);
     for (auto node : ancestors) {
-      auto band = node->elemAs<detail::ScheduleTreeElemBand>();
+      auto band = node->as<detail::ScheduleTreeBand>();
       if (!band) {
         continue;
       }
@@ -455,7 +455,7 @@ std::vector<detail::ScheduleTree*> bandsContainingScheduleDepth(
       ScheduleTree::collectDFSPreorder(root, detail::ScheduleTreeType::Band);
   std::function<bool(ScheduleTree * st)> containsDepth = [&](ScheduleTree* st) {
     auto depthBefore = st->scheduleDepth(root);
-    auto band = st->elemAs<ScheduleTreeElemBand>();
+    auto band = st->as<ScheduleTreeBand>();
     auto depthAfter = depthBefore + band->nMember();
     return depthBefore < depth && depthAfter >= depth;
   };
@@ -474,7 +474,7 @@ std::vector<detail::ScheduleTree*> bandsSplitAfterDepth(
 
   std::function<ScheduleTree*(ScheduleTree*)> splitAtDepth =
       [&](ScheduleTree* st) {
-        auto nMember = st->elemAs<ScheduleTreeElemBand>()->nMember();
+        auto nMember = st->as<ScheduleTreeBand>()->nMember();
         auto scheduleDepth = st->scheduleDepth(root);
         auto depthAfter = scheduleDepth + nMember;
         return depthAfter == depth ? st
@@ -641,16 +641,15 @@ void promoteGreedilyAtDepth(
 void promoteToRegistersBelow(MappedScop& mscop, detail::ScheduleTree* scope) {
   // Cannot promote below a sequence or a set node.  Promotion may insert an
   // extension node, but sequence/set must be followed by filters.
-  if (scope->elemAs<detail::ScheduleTreeElemSequence>() ||
-      scope->elemAs<detail::ScheduleTreeElemSet>()) {
+  if (scope->as<detail::ScheduleTreeSequence>() ||
+      scope->as<detail::ScheduleTreeSet>()) {
     throw promotion::IncorrectScope("cannot promote under a sequence/set node");
   }
   // Cannot promote between a thread-mapped band and a thread-specific marker
   // node because the latter is used to identify thread-mapped bands as
   // immediate ancestors.
   if (scope->numChildren() == 1 &&
-      scope->child({0})
-          ->elemAs<detail::ScheduleTreeElemThreadSpecificMarker>()) {
+      scope->child({0})->as<detail::ScheduleTreeThreadSpecificMarker>()) {
     throw promotion::IncorrectScope(
         "cannot promote above a thread-specific marker node");
   }
@@ -757,7 +756,7 @@ void promoteToRegistersAtDepth(MappedScop& mscop, size_t depth) {
   auto bands = bandsContainingScheduleDepth(root, depth);
   bands = functional::Filter(
       [root, depth](ScheduleTree* tree) {
-        auto band = tree->elemAs<ScheduleTreeElemBand>();
+        auto band = tree->as<ScheduleTreeBand>();
         return !isThreadMappedBand(tree) ||
             tree->scheduleDepth(root) + band->nMember() == depth;
       },

tc/core/polyhedral/cuda/tighten_launch_bounds.cc

@@ -77,7 +77,7 @@ size_t maxValue(const Scop& scop, const MappingIdType& id) {
   auto filters = root->collect(root, ScheduleTreeType::Mapping);
   filters = functional::Filter(isMappingTo<MappingIdType>, filters);
   for (auto p : filters) {
-    auto mappingNode = p->elemAs<ScheduleTreeElemMapping>();
+    auto mappingNode = p->as<ScheduleTreeMapping>();
     auto active = activeDomainPoints(root, p).intersect_params(params);
     active = active.intersect(mappingNode->filter_);
     auto range = rangeOfMappingParameter(active.params(), id);

tc/core/polyhedral/memory_promotion.cc

@@ -453,7 +453,7 @@ isl::multi_aff dropDummyTensorDimensions(
   return isl::multi_aff(space, list);
 }
 
-inline void unrollAllMembers(detail::ScheduleTreeElemBand* band) {
+inline void unrollAllMembers(detail::ScheduleTreeBand* band) {
   band->unroll_ = std::vector<bool>(band->nMember(), true);
 }
 
@@ -494,8 +494,8 @@ ScheduleTree* insertCopiesUnder(
   auto writeBandNode = ScheduleTree::makeBand(writeSchedule);
 
   if (unrollAllCopies) {
-    unrollAllMembers(readBandNode->elemAs<detail::ScheduleTreeElemBand>());
-    unrollAllMembers(writeBandNode->elemAs<detail::ScheduleTreeElemBand>());
+    unrollAllMembers(readBandNode->as<detail::ScheduleTreeBand>());
+    unrollAllMembers(writeBandNode->as<detail::ScheduleTreeBand>());
   }
 
   auto extension =

tc/core/polyhedral/reduction_matcher.cc

@@ -26,8 +26,8 @@ namespace tc {
 namespace polyhedral {
 
 using detail::ScheduleTree;
-using detail::ScheduleTreeElemBand;
-using detail::ScheduleTreeElemFilter;
+using detail::ScheduleTreeBand;
+using detail::ScheduleTreeFilter;
 
 namespace {