insertCopiesUnder: forcibly unroll copy loops for register promotion

When copying to registers, loop unrolling is mandatory.  Otherwise,
promoted arrays will be accessed with non-constant subscripts, which
prevents CUDA from placing them to registers as they are not indirectly
addressable.

Request unrolling when inserting copies to registers.  This has not been
an issue until now since we only promote groups of 1 element to
registers.  This will be relaxed in the following commits.

Author: ftynse

tc/core/polyhedral/memory_promotion.cc

@@ -459,14 +459,20 @@ isl::multi_aff dropDummyTensorDimensions(
   space = add_range(space, list.n());
   return isl::multi_aff(space, list);
 }
+
+inline void unrollAllMembers(detail::ScheduleTreeElemBand* band) {
+  band->unroll_ = std::vector<bool>(band->nMember(), true);
+}
+
 } // namespace
 
 ScheduleTree* insertCopiesUnder(
     Scop& scop,
     ScheduleTree* tree,
     const TensorReferenceGroup& group,
     isl::id tensorId,
-    isl::id groupId) {
+    isl::id groupId,
+    bool unrollAllCopies) {
   const ScheduleTree* root = scop.scheduleRoot();
   auto ctx = root->ctx_;
   isl::id readId = isl::id(ctx, std::string(kReadIdName));
@@ -494,6 +500,11 @@ ScheduleTree* insertCopiesUnder(
   auto readBandNode = ScheduleTree::makeBand(readSchedule);
   auto writeBandNode = ScheduleTree::makeBand(writeSchedule);
 
+  if (unrollAllCopies) {
+    unrollAllMembers(readBandNode->elemAs<detail::ScheduleTreeElemBand>());
+    unrollAllMembers(writeBandNode->elemAs<detail::ScheduleTreeElemBand>());
+  }
+
   auto extension =
       promotion.wrap().identity().domain_factor_domain().domain_factor_domain();
 

tc/core/polyhedral/memory_promotion.h

@@ -214,6 +214,7 @@ detail::ScheduleTree* insertCopiesUnder(
     detail::ScheduleTree* tree,
     const TensorReferenceGroup& group,
     isl::id tensorId,
-    isl::id groupId);
+    isl::id groupId,
+    bool unrollAllCopies);
 } // namespace polyhedral
 } // namespace tc

tc/core/polyhedral/scop.cc

@@ -192,7 +192,13 @@ void Scop::promoteGroup(
     sizes.back() += 1;
   }
   promotedDecls_[groupId] = PromotedDecl{tensorId, sizes, kind};
-  insertCopiesUnder(*this, tree, *gr, tensorId, groupId);
+  insertCopiesUnder(
+      *this,
+      tree,
+      *gr,
+      tensorId,
+      groupId,
+      kind == PromotedDecl::Kind::Register);
 
   // FIXME: we can now store a unique pointer...
   auto group = std::shared_ptr<TensorReferenceGroup>(std::move(gr));