register promotion: use promoteToRegisterAtDepth

Introduce a mapping option "privateDepth" to control the schedule depth
at which register promotion is attempted.  Pass this option to
promoteToRegisterAtDepth instead of calling
promoteToRegistersBelowThreads in the mapping strategy.

Author: ftynse

tc/core/cuda/cuda_mapping_options.cc

@@ -289,6 +289,11 @@ CudaMappingOptions& CudaMappingOptions::useReadOnlyCache(bool b) {
   return *this;
 }
 
+CudaMappingOptions& CudaMappingOptions::privateDepth(uint32_t depth) {
+  ownedProto_.set_private_depth(depth);
+  return *this;
+}
+
 CudaMappingOptions& CudaMappingOptions::mapToThreads(
     const std::string& commaSeparatedSizes) {
   auto sizes = parseCommaSeparatedIntegers<uint64_t>(commaSeparatedSizes);

tc/core/cuda/cuda_mapping_options.h

@@ -195,6 +195,7 @@ class CudaMappingOptions {
   CudaMappingOptions& maxSharedMemory(uint64_t size);
   CudaMappingOptions& unrollCopyShared(bool b);
   CudaMappingOptions& useReadOnlyCache(bool b);
+  CudaMappingOptions& privateDepth(uint32_t depth);
   ///@}
 
   /// Static constructors for predefined strategies.

tc/core/cuda/cuda_mapping_options_cpp_printer.cc

@@ -38,6 +38,7 @@ CudaMappingOptionsCppPrinter& operator<<(
     prn.printValueOption(
         "maxSharedMemory", cudaOptions.proto().max_shared_memory());
   }
+  prn.printValueOption("privateDepth", cudaOptions.proto().private_depth());
   prn.endStmt();
   return prn;
 }

tc/core/polyhedral/cuda/mapped_scop.cc

@@ -1068,7 +1068,7 @@ std::unique_ptr<MappedScop> MappedScop::makeWithOuterBlockInnerThreadStrategy(
 
   // 9. Promote to registers below the loops mapped to threads.
   if (cudaOptions.proto().use_private_memory()) {
-    promoteToRegistersBelowThreads(*mappedScop, -1ull);
+    promoteToRegistersAtDepth(*mappedScop, cudaOptions.proto().private_depth());
   }
 
   LOG_IF(INFO, FLAGS_debug_tc_mapper)

tc/core/polyhedral/cuda/memory_promotion_heuristic.cc

@@ -779,16 +779,5 @@ void promoteToRegistersAtDepth(MappedScop& mscop, size_t depth) {
   }
 }
 
-// Promote at the positions of the thread specific markers.
-void promoteToRegistersBelowThreads(MappedScop& mscop, size_t nRegisters) {
-  auto& scop = mscop.scop();
-  auto root = scop.scheduleRoot();
-  auto markers = findThreadSpecificMarkers(root);
-
-  for (auto marker : markers) {
-    promoteToRegistersBelow(mscop, marker);
-  }
-}
-
 } // namespace polyhedral
 } // namespace tc

tc/core/polyhedral/cuda/memory_promotion_heuristic.h

@@ -44,7 +44,6 @@ void promoteGreedilyAtDepth(
 
 void promoteToRegistersBelow(MappedScop& mscop, detail::ScheduleTree* scope);
 
-void promoteToRegistersBelowThreads(MappedScop& scop, std::size_t nRegisters);
 void promoteToRegistersAtDepth(MappedScop& scop, std::size_t depth);
 
 } // namespace polyhedral

tc/proto/mapping_options.proto

@@ -70,6 +70,8 @@ message CudaMappingOptionsProto {
   optional uint64 max_shared_memory = 7;
   // Use the readonly cache (i.e. emit __ldg loads)
   required bool use_readonly_cache = 8;
+  // Depth of promotion to private memory, ignored if use_private_memory is false.
+  optional uint32 private_depth = 9;
 }
 
 message CpuMappingOptionsProto {

test/test_cuda_mapper_memory_promotion.cc

@@ -531,8 +531,12 @@ def fun(float(N,K) A, float(K,M) B, float(N,M) C) -> (O) {
 };
 
 TEST_F(MatMulBias, RegisterPromotion) {
+  // Scheduled code has three loops, tile all, the top band gets mapped to 2
+  // blocks, the bottom band gets mapped to 2 threads, test promotion before
+  // thread-mapped loops, that is at depth 3 + 2 = 5.
   auto mappingOptions = CudaMappingOptions::makeNaiveMappingOptions()
                             .tile(32, 32, 32)
+                            .privateDepth(5)
                             .useSharedMemory(false)
                             .usePrivateMemory(true);