extractDomainToIds: take into account filters

The original implementation of extractDomainToThreads was assuming that
thread mapping filter nodes also include constraints on the domain
elements active at that position in the schedule tree.  This is not the
case in practice, the mapping filters are defined for statements without
restricting the instance sets.  Therefore, the thread mapping functions
computed by extractDomainToThreads are incorrect for trees that schedule
different instances of the same statement in different branches. In
particular, calling "union_add" takes a sum of the two mappings for the
statements whose instances appear in different branches.

Intersect the domain of the thread mapping function with the set of
points active at the mapping node.  Introduce a check that the result of
this intersection and the set of instances with known mapping is
disjoint.

Note that, to enable the check that all domain points were mapped, the
intersection must be made with the active domain points without
accounting for ancestor mapping filters.  In particular, if
extractDomainToIds is called on the root of a mapped tree and looks for
thread ids, the mapping to blocks will be considered to filter out some
domain points and the thread mapping will not cover the entire domain.
Introduce "activeDomainPointsNoMapping" to avoid this problem.  Unlike
"activeDomainPoints", this function does not take into account extension
nodes.  This is currently safe because the callers of
"extractDomainToIds" do not rely on it returning the mapping of the
elements introduced by extension nodes.  Comment them accordingly.
In the future, we should reconsider activeDomainPoints interpreting
MappingFilter nodes as regular filters, or mappings being filters in
general, as well as "activeDomainPointsNoMapping" ignoring extension
nodes.

Author: ftynse

tc/core/polyhedral/cuda/mapped_scop.h

@@ -190,13 +190,17 @@ class MappedScop {
   // to the thread identifiers, where all branches in "tree"
   // are assumed to have been mapped to thread identifiers.
   // The result lives in a space of the form block[x, ...].
+  //
+  // Note: this function ignores statements introduced by extension nodes.
   isl::multi_union_pw_aff threadMappingSchedule(
       const detail::ScheduleTree* tree) const;
 
   // Extract a mapping from the domain elements active at "tree"
   // to the block identifiers, where all branches in "tree"
   // are assumed to have been mapped to block identifiers.
   // The result lives in a space of the form grid[x, ...].
+  //
+  // Note: this function ignores statements introduced by extension nodes.
   isl::multi_union_pw_aff blockMappingSchedule(
       const detail::ScheduleTree* tree) const;
 

tc/core/polyhedral/schedule_transforms.cc

@@ -145,6 +145,18 @@ isl::union_set activeDomainPointsBelow(
   return activeDomainPointsHelper(root, ancestors);
 }
 
+isl::union_set activeDomainPointsNoMappingNoExtension(
+    const detail::ScheduleTree* root,
+    const detail::ScheduleTree* tree) {
+  auto domain = root->elemAs<detail::ScheduleTreeElemDomain>()->domain_;
+  for (auto t : tree->ancestors(root)) {
+    if (auto f = t->elemAs<detail::ScheduleTreeElemFilter>()) {
+      domain = domain.intersect(f->filter_);
+    }
+  }
+  return domain;
+}
+
 vector<ScheduleTree*> collectScheduleTreesPath(
     std::function<ScheduleTree*(ScheduleTree*)> next,
     ScheduleTree* start) {
@@ -804,6 +816,10 @@ void orderAfter(ScheduleTree* root, ScheduleTree* tree, isl::union_set filter) {
  * 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"...].
+ *
+ * Note: this function only takes into account points that are present in the
+ * root domain node.  Those introduced by extension nodes are ignored.  This
+ * behavior can change in the future.
  */
 isl::multi_union_pw_aff extractDomainToIds(
     const detail::ScheduleTree* root,
@@ -833,12 +849,17 @@ isl::multi_union_pw_aff extractDomainToIds(
       continue;
     }
     auto nodeToIds = isl::multi_union_pw_aff(space, list);
+    auto active = activeDomainPointsNoMappingNoExtension(root, mapping);
+    TC_CHECK(active.intersect(domainToIds.domain()).is_empty())
+        << "conflicting mappings; are the filters in the tree disjoint?";
+    nodeToIds = nodeToIds.intersect_domain(active);
     domainToIds = domainToIds.union_add(nodeToIds);
   }
 
-  TC_CHECK(activeDomainPoints(root, tree).is_subset(domainToIds.domain()))
-      << "not all domain points of" << activeDomainPoints(root, tree)
-      << "were mapped to the required ids";
+  auto active = activeDomainPointsNoMappingNoExtension(root, tree);
+  TC_CHECK(active.is_subset(domainToIds.domain()))
+      << "not all domain points of\n"
+      << active << "\nwere mapped to the required ids";
 
   return domainToIds;
 }

tc/core/polyhedral/schedule_transforms.h

@@ -326,6 +326,15 @@ isl::union_set activeDomainPointsBelow(
     const detail::ScheduleTree* root,
     const detail::ScheduleTree* node);
 
+// Get the set of domain points active below the given node without including
+// the points introduced by extension nodes and without treating mapping nodes
+// as filters.  A point is considered active at a schedule node "tree" if it is
+// present in the "root" domain node and was not filtered away on the path from
+// "root" to "tree".  The root must be a domain element.
+isl::union_set activeDomainPointsNoMappingNoExtension(
+    const detail::ScheduleTree* root,
+    const detail::ScheduleTree* tree);
+
 // Extract a mapping from the domain elements active at "tree" (in a tree
 // rooted at "root") to identifiers "ids", where all branches in "tree" are
 // assumed to have been mapped to these identifiers.  The result lives in a