MappedScop::reductionMapSchedule: extract out ScheduleTreeElemBand::memberRange

This function will be reused in the next commit.

Author: Sven Verdoolaege

tc/core/polyhedral/cuda/mapped_scop.cc

@@ -299,19 +299,13 @@ isl::multi_union_pw_aff MappedScop::reductionMapSchedule(
   auto reductionBand = st->elemAs<detail::ScheduleTreeElemBand>();
   TC_CHECK(reductionBand);
 
-  // Drop band members following the reduction dimension and preceding those
-  // mapped to threads.
-  auto reductionSchedule = reductionBand->mupa_;
   auto nMember = reductionBand->nMember();
   auto reductionDim = reductionBand->nOuterCoincident();
   auto nMappedThreads = numThreads.view.size();
   TC_CHECK_GE(nMember, reductionDim + 1);
-  reductionSchedule = reductionSchedule.drop_dims(
-      isl::dim_type::set, reductionDim + 1, nMember - (reductionDim + 1));
-  reductionSchedule = reductionSchedule.drop_dims(
-      isl::dim_type::set, 0, reductionDim - nMappedThreads + 1);
 
-  return reductionSchedule;
+  auto first = reductionDim + 1 - nMappedThreads;
+  return reductionBand->memberRange(first, nMappedThreads);
 }
 
 detail::ScheduleTree* MappedScop::separateReduction(detail::ScheduleTree* st) {

tc/core/polyhedral/schedule_tree_elem.cc

@@ -175,6 +175,18 @@ void ScheduleTreeElemBand::drop(size_t pos, size_t n) {
   TC_CHECK_EQ(nBegin - n, nMember());
 }
 
+isl::multi_union_pw_aff ScheduleTreeElemBand::memberRange(
+    size_t first,
+    size_t n) const {
+  auto list = mupa_.get_union_pw_aff_list();
+  auto space = addRange(mupa_.get_space().domain(), n);
+  auto end = first + n;
+  TC_CHECK_LE(end, nMember());
+  list = list.drop(end, nMember() - end);
+  list = list.drop(0, first);
+  return isl::multi_union_pw_aff(space, list);
+}
+
 bool ScheduleTreeElemBand::operator==(const ScheduleTreeElemBand& other) const {
   if (permutable_ != other.permutable_) {
     return false;

tc/core/polyhedral/schedule_tree_elem.h

@@ -248,6 +248,10 @@ struct ScheduleTreeElemBand : public ScheduleTreeElemBase {
   // behavior with respect to changes in the schedule space.
   void drop(size_t pos, size_t n);
 
+  // Extract the range of "n" members starting at "first"
+  // (in an anonymous space).
+  isl::multi_union_pw_aff memberRange(size_t first, size_t n) const;
+
  public:
   bool permutable_{false};
   isl::multi_union_pw_aff mupa_;