extract "extractDomainToIds" from extractMappingToThreads

The core functionality for extracting mapped affine expressions can be
reused for individual mapping ids or other types of mapping ids, e.g.,
block ids.  Therefore, it should live with other schedule operations.

Author: ftynse

tc/core/polyhedral/cuda/mapped_scop.cc

@@ -476,28 +476,12 @@ constexpr auto kWarp = "warp";
 isl::multi_union_pw_aff extractDomainToThread(
     const detail::ScheduleTree* tree,
     size_t nThread) {
-  using namespace polyhedral::detail;
-
-  auto space = isl::space(tree->ctx_, 0);
-  auto empty = isl::union_set::empty(space);
-  auto id = isl::id(tree->ctx_, kBlock);
-  space = space.named_set_from_params_id(id, nThread);
-  auto zero = isl::multi_val::zero(space);
-  auto domainToThread = isl::multi_union_pw_aff(empty, zero);
-
-  for (auto mapping : tree->collect(tree, ScheduleTreeType::MappingFilter)) {
-    auto mappingNode = mapping->elemAs<ScheduleTreeElemMappingFilter>();
-    auto list = isl::union_pw_aff_list(tree->ctx_, nThread);
-    for (size_t i = 0; i < nThread; ++i) {
-      auto threadId = mapping::ThreadId::makeId(i);
-      auto threadMap = mappingNode->mapping.at(threadId);
-      list = list.add(threadMap);
-    }
-    auto nodeToThread = isl::multi_union_pw_aff(space, list);
-    domainToThread = domainToThread.union_add(nodeToThread);
+  std::vector<mapping::MappingId> ids;
+  for (size_t i = 0; i < nThread; ++i) {
+    ids.emplace_back(mapping::ThreadId::makeId(i));
   }
-
-  return domainToThread;
+  auto tupleId = isl::id(tree->ctx_, kBlock);
+  return extractDomainToIds(tree, ids, tupleId);
 }
 
 /*

tc/core/polyhedral/schedule_transforms.cc

@@ -31,6 +31,7 @@
 #include "tc/core/check.h"
 #include "tc/core/constants.h"
 #include "tc/core/polyhedral/functional.h"
+#include "tc/core/polyhedral/mapping_types.h"
 #include "tc/core/polyhedral/schedule_tree_elem.h"
 #include "tc/core/polyhedral/schedule_tree_matcher.h"
 #include "tc/core/scope_guard.h"
@@ -798,5 +799,38 @@ void orderAfter(ScheduleTree* root, ScheduleTree* tree, isl::union_set filter) {
   parent->insertChild(childPos, std::move(seq));
 }
 
+/*
+ * Extract a mapping from the domain elements active at "tree"
+ * to identifiers "ids", where all branches in "tree"
+ * are assumed to have been mapped to these identifiers.
+ * The result lives in a space of the form "tupleId"["ids"...].
+ */
+isl::multi_union_pw_aff extractDomainToIds(
+    const detail::ScheduleTree* tree,
+    const std::vector<mapping::MappingId>& ids,
+    isl::id tupleId) {
+  using namespace polyhedral::detail;
+
+  auto space = isl::space(tree->ctx_, 0);
+  auto empty = isl::union_set::empty(space);
+  space = space.named_set_from_params_id(tupleId, ids.size());
+  auto zero = isl::multi_val::zero(space);
+  auto domainToIds = isl::multi_union_pw_aff(empty, zero);
+
+  for (auto mapping : tree->collect(tree, ScheduleTreeType::MappingFilter)) {
+    auto mappingNode = mapping->elemAs<ScheduleTreeElemMappingFilter>();
+    auto list = isl::union_pw_aff_list(tree->ctx_, ids.size());
+    for (auto id : ids) {
+      CHECK_GT(mappingNode->mapping.count(id), 0) << "no mapping to id " << id;
+      auto idMap = mappingNode->mapping.at(id);
+      list = list.add(idMap);
+    }
+    auto nodeToIds = isl::multi_union_pw_aff(space, list);
+    domainToIds = domainToIds.union_add(nodeToIds);
+  }
+
+  return domainToIds;
+}
+
 } // namespace polyhedral
 } // namespace tc

tc/core/polyhedral/schedule_transforms.h

@@ -22,6 +22,7 @@
 #include <vector>
 
 #include "tc/core/polyhedral/functional.h"
+#include "tc/core/polyhedral/mapping_types.h"
 #include "tc/core/polyhedral/options.h"
 #include "tc/core/polyhedral/schedule_tree.h"
 #include "tc/external/isl.h"
@@ -325,6 +326,15 @@ isl::union_set activeDomainPointsBelow(
     const detail::ScheduleTree* root,
     const detail::ScheduleTree* node);
 
+// Extract a mapping from the domain elements active at "tree"
+// to identifiers "ids", where all branches in "tree"
+// are assumed to have been mapped to these identifiers.
+// The result lives in a space of the form "tupleId"["ids"...].
+isl::multi_union_pw_aff extractDomainToIds(
+    const detail::ScheduleTree* tree,
+    const std::vector<mapping::MappingId>& ids,
+    isl::id tupleId);
+
 } // namespace polyhedral
 } // namespace tc