promoteToRegistersAtDepth: limit the total number of elements promoted

The limit applies per thread and is cumulated for all subtrees where
promotion is performed.  By default, it is set to SIZE_MAX, which
ensures backwards-compatible behavior for all sensible cases (if
something had required more than SIZE_MAX registers, it would have been
spilled to global memory and still would not have fit).  This limit will
be exposed as a mapping option in an upcoming commit.

Author: ftynse

tc/core/polyhedral/cuda/memory_promotion_heuristic.cc

@@ -748,9 +748,12 @@ size_t promoteToRegistersBelow(
  * Promote to registers below "depth" schedule dimensions.  Split bands if
  * necessary to create promotion scopes.  Do not promote if it would require
  * splitting the band mapped to threads as we assume only one band can be
- * mapped.
+ * mapped.  Use at most "maxElements" per thread in all promoted subtrees.
  */
-void promoteToRegistersAtDepth(MappedScop& mscop, size_t depth) {
+void promoteToRegistersAtDepth(
+    MappedScop& mscop,
+    size_t depth,
+    size_t maxElements) {
   using namespace detail;
 
   auto root = mscop.scop().scheduleRoot();
@@ -784,7 +787,7 @@ void promoteToRegistersAtDepth(MappedScop& mscop, size_t depth) {
   auto scopes = functional::Map(findScope, bands);
 
   for (auto scope : scopes) {
-    promoteToRegistersBelow(mscop, scope);
+    maxElements = promoteToRegistersBelow(mscop, scope, maxElements);
   }
 }
 

tc/core/polyhedral/cuda/memory_promotion_heuristic.h

@@ -46,7 +46,10 @@ size_t promoteToRegistersBelow(
     detail::ScheduleTree* scope,
     std::size_t maxElements = SIZE_MAX);
 
-void promoteToRegistersAtDepth(MappedScop& scop, std::size_t depth);
+void promoteToRegistersAtDepth(
+    MappedScop& scop,
+    std::size_t depth,
+    std::size_t maxElements = SIZE_MAX);
 
 } // namespace cuda
 } // namespace polyhedral