ScheduleTreeElemMappingFilter: store mapping between identifiers and functions

Store this mapping in addition to the filter derived from the mapping.
The mapping is needed in the next commit.

The mapping is constructed in ScheduleTree::makeMappingFilter and
passed to the ScheduleTreeElemMappingFilter constructor.
Some of the sanity checks are therefore moved from the constructor
to ScheduleTree::makeMappingFilter.

Author: Sven Verdoolaege

tc/core/polyhedral/cuda/memory_promotion_heuristic.cc

@@ -46,8 +46,8 @@ bool isThreadMapping(const detail::ScheduleTree* tree) {
   using namespace detail;
 
   if (auto filterNode = tree->elemAs<ScheduleTreeElemMappingFilter>()) {
-    for (auto id : filterNode->mappingIds) {
-      if (isThreadId(id)) {
+    for (auto& kvp : filterNode->mapping) {
+      if (isThreadId(kvp.first)) {
         return true;
       }
     }

tc/core/polyhedral/schedule_print.cc

@@ -202,8 +202,8 @@ std::ostream& ScheduleTreeElemFilter::write(std::ostream& os) const {
 std::ostream& ScheduleTreeElemMappingFilter::write(std::ostream& os) const {
   WS w;
   os << w.tab() << "mapping_filter(ids(";
-  for (const auto& id : mappingIds) {
-    os << id << ", ";
+  for (auto& kvp : mapping) {
+    os << kvp.first << ", ";
   }
   os << ")";
   for (const auto& u : filter_.get_set_list()) {

tc/core/polyhedral/schedule_tree-inl.h

@@ -23,15 +23,19 @@ inline ScheduleTreeUPtr ScheduleTree::makeMappingFilter(
     const std::vector<MappingIdType>& mappedIds,
     isl::union_pw_aff_list mappedAffs,
     std::vector<ScheduleTreeUPtr>&& children) {
-  std::vector<mapping::MappingId> ids;
-  for (auto id : mappedIds) {
-    ids.push_back(id);
+  CHECK_EQ(mappedIds.size(), static_cast<size_t>(mappedAffs.n()))
+      << "expected as many mapped ids as affs";
+  ScheduleTreeElemMappingFilter::Mapping mapping;
+  for (size_t i = 0, n = mappedAffs.n(); i < n; ++i) {
+    mapping.emplace(mappedIds.at(i), mappedAffs.get(i));
   }
-  CHECK_GE(ids.size(), 1u) << "empty mapping";
+  CHECK_GE(mapping.size(), 1u) << "empty mapping";
+  CHECK_EQ(mappedIds.size(), mapping.size())
+      << "some id is used more than once in the mapping";
   auto ctx = mappedIds[0].get_ctx();
   ScheduleTreeUPtr res(new ScheduleTree(ctx));
   res->elem_ = std::unique_ptr<ScheduleTreeElemMappingFilter>(
-      new ScheduleTreeElemMappingFilter(ids, mappedAffs));
+      new ScheduleTreeElemMappingFilter(mapping));
   res->type_ = ScheduleTreeType::MappingFilter;
   res->appendChildren(std::move(children));
   return res;

tc/core/polyhedral/schedule_tree_elem.h

@@ -17,7 +17,7 @@
 
 #include <memory>
 #include <sstream>
-#include <unordered_set>
+#include <unordered_map>
 #include <vector>
 
 #include "tc/external/isl.h"
@@ -139,34 +139,29 @@ struct ScheduleTreeElemFilter : public ScheduleTreeElemBase {
 };
 
 struct ScheduleTreeElemMappingFilter : public ScheduleTreeElemFilter {
+  using Mapping = std::unordered_map<
+      mapping::MappingId,
+      isl::union_pw_aff,
+      typename mapping::MappingId::Hash>;
   static constexpr std::initializer_list<detail::ScheduleTreeType>
       NodeDerivedTypes{detail::ScheduleTreeType::None};
   static constexpr detail::ScheduleTreeType NodeType =
       detail::ScheduleTreeType::MappingFilter;
   ScheduleTreeElemMappingFilter() = delete;
   ScheduleTreeElemMappingFilter(const ScheduleTreeElemMappingFilter& eb)
-      : ScheduleTreeElemFilter(eb.filter_), mappingIds(eb.mappingIds) {}
-  ScheduleTreeElemMappingFilter(
-      const std::vector<mapping::MappingId>& mappedIds,
-      isl::union_pw_aff_list mappedAffs)
-      : ScheduleTreeElemFilter(isl::union_set()),
-        mappingIds(mappedIds.begin(), mappedIds.end()) {
-    // Check that ids are unique.
-    CHECK_EQ(mappedIds.size(), mappingIds.size())
-        << "some id is used more than once in the mapping";
+      : ScheduleTreeElemFilter(eb.filter_), mapping(eb.mapping) {}
+  ScheduleTreeElemMappingFilter(const Mapping& mapping)
+      : ScheduleTreeElemFilter(isl::union_set()), mapping(mapping) {
+    CHECK_GT(mapping.size(), 0u) << "empty mapping filter";
 
-    CHECK_EQ(mappedIds.size(), static_cast<size_t>(mappedAffs.n()))
-        << "expected as many mapped ids as affs";
-    CHECK_GT(mappedIds.size(), 0u) << "empty mapping filter";
-
-    auto domain = mappedAffs.get(0).domain();
-    for (size_t i = 1, n = mappedAffs.n(); i < n; ++i) {
-      CHECK(domain.is_equal(mappedAffs.get(i).domain()));
+    auto domain = mapping.cbegin()->second.domain();
+    for (auto& kvp : mapping) {
+      CHECK(domain.is_equal(kvp.second.domain()));
     }
     filter_ = domain.universe();
-    for (size_t i = 0, n = mappedAffs.n(); i < n; ++i) {
-      auto upa = mappedAffs.get(i);
-      auto id = mappedIds.at(i);
+    for (auto& kvp : mapping) {
+      auto upa = kvp.second;
+      auto id = kvp.first;
       // Create mapping filter by equating the
       // parameter mappedIds[i] to the "i"-th affine function.
       upa = upa.sub(isl::union_pw_aff::param_on_domain(domain.universe(), id));
@@ -183,9 +178,8 @@ struct ScheduleTreeElemMappingFilter : public ScheduleTreeElemFilter {
     return NodeType;
   }
 
-  const std::
-      unordered_set<mapping::MappingId, typename mapping::MappingId::Hash>
-          mappingIds;
+  // Mapping from identifiers to affine functions on domain elements.
+  const Mapping mapping;
 };
 
 struct ScheduleTreeElemSequence : public ScheduleTreeElemBase {