promoteToSharedAtDepth: ignore scopes below thread mapping

Band nodes at the given depth may be above or below the thread mapping.
In the latter case, promoteToSharedBelow would throw an IncorrectScope
exception because promotion to shared below a thread mapping node is
forbidden and is generally nonsensical.  In a branching tree, there is
no guarantee that nodes at the given schedule depth are all above or
below the thread mapping.  Simply ignore the subtrees that have a
thread-mapping ancestor node when performing shared memory promotion at
depth.

Leaving promoteToSharedBelow as is (throwing an exception) because it
may end up being reused by different callers.  Arguably,
promoteToSharedAtDepth is a higher-level API and should not throw if
only one subtree is problematic.

Author: ftynse

tc/core/polyhedral/cuda/memory_promotion_heuristic.cc

@@ -565,7 +565,8 @@ inline bool isThreadMappedBand(const detail::ScheduleTree* tree) {
 /*
  * For every place in the schedule tree where schedule depth (i.e., the number
  * of preceding band members) is "depth", promote tensor reference groups to
- * shared memory.  Split bands if necessary to insert promotions.
+ * shared memory if there is no thread mapping above this place.  Split bands
+ * if necessary to insert promotions.
  *
  * Use at most "maxMemory" bytes.  If a groups does not fit the remaining
  * memory, do not promote it and keep looking for a smaller group.
@@ -600,9 +601,16 @@ void promoteToSharedAtDepth(
   // immediately below it in the tree.  In particular, promote if the
   // approximated footprint fits into the remaining memory, and the reference
   // group either features reuse or is accessed in a non-coalesced way, or
-  // both.
+  // both.  Do not promote if the band node is located below the thread mapping
+  // as promotion to shared is not allowed in this context.
   size_t remainingMemory = maxMemory;
   for (auto bandNode : bands) {
+    if (isInThreadMappedScope(root, bandNode)) {
+      LOG_IF(INFO, FLAGS_debug_tc_mapper)
+          << "not promoting subtree to shared because it is below "
+          << "a thread mapping node";
+      continue;
+    }
     promoteToSharedBelow(scop, bandNode, remainingMemory);
   }
 

test/test_cuda_mapper_memory_promotion.cc

@@ -437,12 +437,14 @@ TEST_F(MapperMemoryPromotionRAW, fitAtOuterDepths) {
       << "expected one reference group to be promoted";
 }
 
-TEST_F(MapperMemoryPromotionRAW, throwIfCopiesBelowThreads) {
-  EXPECT_THROW(
-      makeWithSharedGreedy(42, 40, 64, 64, 3, 8192), promotion::IncorrectScope);
-
-  EXPECT_THROW(
-      makeWithSharedGreedy(42, 40, 64, 64, 4, 8192), promotion::IncorrectScope);
+TEST_F(MapperMemoryPromotionRAW, noSharedPromotionBelowThreads) {
+  auto mscop1 = makeWithSharedGreedy(42, 40, 64, 64, 3, 8192);
+  EXPECT_EQ(mscop1->scop().promotedDecls().size(), 0u)
+      << "expected no promotion below threads";
+
+  auto mscop2 = makeWithSharedGreedy(42, 40, 64, 64, 4, 8192);
+  EXPECT_EQ(mscop2->scop().promotedDecls().size(), 0u)
+      << "expected no promotion below threads";
 }
 
 class MatMulBias : public TestMapper {