ScheduleTree: rename "elemAs" to "as"

The preceding commits removed the notion of schedule tree element (an
object contained in the tree nodes), making specific tree node types
inherit from a generic one.  Rename the function to make it clear to the
caller that it casts the object it is being called on.

Author: ftynse

tc/core/polyhedral/codegen.cc

@@ -31,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::ScheduleTreeElemBand>();
     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::ScheduleTreeElemDomain>());
+  TC_CHECK(mscop.schedule()->child({0})->as<detail::ScheduleTreeElemContext>());
   const auto& scop = mscop.scop();
 
   // Make a map of the specialized scalar parameter values

tc/core/polyhedral/cuda/mapped_scop.cc

@@ -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::ScheduleTreeElemBand>();
   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<ScheduleTreeElemBand>();
   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::ScheduleTreeElemBand>();
   // 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::ScheduleTreeElemBand>();
   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::ScheduleTreeElemBand>();
   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<ScheduleTreeElemBand>();
   // 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<ScheduleTreeElemBand>();
   }
 
   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::ScheduleTreeElemBand>();
     if (band && band->nMember() > band->nOuterCoincident()) {
       return true;
     }
-    if (a->elemAs<detail::ScheduleTreeElemSequence>()) {
+    if (a->as<detail::ScheduleTreeElemSequence>()) {
       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::ScheduleTreeElemContext>();
   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::ScheduleTreeElemSequence>());
 
   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::ScheduleTreeElemSequence>() && 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::ScheduleTreeElemBand>()) {
     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<ScheduleTreeElemBand>();
     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<ScheduleTreeElemBand>();
     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<ScheduleTreeElemDomain>()) {
     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<ScheduleTreeElemDomain>()->domain_.get_space());
   for (auto node : leaves) {
-    auto domain = root->elemAs<ScheduleTreeElemDomain>()->domain_;
+    auto domain = root->as<ScheduleTreeElemDomain>()->domain_;
     auto prefixMupa = prefixScheduleMupa(root, node);
-    if (auto band = node->elemAs<ScheduleTreeElemBand>()) {
+    if (auto band = node->as<ScheduleTreeElemBand>()) {
       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<ScheduleTreeElemFilter>()) {
         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::ScheduleTreeElemMapping>();
     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::ScheduleTreeElemDomain>();
   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::ScheduleTreeElemBand>();
       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<ScheduleTreeElemBand>();
     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<ScheduleTreeElemBand>()->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::ScheduleTreeElemSequence>() ||
+      scope->as<detail::ScheduleTreeElemSet>()) {
     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::ScheduleTreeElemThreadSpecificMarker>()) {
     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<ScheduleTreeElemBand>();
         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<ScheduleTreeElemMapping>();
     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

@@ -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::ScheduleTreeElemBand>());
+    unrollAllMembers(writeBandNode->as<detail::ScheduleTreeElemBand>());
   }
 
   auto extension =

tc/core/polyhedral/schedule_isl_conversion.cc

@@ -48,11 +48,11 @@ isl::schedule_node insertBranch(
     const std::vector<size_t>& pos) {
   auto filters = isl::union_set_list(node.get_ctx(), st->numChildren());
   for (size_t i = 0; i < pos.size(); ++i) {
-    auto filter = st->child({pos[i]})->elemAs<ScheduleTreeElemFilter>();
+    auto filter = st->child({pos[i]})->as<ScheduleTreeElemFilter>();
     TC_CHECK(filter);
     filters = filters.add(filter->filter_);
   }
-  if (st->elemAs<ScheduleTreeElemSet>()) {
+  if (st->as<ScheduleTreeElemSet>()) {
     node = node.insert_set(filters);
   } else {
     node = node.insert_sequence(filters);
@@ -83,9 +83,9 @@ std::vector<size_t> findCorePositions(
     const ScheduleTree* st,
     isl::union_set domain) {
   std::vector<size_t> positions;
-  TC_CHECK(st->elemAs<ScheduleTreeElemSequence>());
+  TC_CHECK(st->as<ScheduleTreeElemSequence>());
   for (size_t i = 0; i < st->numChildren(); ++i) {
-    auto filter = st->child({i})->elemAs<ScheduleTreeElemFilter>();
+    auto filter = st->child({i})->as<ScheduleTreeElemFilter>();
     TC_CHECK(filter);
     if (!filter->filter_.intersect(domain).is_empty()) {
       positions.emplace_back(i);
@@ -103,7 +103,7 @@ std::vector<size_t> findCorePositions(
 isl::schedule_node graftFromFilterSubtree(
     const ScheduleTree* st,
     isl::union_map extension) {
-  auto filter = st->elemAs<ScheduleTreeElemFilter>();
+  auto filter = st->as<ScheduleTreeElemFilter>();
   TC_CHECK(filter);
   auto filterExtension = extension.intersect_range(filter->filter_);
   auto extensionNode = isl::schedule_node::from_extension(filterExtension);
@@ -131,7 +131,7 @@ isl::schedule_node insertExtension(
   TC_CHECK(!corePos.empty());
   node = insertBranch(node, child, corePos);
 
-  auto extension = st->elemAs<ScheduleTreeElemExtension>()->extension_;
+  auto extension = st->as<ScheduleTreeElemExtension>()->extension_;
   for (size_t i = 0; i < corePos.size() - 1; ++i) {
     auto depth0 = node.get_tree_depth();
     node = node.child(i).child(0);
@@ -162,7 +162,7 @@ isl::schedule_node insertExtension(
  * some extra functionality in isl.
  */
 isl::schedule_node insert(isl::schedule_node node, const ScheduleTree* st) {
-  if (auto band = st->elemAs<ScheduleTreeElemBand>()) {
+  if (auto band = st->as<ScheduleTreeElemBand>()) {
     node = node.insert_partial_schedule(band->mupa_);
     auto bandNode = node.as<isl::schedule_node_band>();
     bandNode = bandNode.set_permutable(band->permutable_);
@@ -176,19 +176,18 @@ isl::schedule_node insert(isl::schedule_node node, const ScheduleTree* st) {
       }
     }
     node = bandNode;
-  } else if (auto context = st->elemAs<ScheduleTreeElemContext>()) {
+  } else if (auto context = st->as<ScheduleTreeElemContext>()) {
     node = node.insert_context(context->context_);
-  } else if (auto filter = st->elemAs<ScheduleTreeElemFilter>()) {
+  } else if (auto filter = st->as<ScheduleTreeElemFilter>()) {
     node = node.insert_filter(filter->filter_);
-  } else if (auto filter = st->elemAs<ScheduleTreeElemMapping>()) {
+  } else if (auto filter = st->as<ScheduleTreeElemMapping>()) {
     node = node.insert_filter(filter->filter_);
   } else if (
-      st->elemAs<ScheduleTreeElemSet>() ||
-      st->elemAs<ScheduleTreeElemSequence>()) {
+      st->as<ScheduleTreeElemSet>() || st->as<ScheduleTreeElemSequence>()) {
     return insertBranch(node, st);
-  } else if (st->elemAs<ScheduleTreeElemExtension>()) {
+  } else if (st->as<ScheduleTreeElemExtension>()) {
     return insertExtension(node, st);
-  } else if (st->elemAs<ScheduleTreeElemThreadSpecificMarker>()) {
+  } else if (st->as<ScheduleTreeElemThreadSpecificMarker>()) {
     return insertChild(node, st);
   } else {
     LOG(FATAL) << "NYI: insert type: " << *st;
@@ -228,7 +227,7 @@ isl::schedule_node extendChild(
  * then recursively add nodes corresponding to the descendants of "root".
  */
 isl::schedule toIslSchedule(const ScheduleTree* root) {
-  auto domain = root->elemAs<ScheduleTreeElemDomain>();
+  auto domain = root->as<ScheduleTreeElemDomain>();
   TC_CHECK(domain) << "Root node should be domain node" << *root;
   auto node = isl::schedule_node::from_domain(domain->domain_);
   node = extendChild(node, root);