insert thread specific markers

These markers are placed right underneath the innermost band member mapped
to a thread identifier and allow the memory promotion to recover
these positions.  They form an alternative to ThreadIdxXScheduleDepthState,
which will gradually be removed.
ThreadIdxXScheduleDepthState maintains duplicate state and
does so in an inconsistent way.  In particular, when at least
one member is mapped to a thread identifier, it holds
the depth of the member mapped to the x thread identifier
(which is one less that the schedule depth of the newly introduced marker),
but if no member is mapped, then it points to the depth
of where the marker is introduced.

For consistency, a marker is also introduced in mapCopiesToThreads,
even though ThreadIdxXScheduleDepthState is not being set there.

Author: Sven Verdoolaege

tc/core/polyhedral/cuda/mapped_scop.cc

@@ -161,6 +161,7 @@ void MappedScop::mapToBlocksAndScaleBand(
  * Given a node in the schedule tree of a mapped scop,
  * insert a mapping filter underneath (if needed) that fixes
  * the remaining thread identifiers starting at "begin" to zero.
+ * Add a marker underneath that marks the subtree that is thread specific.
  */
 void fixThreadsBelow(
     MappedScop& mscop,
@@ -173,6 +174,9 @@ void fixThreadsBelow(
 
   auto band = detail::ScheduleTree::makeEmptyBand(mscop.scop().scheduleRoot());
   auto bandTree = insertNodeBelow(tree, std::move(band));
+  auto ctx = tree->ctx_;
+  insertNodeBelow(
+      bandTree, detail::ScheduleTree::makeThreadSpecificMarker(ctx));
   mscop.mapRemaining<mapping::ThreadId>(bandTree, begin);
 }
 
@@ -345,6 +349,9 @@ size_t MappedScop::mapToThreads(detail::ScheduleTree* band) {
     bandSplit(scop_->scheduleRoot(), band, nCanMap);
   }
 
+  auto ctx = band->ctx_;
+  insertNodeBelow(band, detail::ScheduleTree::makeThreadSpecificMarker(ctx));
+
   CHECK_GT(nMappedThreads, 0) << "not mapping to threads";
   CHECK_LE(nMappedThreads, 3) << "mapping to too many threads";
 

tc/core/polyhedral/cuda/mapped_scop.h

@@ -162,6 +162,8 @@ class MappedScop {
   // coincident dimensions (plus reduction dimension, if any),
   // insert synchronization in case of a reduction, and
   // return the number of mapped thread identifiers.
+  // A marker is added to mark the part of the tree that is thread specific
+  // (right underneath the innermost band member mapped to a thread identifier).
   size_t mapToThreads(detail::ScheduleTree* band);
   // Map innermost bands to thread identifiers,
   // inserting synchronization in case of a reduction, and

tc/core/polyhedral/cuda/memory_promotion_heuristic.cc

@@ -65,6 +65,10 @@ void mapCopiesToThreads(MappedScop& mscop, bool unroll) {
       throw promotion::PromotionLogicError("no copy band");
     }
 
+    auto ctx = node->ctx_;
+    insertNodeBelow(
+        bandNode, detail::ScheduleTree::makeThreadSpecificMarker(ctx));
+
     // Check that we are not mapping to threads below other thread mappings.
     std::unordered_set<mapping::ThreadId, mapping::ThreadId::Hash> usedThreads;
     for (auto n : node->ancestors(root)) {

tc/core/polyhedral/schedule_isl_conversion.cc

@@ -203,6 +203,8 @@ isl::schedule_node insert(isl::schedule_node node, const ScheduleTree* st) {
     return insertBranch(node, st);
   } else if (st->elemAs<ScheduleTreeElemExtension>()) {
     return insertExtension(node, st);
+  } else if (st->elemAs<ScheduleTreeElemThreadSpecificMarker>()) {
+    return insertChild(node, st);
   } else {
     LOG(FATAL) << "NYI: insert type: " << *st;
   }
@@ -329,6 +331,7 @@ isl::space definitionParamSpace(const ScheduleTree* node) {
     case detail::ScheduleTreeType::None:
     case detail::ScheduleTreeType::Set:
     case detail::ScheduleTreeType::Sequence:
+    case detail::ScheduleTreeType::ThreadSpecificMarker:
       break;
   }
   return space;

tc/core/polyhedral/schedule_print.cc

@@ -127,6 +127,8 @@ std::ostream& operator<<(std::ostream& os, detail::ScheduleTreeType nt) {
     os << "sequence";
   } else if (nt == detail::ScheduleTreeType::Set) {
     os << "seq";
+  } else if (nt == detail::ScheduleTreeType::ThreadSpecificMarker) {
+    os << "thread_specific";
   } else {
     LOG(FATAL) << "NYI: print type: " << static_cast<int>(nt);
   }
@@ -224,6 +226,13 @@ std::ostream& ScheduleTreeElemSet::write(std::ostream& os) const {
   return os;
 }
 
+std::ostream& ScheduleTreeElemThreadSpecificMarker::write(
+    std::ostream& os) const {
+  WS w;
+  os << w.tab() << "thread_specific()";
+  return os;
+}
+
 std::ostream& operator<<(
     std::ostream& os,
     const std::vector<ScheduleTreeUPtr>& vst) {

tc/core/polyhedral/schedule_tree.cc

@@ -264,6 +264,17 @@ std::unique_ptr<ScheduleTree> ScheduleTree::makeExtension(
   return res;
 }
 
+std::unique_ptr<ScheduleTree> ScheduleTree::makeThreadSpecificMarker(
+    isl::ctx ctx,
+    std::vector<ScheduleTreeUPtr>&& children) {
+  ScheduleTreeUPtr res(new ScheduleTree(ctx));
+  res->elem_ = std::unique_ptr<ScheduleTreeElemThreadSpecificMarker>(
+      new ScheduleTreeElemThreadSpecificMarker());
+  res->type_ = detail::ScheduleTreeType::ThreadSpecificMarker;
+  res->appendChildren(std::move(children));
+  return res;
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 //                        Collector member functions
 ////////////////////////////////////////////////////////////////////////////////

tc/core/polyhedral/schedule_tree.h

@@ -311,6 +311,10 @@ struct ScheduleTree {
       isl::union_map extension,
       std::vector<ScheduleTreeUPtr>&& children = {});
 
+  static ScheduleTreeUPtr makeThreadSpecificMarker(
+      isl::ctx ctx,
+      std::vector<ScheduleTreeUPtr>&& children = {});
+
   template <typename... Args>
   static ScheduleTreeUPtr makeBand(
       isl::multi_union_pw_aff mupa,

tc/core/polyhedral/schedule_tree_elem.cc

@@ -113,6 +113,7 @@ std::unique_ptr<ScheduleTreeElemBase> ScheduleTreeElemBase::make(
   ELEM_MAKE_CASE(ScheduleTreeElemMappingFilter)
   ELEM_MAKE_CASE(ScheduleTreeElemSequence)
   ELEM_MAKE_CASE(ScheduleTreeElemSet)
+  ELEM_MAKE_CASE(ScheduleTreeElemThreadSpecificMarker)
 
 #undef ELEM_MAKE_CASE
 

tc/core/polyhedral/schedule_tree_elem.h

@@ -38,6 +38,7 @@ enum class ScheduleTreeType {
   Sequence,
   Set,
   MappingFilter,
+  ThreadSpecificMarker,
   Any,
 };
 
@@ -279,6 +280,30 @@ struct ScheduleTreeElemBand : public ScheduleTreeElemBase {
   std::vector<bool> unroll_;
 };
 
+/*
+ * A node of type ThreadSpecificMarker marks part of a schedule tree
+ * that is specific to a thread.  That is, the marker appears right
+ * underneath the innermost band member mapped to threads.
+ */
+struct ScheduleTreeElemThreadSpecificMarker : public ScheduleTreeElemBase {
+  static constexpr std::initializer_list<detail::ScheduleTreeType>
+      NodeDerivedTypes{detail::ScheduleTreeType::None};
+  static constexpr detail::ScheduleTreeType NodeType =
+      detail::ScheduleTreeType::ThreadSpecificMarker;
+  explicit ScheduleTreeElemThreadSpecificMarker() {}
+  virtual ~ScheduleTreeElemThreadSpecificMarker() override {}
+  bool operator==(const ScheduleTreeElemThreadSpecificMarker& other) const {
+    return true;
+  }
+  bool operator!=(const ScheduleTreeElemThreadSpecificMarker& other) const {
+    return !(*this == other);
+  }
+  virtual std::ostream& write(std::ostream& os) const override;
+  virtual detail::ScheduleTreeType type() const override {
+    return NodeType;
+  }
+};
+
 bool elemEquals(
     const ScheduleTreeElemBase* e1,
     const ScheduleTreeElemBase* e2,