Merge pull request #302 from facebookresearch/pr/reuse_within

memory promotion: check for any reuse within outer schedule

Author: ftynse

tc/core/polyhedral/cuda/memory_promotion_heuristic.cc

@@ -210,25 +210,19 @@ isl::map fixOuterInputDimsAsParameters(isl::map map, int nDims) {
 }
 
 /*
- * Check if a reference group features reuse at "depth" after applying
- * "schedule". In particular, consider first depth schedule dimensions as fixed
- * by equating them to parameters and check if the resulting relation is not
- * injective.
+ * Check if a reference group features reuse within the "outer" schedule.
+ * In particular, check that for some given point in the outer schedule and
+ * some given group element, there is more than one statement instance
+ * accessing the element within the point in the outer schedule.
+ * In other words, check that the mapping from statement instances
+ * to pairs of outer schedule points and group elements is not injective.
  */
-bool hasReuse(
+bool hasReuseWithin(
     const TensorReferenceGroup& group,
-    isl::union_map schedule,
-    size_t depth) {
-  auto scheduledAccessesUMap = group.originalAccesses().apply_domain(schedule);
-  auto scheduledAccessMaps =
-      isl::UnionAsVector<isl::union_map>(scheduledAccessesUMap);
-  return std::any_of(
-      scheduledAccessMaps.begin(),
-      scheduledAccessMaps.end(),
-      [schedule, depth](isl::map access) {
-        access = fixOuterInputDimsAsParameters(access, static_cast<int>(depth));
-        return !access.is_injective();
-      });
+    isl::multi_union_pw_aff outer) {
+  auto map = isl::union_map::from(outer);
+  map = map.range_product(group.originalAccesses());
+  return !map.is_injective();
 }
 
 /*
@@ -463,6 +457,8 @@ void promoteToSharedGreedy(
   for (auto bandNode : bands) {
     auto groupMap = TensorReferenceGroup::accessedBySubtree(bandNode, scop);
     auto partialSched = partialSchedule(root, bandNode);
+    // Pure affine schedule without (mapping) filters.
+    auto partialSchedMupa = partialScheduleMupa(root, bandNode);
     auto activePoints = activeDomainPoints(root, bandNode);
 
     // Prepare groups for sorting, to have specified order necessary for
@@ -522,7 +518,7 @@ void promoteToSharedGreedy(
         }
         // Do not promote if the group features no reuse and is accessed in a
         // coalesced way.
-        if (!hasReuse(*group, fullSched, depth) &&
+        if (!hasReuseWithin(*group, partialSchedMupa) &&
             isCoalesced(
                 threadIdxXScheduleDepthState,
                 *group,
@@ -606,6 +602,8 @@ void promoteToRegistersBelowThreads(
       // per-thread-group access relations.
       auto points = activeDomainPoints(root, band);
       auto partialSched = partialSchedule(root, band);
+      // Pure affine schedule without (mapping) filters.
+      auto partialSchedMupa = partialScheduleMupa(root, band);
 
       size_t nMappedThreads = 0;
       for (int j = 0; j < points.dim(isl::dim_type::param); ++j) {
@@ -643,7 +641,7 @@ void promoteToRegistersBelowThreads(
                   points)) {
             continue;
           }
-          if (!hasReuse(*group, fullSched, depth)) {
+          if (!hasReuseWithin(*group, partialSchedMupa)) {
             continue;
           }
           // TODO: if something is already in shared, but reuse it within one

tc/core/polyhedral/schedule_transforms.cc

@@ -454,6 +454,14 @@ isl::multi_union_pw_aff prefixScheduleMupa(
   return prefix;
 }
 
+isl::multi_union_pw_aff partialScheduleMupa(
+    const detail::ScheduleTree* root,
+    const detail::ScheduleTree* tree) {
+  auto band = tree->elemAs<ScheduleTreeElemBand>();
+  CHECK(band);
+  return prefixScheduleMupa(root, tree).flat_range_product(band->mupa_);
+}
+
 ScheduleTree* insertBandAbove(
     ScheduleTree* root,
     ScheduleTree* tree,

tc/core/polyhedral/schedule_transforms.h

@@ -276,6 +276,14 @@ isl::multi_union_pw_aff prefixScheduleMupa(
     const detail::ScheduleTree* root,
     const detail::ScheduleTree* tree);
 
+// Return the concatenation of all outer band node partial schedules,
+// including that of the node itself.
+// Note that this function does not take into account
+// any intermediate filter nodes.
+isl::multi_union_pw_aff partialScheduleMupa(
+    const detail::ScheduleTree* root,
+    const detail::ScheduleTree* tree);
+
 // Get the set of domain points active at the given node.  A domain
 // point is active if it was not filtered away on the path from the
 // root to the node.  The root must be a domain element, otherwise no