Merge pull request #320 from facebookresearch/pr/copies_under

insertCopiesUnder: use insertExtension{Before,After}

Author: skimo-openhub

tc/core/polyhedral/memory_promotion.cc

@@ -391,34 +391,6 @@ TensorReferenceGroup::referenceIds() const {
 }
 
 namespace {
-bool hasCopyExtensionSingleChild(const ScheduleTree* tree) {
-  if (tree->numChildren() != 1) {
-    return false;
-  }
-
-  auto extensionNode =
-      tree->child({0})->elemAs<detail::ScheduleTreeElemExtension>();
-  if (!extensionNode) {
-    return false;
-  }
-
-  if ((tree->child({0})->numChildren() != 1) &&
-      (tree->child({0, 0})->elemAs<detail::ScheduleTreeElemSequence>())) {
-    return false;
-  }
-
-  for (auto e : isl::UnionAsVector<isl::union_map>(extensionNode->extension_)) {
-    if (!e.has_tuple_name(isl::dim_type::out)) {
-      return false;
-    }
-    if (e.get_tuple_name(isl::dim_type::out) != kReadIdName &&
-        e.get_tuple_name(isl::dim_type::out) != kWriteIdName) {
-      return false;
-    }
-  }
-  return true;
-}
-
 // Construct the set containing all tensor elements.
 //
 // Find the Halide image corresponding to the given tensorId.  Transform its
@@ -524,48 +496,26 @@ ScheduleTree* insertCopiesUnder(
   bool reads = !group.scopedReads().is_empty();
   bool writes = !group.scopedWrites().is_empty();
 
-  if (hasCopyExtensionSingleChild(tree)) {
-    auto extensionNode = tree->child({0});
-    auto sequenceNode = tree->child({0, 0});
-
-    auto& ext =
-        extensionNode->elemAs<detail::ScheduleTreeElemExtension>()->extension_;
-    if (reads) {
-      ext = ext.unite(isl::union_map(readExtension));
-      sequenceNode->insertChild(0, std::move(readFilterNode));
-    }
-    if (writes) {
-      ext = ext.unite(isl::union_map(writeExtension));
-      sequenceNode->appendChild(std::move(writeFilterNode));
-    }
-    return tree;
+  if (tree->numChildren() == 0) {
+    // The point underneath a leaf node cannot be referenced,
+    // so insert a dummy sequence first.  It will be extended
+    // with the reads and/or writes.
+    insertSequenceBelow(root, tree);
   }
 
-  auto mainCompFilter = activeDomainPoints(root, tree).universe();
-  auto mainCompFilterNode =
-      ScheduleTree::makeFilter(mainCompFilter, tree->detachChildren());
-
-  // XXX: I don't really like the syntax-imposed impossibility to create a
-  // sequence node with no children.
-  auto sequenceNode = ScheduleTree::makeSequence(
-      reads ? std::move(readFilterNode) : std::move(mainCompFilterNode));
   if (reads) {
-    sequenceNode->appendChild(std::move(mainCompFilterNode));
+    insertExtensionBefore(
+        root, tree, tree->child({0}), readExtension, std::move(readFilterNode));
   }
   if (writes) {
-    sequenceNode->appendChild(std::move(writeFilterNode));
+    insertExtensionAfter(
+        root,
+        tree,
+        tree->child({0}),
+        writeExtension,
+        std::move(writeFilterNode));
   }
 
-  auto extensionUmap = isl::union_map::empty(promotionSpace.params());
-  if (reads) {
-    extensionUmap = extensionUmap.unite(readExtension);
-  }
-  if (writes) {
-    extensionUmap = extensionUmap.unite(writeExtension);
-  }
-  auto extensionNode =
-      ScheduleTree::makeExtension(extensionUmap, std::move(sequenceNode));
-  tree->appendChild(std::move(extensionNode));
   return tree;
 }
 } // namespace polyhedral

tc/core/polyhedral/schedule_isl_conversion.cc

@@ -22,6 +22,7 @@
 #include "tc/external/isl.h"
 
 #include "tc/core/flags.h"
+#include "tc/core/polyhedral/schedule_transforms.h"
 #include "tc/external/isl.h"
 
 using namespace std;

tc/core/polyhedral/schedule_transforms.cc

@@ -99,23 +99,22 @@ isl::union_map partialSchedule(
   return partialScheduleImpl(root, node, true);
 }
 
-// Get a set of domain elements that are active at the given node.
+namespace {
+// Get a set of domain elements that are active below
+// the given branch of nodes.
 //
 // Domain elements are introduced by the root domain node.  Filter nodes
 // disable the points that do not intersect with the filter.  Extension nodes
 // are considered to introduce additional domain points.
-isl::union_set activeDomainPoints(
+isl::union_set activeDomainPointsHelper(
     const ScheduleTree* root,
-    const ScheduleTree* node) {
+    const vector<const ScheduleTree*>& nodes) {
   auto domainElem = root->elemAs<ScheduleTreeElemDomain>();
   CHECK(domainElem) << "root must be a Domain node" << *root;
 
   auto domain = domainElem->domain_;
-  if (root == node) {
-    return domain;
-  }
 
-  for (auto anc : node->ancestors(root)) {
+  for (auto anc : nodes) {
     if (auto filterElem = anc->elemAsBase<ScheduleTreeElemFilter>()) {
       domain = domain.intersect(filterElem->filter_);
     } else if (auto extensionElem = anc->elemAs<ScheduleTreeElemExtension>()) {
@@ -134,6 +133,21 @@ isl::union_set activeDomainPoints(
   }
   return domain;
 }
+} // namespace
+
+isl::union_set activeDomainPoints(
+    const ScheduleTree* root,
+    const ScheduleTree* node) {
+  return activeDomainPointsHelper(root, node->ancestors(root));
+}
+
+isl::union_set activeDomainPointsBelow(
+    const ScheduleTree* root,
+    const ScheduleTree* node) {
+  auto ancestors = node->ancestors(root);
+  ancestors.emplace_back(node);
+  return activeDomainPointsHelper(root, ancestors);
+}
 
 vector<ScheduleTree*> collectScheduleTreesPath(
     std::function<ScheduleTree*(ScheduleTree*)> next,
@@ -473,8 +487,16 @@ void updateTopLevelContext(detail::ScheduleTree* root, isl::set context) {
   contextElem->context_ = contextElem->context_ & context;
 }
 
-ScheduleTree* insertSequenceAbove(ScheduleTree* root, ScheduleTree* tree) {
-  auto parent = tree->ancestor(root, 1);
+namespace {
+
+// In a tree starting at "root", insert a sequence node with
+// as only child the node identified by "tree"
+// within the subtree at "relativeRoot".
+ScheduleTree* insertSequenceAbove(
+    const ScheduleTree* root,
+    ScheduleTree* relativeRoot,
+    ScheduleTree* tree) {
+  auto parent = tree->ancestor(relativeRoot, 1);
   auto childPos = tree->positionInParent(parent);
   auto filter = activeDomainPoints(root, tree).universe();
   parent->insertChild(
@@ -484,11 +506,27 @@ ScheduleTree* insertSequenceAbove(ScheduleTree* root, ScheduleTree* tree) {
   return parent->child({childPos});
 }
 
+} // namespace
+
+ScheduleTree* insertSequenceAbove(ScheduleTree* root, ScheduleTree* tree) {
+  return insertSequenceAbove(root, root, tree);
+}
+
+void insertSequenceBelow(
+    const detail::ScheduleTree* root,
+    detail::ScheduleTree* tree) {
+  auto numChildren = tree->numChildren();
+  CHECK_LE(numChildren, 1u);
+  auto filter = activeDomainPointsBelow(root, tree).universe();
+  auto node = ScheduleTree::makeFilter(filter, tree->detachChildren());
+  tree->appendChild(ScheduleTree::makeSequence(std::move(node)));
+}
+
 ScheduleTree* insertExtensionAbove(
-    ScheduleTree* root,
+    ScheduleTree* relativeRoot,
     ScheduleTree* tree,
     isl::union_map extension) {
-  auto parent = tree->ancestor(root, 1);
+  auto parent = tree->ancestor(relativeRoot, 1);
   auto childPos = tree->positionInParent(parent);
   auto child = parent->detachChild(childPos);
   parent->insertChild(
@@ -500,85 +538,153 @@ namespace {
 /*
  * Insert an empty extension node above "st" in a tree with the given root and
  * return a pointer to the inserted extension node.
+ * The modification is performed within the subtree at "relativeRoot".
  */
 detail::ScheduleTree* insertEmptyExtensionAbove(
-    ScheduleTree* root,
+    const ScheduleTree* root,
+    ScheduleTree* relativeRoot,
     ScheduleTree* st) {
   auto domain = root->elemAs<ScheduleTreeElemDomain>();
   CHECK(domain);
   auto space = domain->domain_.get_space();
   auto extension = isl::union_map::empty(space);
-  return insertExtensionAbove(root, st, extension);
+  return insertExtensionAbove(relativeRoot, st, extension);
 }
-} // namespace
 
-void insertExtensionLabelAt(
-    ScheduleTree* root,
+/*
+ * Construct an extension map for a zero-dimensional statement
+ * with the given identifier.
+ */
+isl::map labelExtension(ScheduleTree* root, ScheduleTree* tree, isl::id id) {
+  auto prefix = prefixScheduleMupa(root, tree);
+  auto scheduleSpace = prefix.get_space();
+  auto space = scheduleSpace.params().set_from_params().set_tuple_id(
+      isl::dim_type::set, id);
+  auto extensionSpace = scheduleSpace.map_from_domain_and_range(space);
+  return isl::map::universe(extensionSpace);
+}
+
+/*
+ * Construct a filter node for a zero-dimensional extension statement
+ * with the given extension map.
+ */
+ScheduleTreeUPtr labelFilterFromExtension(isl::map extension) {
+  return detail::ScheduleTree::makeFilter(extension.range());
+}
+
+/*
+ * Given a sequence node in the schedule tree, insert
+ * an extension with the given extension map and extension filter node
+ * before the child at position "pos".
+ * If "pos" is equal to the number of children, then
+ * the statement is added after the last child.
+ * The modification is performed within the subtree at "relativeRoot".
+ */
+void insertExtensionAt(
+    const ScheduleTree* root,
+    ScheduleTree* relativeRoot,
     ScheduleTree* seqNode,
     size_t pos,
-    isl::id id) {
-  auto extensionTree = seqNode->ancestor(root, 1);
+    isl::union_map extension,
+    ScheduleTreeUPtr&& filterNode) {
+  auto extensionTree = seqNode->ancestor(relativeRoot, 1);
   auto extensionNode =
       extensionTree->elemAs<detail::ScheduleTreeElemExtension>();
   if (!extensionNode) {
-    extensionTree = insertEmptyExtensionAbove(root, seqNode);
+    extensionTree = insertEmptyExtensionAbove(root, relativeRoot, seqNode);
     extensionNode = extensionTree->elemAs<detail::ScheduleTreeElemExtension>();
   }
   CHECK(extensionNode);
   CHECK(seqNode->elemAs<detail::ScheduleTreeElemSequence>());
-  auto prefix = prefixScheduleMupa(root, extensionTree);
-  auto scheduleSpace = prefix.get_space();
-  auto space = scheduleSpace.params().set_from_params().set_tuple_id(
-      isl::dim_type::set, id);
-  auto extensionSpace = scheduleSpace.map_from_domain_and_range(space);
-  auto extension = isl::map::universe(extensionSpace);
   extensionNode->extension_ = extensionNode->extension_.unite(extension);
-  auto filterNode = detail::ScheduleTree::makeFilter(extension.range());
   seqNode->insertChild(pos, std::move(filterNode));
 }
+} // namespace
 
-void insertExtensionLabelBefore(
-    ScheduleTree* root,
+void insertExtensionBefore(
+    const ScheduleTree* root,
+    ScheduleTree* relativeRoot,
     ScheduleTree* tree,
-    isl::id id) {
+    isl::union_map extension,
+    ScheduleTreeUPtr&& filterNode) {
   size_t pos;
-  auto parent = tree->ancestor(root, 1);
+  auto parent = tree->ancestor(relativeRoot, 1);
   ScheduleTree* seqTree;
+  if (tree->elemAs<detail::ScheduleTreeElemExtension>()) {
+    tree = tree->child({0});
+    parent = tree;
+  }
   if (tree->elemAs<detail::ScheduleTreeElemSequence>()) {
     seqTree = tree;
     pos = 0;
   } else if (
       parent->elemAs<detail::ScheduleTreeElemFilter>() &&
       parent->ancestor(root, 1)->elemAs<detail::ScheduleTreeElemSequence>()) {
-    seqTree = parent->ancestor(root, 1);
+    seqTree = parent->ancestor(relativeRoot, 1);
     pos = parent->positionInParent(seqTree);
   } else {
-    seqTree = insertSequenceAbove(root, tree);
+    seqTree = insertSequenceAbove(root, relativeRoot, tree);
     pos = 0;
   }
-  insertExtensionLabelAt(root, seqTree, pos, id);
+  insertExtensionAt(
+      root, relativeRoot, seqTree, pos, extension, std::move(filterNode));
 }
 
-void insertExtensionLabelAfter(
-    ScheduleTree* root,
+void insertExtensionAfter(
+    const ScheduleTree* root,
+    ScheduleTree* relativeRoot,
     ScheduleTree* tree,
-    isl::id id) {
+    isl::union_map extension,
+    ScheduleTreeUPtr&& filterNode) {
   size_t pos;
-  auto parent = tree->ancestor(root, 1);
+  auto parent = tree->ancestor(relativeRoot, 1);
   ScheduleTree* seqTree;
+  if (tree->elemAs<detail::ScheduleTreeElemExtension>()) {
+    tree = tree->child({0});
+    parent = tree;
+  }
   if (tree->elemAs<detail::ScheduleTreeElemSequence>()) {
     seqTree = tree;
     pos = tree->numChildren();
   } else if (
       parent->elemAs<detail::ScheduleTreeElemFilter>() &&
       parent->ancestor(root, 1)->elemAs<detail::ScheduleTreeElemSequence>()) {
-    seqTree = parent->ancestor(root, 1);
+    seqTree = parent->ancestor(relativeRoot, 1);
     pos = parent->positionInParent(seqTree) + 1;
   } else {
-    seqTree = insertSequenceAbove(root, tree);
+    seqTree = insertSequenceAbove(root, relativeRoot, tree);
     pos = 1;
   }
-  insertExtensionLabelAt(root, seqTree, pos, id);
+  insertExtensionAt(
+      root, relativeRoot, seqTree, pos, extension, std::move(filterNode));
+}
+
+void insertExtensionLabelAt(
+    ScheduleTree* root,
+    ScheduleTree* seqNode,
+    size_t pos,
+    isl::id id) {
+  auto extension = labelExtension(root, seqNode, id);
+  auto filterNode = labelFilterFromExtension(extension);
+  insertExtensionAt(root, root, seqNode, pos, extension, std::move(filterNode));
+}
+
+void insertExtensionLabelBefore(
+    ScheduleTree* root,
+    ScheduleTree* tree,
+    isl::id id) {
+  auto extension = labelExtension(root, tree, id);
+  auto filterNode = labelFilterFromExtension(extension);
+  insertExtensionBefore(root, root, tree, extension, std::move(filterNode));
+}
+
+void insertExtensionLabelAfter(
+    ScheduleTree* root,
+    ScheduleTree* tree,
+    isl::id id) {
+  auto extension = labelExtension(root, tree, id);
+  auto filterNode = labelFilterFromExtension(extension);
+  insertExtensionAfter(root, root, tree, extension, std::move(filterNode));
 }
 
 namespace {