Merge pull request #360 from facebookresearch/pr/single

create a single block mapping filter and a single thread mapping filter per branch

Author: ftynse

tc/core/polyhedral/cuda/mapped_scop.cc

@@ -108,22 +108,69 @@ const CudaDim& mappingSize<mapping::ThreadId>(const MappedScop* mscop) {
 }
 } // namespace
 
+// Map the elements in "list" to successive blocks or thread identifiers,
+// with the first element mapped to identifier X.  The extents are obtained
+// from the initial elements of numBlocks or numThreads.  The identifiers
+// must not be present in the space of the partial schedules in "list" and
+// extents must be non-zero.  The mapping corresponds to inserting a filter
+// node with condition 'list % extent = ids'.
+// The mapping is inserted above "tree".
+//
+// Return a pointer to the updated node (below the inserted filter)
+// for call chaining purposes.
 template <typename MappingTypeId>
-void MappedScop::mapRemaining(detail::ScheduleTree* tree, size_t nMapped) {
-  size_t nToMap = mappingSize<MappingTypeId>(this).view.size();
-  if (nMapped >= nToMap) {
-    return;
+detail::ScheduleTree* MappedScop::map(
+    detail::ScheduleTree* tree,
+    isl::union_pw_aff_list list) {
+  size_t nToMap = list.n();
+  const auto& extent = mappingSize<MappingTypeId>(this).view;
+  CHECK_LE(nToMap, extent.size()) << "dimension overflow";
+
+  auto root = scop_->scheduleRoot();
+  auto domain = activeDomainPoints(root, tree).universe();
+  auto filter = domain;
+
+  std::unordered_set<MappingTypeId, typename MappingTypeId::Hash> idSet;
+  for (size_t i = 0; i < nToMap; ++i) {
+    auto id = MappingTypeId::makeId(i);
+    auto upa = list.get(i);
+    // Introduce the "mapping" parameter after checking it is not already
+    // present in the schedule space.
+    CHECK(not upa.involves_param(id));
+    CHECK_NE(extent[i], 0u) << "NYI: mapping to 0";
+
+    // Create mapping filter by equating the newly introduced
+    // parameter ids[i] to the "i"-th affine function modulo its extent.
+    upa = upa.mod_val(isl::val(tree->ctx_, extent[i]));
+    upa = upa.sub(isl::union_pw_aff::param_on_domain(domain, id));
+    filter = filter.intersect(upa.zero_union_set());
+
+    idSet.emplace(id);
   }
 
-  std::unordered_set<MappingTypeId, typename MappingTypeId::Hash> ids;
-  for (size_t i = nMapped; i < nToMap; ++i) {
-    ids.insert(MappingTypeId::makeId(i));
+  std::unordered_set<MappingTypeId, typename MappingTypeId::Hash> unmapped;
+  for (size_t i = nToMap; i < extent.size(); ++i) {
+    auto id = MappingTypeId::makeId(i);
+    unmapped.emplace(id);
+    idSet.emplace(id);
   }
-  auto root = scop_->scheduleRoot();
-  auto domain = activeDomainPoints(root, tree);
-  auto filter = makeFixRemainingZeroFilter(domain, ids);
-  auto mapping = detail::ScheduleTree::makeMappingFilter(filter, ids);
-  insertNodeAbove(root, tree, std::move(mapping));
+  filter = filter.intersect(makeFixRemainingZeroFilter(domain, unmapped));
+
+  auto mapping = detail::ScheduleTree::makeMappingFilter(filter, idSet);
+  tree = insertNodeAbove(root, tree, std::move(mapping))->child({0});
+
+  return tree;
+}
+
+detail::ScheduleTree* MappedScop::mapBlocksForward(
+    detail::ScheduleTree* band,
+    size_t nToMap) {
+  auto bandNode = band->elemAs<detail::ScheduleTreeElemBand>();
+  CHECK(bandNode) << "expected a band, got " << *band;
+
+  auto list = bandNode->mupa_.get_union_pw_aff_list();
+  list = list.drop(nToMap, list.n() - nToMap);
+  return map<mapping::BlockId>(band, list);
 }
 
 // Uses as many blockSizes elements as outer coincident dimensions in the
@@ -142,10 +189,7 @@ void MappedScop::mapToBlocksAndScaleBand(
   // and no more than block dimensions to be mapped
   nBlocksToMap = std::min(nBlocksToMap, numBlocks.view.size());
 
-  for (size_t i = 0; i < nBlocksToMap; ++i) {
-    band = map(band, i, mapping::BlockId::makeId(i));
-  }
-  mapRemaining<mapping::BlockId>(band, nBlocksToMap);
+  mapBlocksForward(band, nBlocksToMap);
   bandScale(band, tileSizes);
 }
 
@@ -166,10 +210,7 @@ void fixThreadsBelow(
 
   auto band = detail::ScheduleTree::makeEmptyBand(mscop.scop().scheduleRoot());
   auto bandTree = insertNodeBelow(tree, std::move(band));
-  auto ctx = tree->ctx_;
-  insertNodeBelow(
-      bandTree, detail::ScheduleTree::makeThreadSpecificMarker(ctx));
-  mscop.mapRemaining<mapping::ThreadId>(bandTree, begin);
+  mscop.mapThreadsBackward(bandTree);
 }
 
 bool MappedScop::detectReductions(detail::ScheduleTree* tree) {
@@ -305,6 +346,22 @@ detail::ScheduleTree* MappedScop::separateReduction(detail::ScheduleTree* st) {
   return st->ancestor(root, 2);
 }
 
+detail::ScheduleTree* MappedScop::mapThreadsBackward(
+    detail::ScheduleTree* band) {
+  auto bandNode = band->elemAs<detail::ScheduleTreeElemBand>();
+  CHECK(bandNode);
+  auto nMember = bandNode->nMember();
+  auto nToMap = std::min(nMember, numThreads.view.size());
+  CHECK_LE(nToMap, 3) << "mapping to too many threads";
+
+  auto ctx = band->ctx_;
+  insertNodeBelow(band, detail::ScheduleTree::makeThreadSpecificMarker(ctx));
+
+  auto list = bandNode->mupa_.get_union_pw_aff_list().reverse();
+  list = list.drop(nToMap, list.n() - nToMap);
+  return map<mapping::ThreadId>(band, list);
+}
+
 size_t MappedScop::mapToThreads(detail::ScheduleTree* band) {
   using namespace tc::polyhedral::detail;
 
@@ -355,20 +412,9 @@ size_t MappedScop::mapToThreads(detail::ScheduleTree* band) {
     bandSplit(scop_->scheduleRoot(), band, nMappedThreads);
   }
 
-  auto ctx = band->ctx_;
-  insertNodeBelow(band, detail::ScheduleTree::makeThreadSpecificMarker(ctx));
-
   CHECK_GT(nMappedThreads, 0) << "not mapping to threads";
-  CHECK_LE(nMappedThreads, 3) << "mapping to too many threads";
 
-  // Map the coincident dimensions to threads starting from the innermost and
-  // from thread x.
-  for (size_t i = 0; i < nMappedThreads; ++i) {
-    auto id = mapping::ThreadId::makeId(i);
-    auto dim = nMappedThreads - 1 - i;
-    band = map(band, dim, id);
-  }
-  mapRemaining<mapping::ThreadId>(band, nMappedThreads);
+  mapThreadsBackward(band);
 
   if (isReduction) {
     splitOutReductionAndInsertSyncs(band, nMappedThreads - 1);

tc/core/polyhedral/cuda/mapped_scop.h

@@ -87,25 +87,17 @@ class MappedScop {
       std::unique_ptr<Scop>&& scopUPtr,
       const CudaMappingOptions& mappingOptions);
 
-  // Map a particular "pos"-th dimension in a _band_ node identified by "tree"
-  // to the block or thread dimension.  Ancestors or descendants of "tree" must
-  // not have a dimension already mapped to the same block or thread.
-  inline detail::ScheduleTree*
-  map(detail::ScheduleTree* tree, int pos, const mapping::BlockId& id) {
-    return mapToParameterWithExtent(
-        scop_->scheduleRoot(), tree, pos, id, id.mappingSize(numBlocks));
-  }
-  inline detail::ScheduleTree*
-  map(detail::ScheduleTree* tree, int pos, const mapping::ThreadId& id) {
-    return mapToParameterWithExtent(
-        scop_->scheduleRoot(), tree, pos, id, id.mappingSize(numThreads));
-  }
-
-  // Given that "nMapped" identifiers of type "MappingTypeId" have already
-  // been mapped, map the remaining ones to zero
-  // for all statement instances.
-  template <typename MappingTypeId>
-  void mapRemaining(detail::ScheduleTree* tree, size_t nMapped);
+  // Map the initial (up to "nToMap") band members of "band"
+  // to successive block identifiers.
+  // This function can only be called once on the entire tree.
+  detail::ScheduleTree* mapBlocksForward(
+      detail::ScheduleTree* band,
+      size_t nToMap);
+  // Map the final band members of "band"
+  // to successive thread identifiers, with the last member mapped
+  // to thread identifier X.
+  // This function can only be called once in any branch of the tree.
+  detail::ScheduleTree* mapThreadsBackward(detail::ScheduleTree* band);
 
   // Fix the values of the specified parameters in the context
   // to the corresponding specified values.
@@ -136,6 +128,14 @@ class MappedScop {
   }
 
  private:
+  // Map the elements in "list" to successive blocks or thread identifiers,
+  // with the first element mapped to identifier X.
+  // Return a pointer to the updated node (below the inserted filter)
+  // for call chaining purposes.
+  template <typename MappingTypeId>
+  detail::ScheduleTree* map(
+      detail::ScheduleTree* tree,
+      isl::union_pw_aff_list list);
   // Map "band" to block identifiers and then scale
   // the band members by "tileSizes".
   void mapToBlocksAndScaleBand(

tc/core/polyhedral/cuda/memory_promotion_heuristic.cc

@@ -83,10 +83,6 @@ void mapCopiesToThreads(MappedScop& mscop, bool unroll) {
       throw promotion::PromotionLogicError("no copy band");
     }
 
-    auto ctx = node->ctx_;
-    insertNodeBelow(
-        bandNode, detail::ScheduleTree::makeThreadSpecificMarker(ctx));
-
     // Check that we are not mapping to threads below other thread mappings.
     std::unordered_set<mapping::ThreadId, mapping::ThreadId::Hash> usedThreads;
     for (auto n : node->ancestors(root)) {
@@ -97,20 +93,7 @@ void mapCopiesToThreads(MappedScop& mscop, bool unroll) {
       }
     }
 
-    // Map band dimensions to threads, in inverse order since the last member
-    // iterates over the last subscript and is likely to result in coalescing.
-    // If not all available thread ids are used, fix remaining to 1 thread.
-    auto nToMap = std::min(band->nMember(), mscop.numThreads.view.size());
-    for (size_t t = 0; t < nToMap; ++t) {
-      auto pos = band->nMember() - 1 - t;
-      mapToParameterWithExtent(
-          root,
-          bandNode,
-          pos,
-          mapping::ThreadId::makeId(t),
-          mscop.numThreads.view[t]);
-    }
-    mscop.mapRemaining<mapping::ThreadId>(bandNode, nToMap);
+    mscop.mapThreadsBackward(bandNode);
 
     // Unroll if requested.
     if (unroll) {

tc/core/polyhedral/schedule_transforms-inl.h

@@ -39,35 +39,5 @@ inline detail::ScheduleTree* insertNodeBelow(
   tree->appendChild(std::move(node));
   return tree->child({0});
 }
-
-template <typename MappingIdType>
-inline detail::ScheduleTree* mapToParameterWithExtent(
-    detail::ScheduleTree* root,
-    detail::ScheduleTree* tree,
-    size_t pos,
-    MappingIdType id,
-    size_t extent) {
-  auto band = tree->elemAs<detail::ScheduleTreeElemBand>();
-  CHECK(band) << "expected a band, got " << *tree;
-  CHECK_GE(pos, 0u) << "dimension underflow";
-  CHECK_LT(pos, band->nMember()) << "dimension overflow";
-  CHECK_NE(extent, 0u) << "NYI: mapping to 0";
-
-  auto domain = activeDomainPoints(root, tree).universe();
-
-  // Introduce the "mapping" parameter after checking it is not already present
-  // in the schedule space.
-  CHECK(not band->mupa_.involves_param(id));
-
-  // Create mapping filter by equating the newly introduced
-  // parameter "id" to the "pos"-th schedule dimension modulo its extent.
-  auto upa =
-      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();
-  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

@@ -754,52 +754,5 @@ void orderAfter(ScheduleTree* root, ScheduleTree* tree, isl::union_set filter) {
   parent->insertChild(childPos, std::move(seq));
 }
 
-detail::ScheduleTree* mergeConsecutiveMappingFilters(
-    detail::ScheduleTree* root,
-    detail::ScheduleTree* node) {
-  CHECK(
-      root->elemAs<ScheduleTreeElemDomain>() ||
-      root->elemAs<ScheduleTreeElemExtension>());
-  bool changed = true;
-  while (changed) {
-    changed = false;
-    auto filterNodes = detail::ScheduleTree::collect(
-        node, detail::ScheduleTreeType::MappingFilter);
-
-    for (auto f : filterNodes) {
-      auto p = f->ancestor(root, 1);
-      auto parentFilter = p->elemAs<ScheduleTreeElemMappingFilter>();
-      if (!parentFilter) {
-        continue;
-      }
-      auto filter = f->elemAs<ScheduleTreeElemMappingFilter>();
-      auto merged = parentFilter->filter_ & filter->filter_;
-      // We can only merge filters that have the same number of tuples
-      if (merged.n_set() != parentFilter->filter_.n_set() ||
-          merged.n_set() != filter->filter_.n_set()) {
-        continue;
-      }
-      p->elemAs<ScheduleTreeElemMappingFilter>()->filter_ = merged;
-      // const cast to replace in place rather than construct a new
-      // ScheduleTree object (which would not be more functional-style anyway)
-      auto& ids = const_cast<std::unordered_set<
-          mapping::MappingId,
-          typename mapping::MappingId::Hash>&>(
-          p->elemAs<ScheduleTreeElemMappingFilter>()->mappingIds);
-      for (auto id : filter->mappingIds) {
-        CHECK_EQ(0u, ids.count(id))
-            << "Error when merging filters\n"
-            << *f << "\nand\n"
-            << *p << "\nid: " << id << " mapped in both!";
-        ids.insert(id);
-      }
-      p->replaceChild(f->positionInParent(p), f->detachChild(0));
-      changed = true;
-      break;
-    }
-  }
-  return node;
-}
-
 } // namespace polyhedral
 } // namespace tc

tc/core/polyhedral/schedule_transforms.h

@@ -114,23 +114,6 @@ detail::ScheduleTree* bandScale(
     detail::ScheduleTree* tree,
     const std::vector<size_t>& scales);
 
-// Map "pos"-th schedule dimension of the band node identified by "tree" to a
-// _new_ parameter identified by "id" and limited by 0 <= id < extent.  The
-// parameter must not be present in the space of partial schedule of "tree" and
-// extent must be non-zero.  The mapping corresponds to inserting a filter
-// node with condition 'dim % extent = id' where dim is "pos"-th
-// schedule dimension.
-//
-// Returns a pointer to the updated band (below the inserted filter)
-// for call chaining purposes.
-template <typename MappingIdType>
-detail::ScheduleTree* mapToParameterWithExtent(
-    detail::ScheduleTree* root,
-    detail::ScheduleTree* tree,
-    size_t pos,
-    MappingIdType id,
-    size_t extent);
-
 // 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.
@@ -314,15 +297,6 @@ isl::union_set activeDomainPointsBelow(
     const detail::ScheduleTree* root,
     const detail::ScheduleTree* node);
 
-////////////////////////////////////////////////////////////////////////////////
-// Experimental
-////////////////////////////////////////////////////////////////////////////////
-// Mapping filters are introduced one mapping dimension at a time.
-// This merges consecutive filters.
-detail::ScheduleTree* mergeConsecutiveMappingFilters(
-    detail::ScheduleTree* root,
-    detail::ScheduleTree* node);
-
 } // namespace polyhedral
 } // namespace tc