use domain sets instead of statement ids in memory promotion

Original implementation of memory promotion (and related parts of code
generation) was assuming schedule tree branches only for entire
statements.  It was therefore using statement ids as a lightweight key
for finding relevant active promotions and schedule parts.  With the
introduction of full/partial tile separation, this assumption no longer
holds.  Make memory promotion use domain sets instead of statement ids,
and find relevant promotions and schedule parts by intersecting these
sets with sets of statement instances active in a particular location
in the schedule tree.

Author: ftynse

include/tc/core/polyhedral/codegen_cuda.h

@@ -96,13 +96,18 @@ struct CodegenStatementContext : CodegenContext {
   isl::id statementId() const {
     return this->iteratorMaps.at(astNodeId).get_tuple_id(isl::dim_type::out);
   }
+  isl::set domain() const {
+    return isl::map::from(this->iteratorMaps.at(astNodeId)).range();
+  }
   std::vector<Scop::PromotionInfo> activePromotions() const {
-    auto stmtId = statementId();
-    const auto& promotions = this->scop().activePromotions();
-    if (promotions.count(stmtId) == 0) {
-      return {};
+    std::vector<Scop::PromotionInfo> result;
+    auto dom = isl::union_set(this->domain());
+    for (const auto& kvp : this->scop().activePromotions()) {
+      if (!kvp.first.intersect(dom).is_empty()) {
+        result.emplace_back(kvp.second);
+      }
     }
-    return promotions.at(stmtId);
+    return result;
   }
 
   isl::id astNodeId;

include/tc/core/polyhedral/scop.h

@@ -324,9 +324,8 @@ struct Scop {
     return promotedDecls_;
   }
 
-  const std::
-      unordered_map<isl::id, std::vector<PromotionInfo>, isl::IslIdIslHash>&
-      activePromotions() const {
+  const std::vector<std::pair<isl::union_set, PromotionInfo>>&
+  activePromotions() const {
     return activePromotions_;
   }
 
@@ -377,7 +376,6 @@ struct Scop {
       isl::id tensorId,
       std::unique_ptr<TensorReferenceGroup>&& gr,
       detail::ScheduleTree* tree,
-      const std::unordered_set<isl::id, isl::IslIdIslHash>& activeStmts,
       isl::union_map schedule,
       bool forceLastExtentOdd = false);
 
@@ -468,9 +466,10 @@ struct Scop {
   std::unordered_map<isl::id, size_t, isl::IslIdIslHash> groupCounts_;
   // groupId -> (tensorId, groupSizes)
   std::unordered_map<isl::id, PromotedDecl, isl::IslIdIslHash> promotedDecls_;
-  // stmtId -> (group, partial schedule, groupId)
-  std::unordered_map<isl::id, std::vector<PromotionInfo>, isl::IslIdIslHash>
-      activePromotions_;
+  // (domain, group, partial schedule, groupId)
+  // Note that domain is a non-unique key, i.e. multiple groups can be listed
+  // for the same domain, or for partially intersecting domains.
+  std::vector<std::pair<isl::union_set, PromotionInfo>> activePromotions_;
 };
 
 std::ostream& operator<<(std::ostream& os, const Scop&);

src/core/polyhedral/codegen_cuda.cc

@@ -503,17 +503,6 @@ isl::space findDomainSpaceById(const CodegenStatementContext& context) {
   return isl::space();
 }
 
-isl::map findScheduleByStmtId(isl::union_map schedule, isl::id stmtId) {
-  for (auto s : isl::UnionAsVector<isl::union_map>(schedule)) {
-    if (s.get_tuple_id(isl::dim_type::in) == stmtId) {
-      return s;
-    }
-  }
-  CHECK(false) << "could not find schedule for " << stmtId << " in "
-               << schedule;
-  return isl::map();
-}
-
 isl::multi_aff makeMultiAffAccess(
     isl::id tensorId,
     const std::vector<Halide::Expr>& subscripts,
@@ -633,9 +622,9 @@ void emitMappedTensorAccess(
   auto promotion = promotionInfo.group->promotion(); // MA :: [S -> O] -> P
   promotion = promotion.set_tuple_id(isl::dim_type::out, promotionInfo.groupId);
   auto iteratorMap = context.iteratorMap(); // PMA :: A -> D
-  auto schedule = findScheduleByStmtId(
-      promotionInfo.outerSchedule,
-      context.statementId()); // map :: D -> S
+  auto schedule =
+      isl::map::from_union_map(promotionInfo.outerSchedule.intersect_domain(
+          context.domain())); // map :: D -> S
 
   CHECK(schedule.is_single_valued())
       << "expected single-valued schedule, got " << schedule;

src/core/polyhedral/memory_promotion_heuristic.cc

@@ -300,9 +300,7 @@ bool isCoalesced(
       auto partialSchedule = isl::map::from_union_map(partialScheduleUMap);
       auto scheduleToNextX = makeNextElementMap(
           partialSchedule.get_space().range(), threadIdxxDepth);
-      auto scheduledAccess = isl::map(access)
-                                 .gist_domain(access.domain())
-                                 .apply_domain(partialSchedule);
+      auto scheduledAccess = isl::map(access).apply_domain(partialSchedule);
       auto accessedByAdjacentX = scheduleToNextX.apply_domain(scheduledAccess)
                                      .apply_range(scheduledAccess);
 
@@ -349,9 +347,7 @@ bool isPromotableToRegisterBelowThreads(
                    threadIdxxScheduleDepthState,
                    originalAccesses.domain().intersect(activePoints));
 
-  auto scheduledAccesses =
-      originalAccesses.gist_domain(originalAccesses.domain())
-          .apply_domain(schedule);
+  auto scheduledAccesses = originalAccesses.apply_domain(schedule);
 
   // Scheduled accesses contain maps from schedule dimensions to tensor
   // subscripts.  Compute the relation that between the schedule dimensions
@@ -460,7 +456,6 @@ void promoteToSharedGreedy(
   size_t remainingMemory = maxMemory;
   for (auto bandNode : bands) {
     auto groupMap = TensorReferenceGroup::accessedBySubtree(bandNode, scop);
-    auto activeStmts = activeStatements(root, bandNode);
     auto partialSched = partialSchedule(root, bandNode);
     auto activePoints = activeDomainPoints(root, bandNode);
 
@@ -535,7 +530,6 @@ void promoteToSharedGreedy(
             tensorId,
             std::move(group),
             bandNode,
-            activeStmts,
             partialSched,
             true);
         remainingMemory -= memoryRequirement;

src/core/polyhedral/scop.cc

@@ -184,9 +184,10 @@ void Scop::promoteGroup(
     isl::id tensorId,
     std::unique_ptr<TensorReferenceGroup>&& gr,
     ScheduleTree* tree,
-    const std::unordered_set<isl::id, isl::IslIdIslHash>& activeStmts,
     isl::union_map schedule,
     bool forceLastExtentOdd) {
+  auto activePoints = activeDomainPoints(scheduleRoot(), tree);
+
   for (const auto& id : activeStmts) {
     for (const auto& prom : activePromotions_[id]) {
       if (promotedDecls_.count(prom.groupId) != 0 &&
@@ -211,10 +212,10 @@ void Scop::promoteGroup(
   }
   promotedDecls_[groupId] = PromotedDecl{tensorId, sizes, kind};
 
+  // FIXME: we can now store a unique pointer...
   auto group = std::shared_ptr<TensorReferenceGroup>(std::move(gr));
-  for (const auto& id : activeStmts) {
-    activePromotions_[id].push_back(PromotionInfo{group, schedule, groupId});
-  }
+  activePromotions_.emplace_back(
+      std::make_pair(activePoints, PromotionInfo{group, schedule, groupId}));
 }
 
 void Scop::insertSyncsAroundCopies(ScheduleTree* tree) {
@@ -259,7 +260,6 @@ void Scop::promoteEverythingAt(std::vector<size_t> pos) {
   auto tree = scheduleRoot()->child(pos);
 
   checkFiltersDisjointStatements(scheduleRoot());
-  auto activeStmts = activeStatements(root, tree);
   auto schedule = partialSchedule(root, tree);
 
   auto groupMap = TensorReferenceGroup::accessedBySubtree(tree, *this);
@@ -270,7 +270,6 @@ void Scop::promoteEverythingAt(std::vector<size_t> pos) {
           p.first,
           std::move(gr),
           tree,
-          activeStmts,
           schedule);
     }
   }