introduce single entry point for mapping to threads

This is a step towards creating a single mapping filter node
for an entire thread mapping within a branch, rather than
having separate nodes for each mapped thread identifier.
In particular, this commit consolidates the introduction
of thread mapping nodes to a single function.

Author: Sven Verdoolaege

tc/core/polyhedral/cuda/mapped_scop.cc

@@ -175,10 +175,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) {
@@ -314,6 +311,27 @@ 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 root = scop_->scheduleRoot();
+  for (size_t i = 0; i < nToMap; ++i) {
+    auto id = mapping::ThreadId::makeId(i);
+    auto pos = nMember - 1 - i;
+    band = mapToParameterWithExtent(root, band, pos, id, numThreads.view[i]);
+  }
+  mapRemaining<mapping::ThreadId>(band, nToMap);
+  return band;
+}
+
 size_t MappedScop::mapToThreads(detail::ScheduleTree* band) {
   using namespace tc::polyhedral::detail;
 
@@ -364,20 +382,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

@@ -93,14 +93,11 @@ class MappedScop {
   detail::ScheduleTree* mapBlocksForward(
       detail::ScheduleTree* band,
       size_t nToMap);
-  // Map a particular "pos"-th dimension in a _band_ node identified by "tree"
-  // to the thread dimension.  Ancestors or descendants of "tree" must
-  // not have a dimension already mapped to the same thread.
-  inline detail::ScheduleTree*
-  map(detail::ScheduleTree* tree, int pos, const mapping::ThreadId& id) {
-    return mapToParameterWithExtent(
-        scop_->scheduleRoot(), tree, pos, id, id.mappingSize(numThreads));
-  }
+  // 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);
 
   // Given that "nMapped" identifiers of type "MappingTypeId" have already
   // been mapped, map the remaining ones to zero

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) {

test/test_cuda_mapper.cc

@@ -86,11 +86,9 @@ struct PolyhedralMapperTest : public ::testing::Test {
     auto band = mscop->mapBlocksForward(root->child({0}), 1);
     bandScale(band, tileSizes);
 
-    USING_MAPPING_SHORT_NAMES(BX, BY, BZ, TX, TY, TZ);
     auto ns = detail::ScheduleTree::collectDFSPostorder(
         root, detail::ScheduleTreeType::Band);
-    mscop->map(ns[1], 1, TX);
-    mscop->map(ns[1], 0, TY);
+    mscop->mapThreadsBackward(ns[1]);
     mscop->insertMappingContext();
     return mscop;
   }
@@ -113,9 +111,7 @@ struct PolyhedralMapperTest : public ::testing::Test {
     auto band = mscop->mapBlocksForward(root->child({0}), 2);
     bandScale(band, tileSizes);
 
-    USING_MAPPING_SHORT_NAMES(BX, BY, BZ, TX, TY, TZ);
-    band = mscop->map(band->child({0}), 1, TX);
-    band = mscop->map(band, 0, TY);
+    band = mscop->mapThreadsBackward(band->child({0}));
     mscop->insertMappingContext();
     return mscop;
   }