ScheduleTreeElem*: inherit from ScheduleTree instead of ScheduleTreeElemBase

The existing implementation of TC schedule trees, inherited from the
prehistory, chose to have an abstract tree node type (ScheduleTree) with
a unique pointer to a specific node element (ScheduleTreeElem*) as
payload.  It is unclear now why this structure was chosen, possible
reasons include: the duality between isl_schedule_tree and
isl_schedule_node in isl internals; gradual movement from raw pointers,
to shared, to unique pointers; an attempt to separate tree memory
management (unique pointers to children in ScheduleTree) from data
(fields in ScheduleTreeElem*).  However, the resulting structure ended
up being unnecessarily complex.  Conceptually, a specific node type
_is-a_ tree node, so it makes sense to have specific nodes inherit from
a generic one.  isl C++ API already provides this structure and it
proved to be more convenient.  In particular, it enables functions to
specify that they expect a certain type of node statically as a part of
their signature rather than dynamically through CHECK macros.

Make ScheduleTreeElem* classes inherit from ScheduleTree instead of
ScheduleTreeElemBase and remove the "elem_" field from ScheduleTree (the
tree class itself can now be casted to a subclass).  This change
resulted in changes to copy and static constructors of the ScheduleTree*
classes.  In particular, Sequence, Set and ThreadSpecificMarker elements
now use the global isl context in construction since the ScheduleTree
constructor requires an isl context.  This will be cleaned up in an
upcoming commit.

Note that this commit does not rename elem-related functions for the
sake of diff minimization.  These changes will be performed in
subsequent commits.

Author: ftynse

tc/core/polyhedral/schedule_isl_conversion.cc

@@ -248,8 +248,7 @@ std::unique_ptr<ScheduleTreeElemBand> fromIslScheduleNodeBand(
   return res;
 }
 
-std::unique_ptr<ScheduleTreeElemBase> elemFromIslScheduleNode(
-    isl::schedule_node node) {
+std::unique_ptr<ScheduleTree> elemFromIslScheduleNode(isl::schedule_node node) {
   if (auto band = node.as<isl::schedule_node_band>()) {
     return fromIslScheduleNodeBand(band);
   } else if (auto context = node.as<isl::schedule_node_context>()) {
@@ -278,15 +277,15 @@ std::unique_ptr<ScheduleTreeElemBase> elemFromIslScheduleNode(
     LOG(FATAL) << "mark nodes not supported";
     return nullptr;
   } else if (node.isa<isl::schedule_node_leaf>()) {
-    LOG(FATAL) << "ScheduleTreeElemBase::make called on explicit leaf";
+    LOG(FATAL) << "ScheduleTree::make called on explicit leaf";
     return nullptr;
   } else if (node.isa<isl::schedule_node_sequence>()) {
     return std::unique_ptr<ScheduleTreeElemSequence>(
         new ScheduleTreeElemSequence());
   } else if (node.isa<isl::schedule_node_set>()) {
     return std::unique_ptr<ScheduleTreeElemSet>(new ScheduleTreeElemSet());
   }
-  LOG(FATAL) << "NYI: ScheduleTreeElemBase from type: "
+  LOG(FATAL) << "NYI: ScheduleTree from type: "
              << isl_schedule_node_get_type(node.get());
   return nullptr;
 }
@@ -299,9 +298,7 @@ std::unique_ptr<ScheduleTreeElemBase> elemFromIslScheduleNode(
  * if this single child node is a leaf.
  */
 std::unique_ptr<ScheduleTree> fromIslScheduleNode(isl::schedule_node node) {
-  unique_ptr<ScheduleTree> res(new ScheduleTree(node.get_ctx()));
-  res->elem_ = elemFromIslScheduleNode(node);
-  res->type_ = res->elem_->type();
+  auto res = elemFromIslScheduleNode(node);
   auto n = node.n_children();
   if (n == 1 && node.child(0).isa<isl::schedule_node_leaf>()) {
     return res;

tc/core/polyhedral/schedule_print.cc

@@ -137,11 +137,6 @@ std::ostream& operator<<(std::ostream& os, detail::ScheduleTreeType nt) {
   return os;
 }
 
-std::ostream& operator<<(std::ostream& os, const ScheduleTreeElemBase& st) {
-  st.write(os);
-  return os;
-}
-
 std::ostream& ScheduleTreeElemBand::write(std::ostream& os) const {
   WS w;
   os << w.tab() << "band(n(" << coincident_.size() << ") permutable(";
@@ -253,8 +248,8 @@ std::ostream& operator<<(
 }
 
 std::ostream& operator<<(std::ostream& os, const ScheduleTree& st) {
-  TC_CHECK(st.elem_.get());
-  os << *st.elem_ << "\n";
+  st.write(os);
+  os << "\n";
   os << st.children_;
 
   return os;

tc/core/polyhedral/schedule_tree.cc

@@ -123,11 +123,10 @@ vector<ScheduleTree*> ancestorsInSubTree(
   return res;
 }
 
-static std::unique_ptr<ScheduleTreeElemBase> makeElem(const ScheduleTree& st) {
-#define ELEM_MAKE_CASE(CLASS)                             \
-  else if (st.type_ == CLASS::NodeType) {                 \
-    return std::unique_ptr<CLASS>(                        \
-        new CLASS(*static_cast<CLASS*>(st.elem_.get()))); \
+static std::unique_ptr<ScheduleTree> makeElem(const ScheduleTree& st) {
+#define ELEM_MAKE_CASE(CLASS)                                                \
+  else if (st.type_ == CLASS::NodeType) {                                    \
+    return std::unique_ptr<CLASS>(new CLASS(static_cast<const CLASS&>(st))); \
   }
 
   if (st.type_ == detail::ScheduleTreeType::None) {
@@ -145,8 +144,7 @@ static std::unique_ptr<ScheduleTreeElemBase> makeElem(const ScheduleTree& st) {
 
 #undef ELEM_MAKE_CASE
 
-  LOG(FATAL) << "NYI: ScheduleTreeElemBase from type: "
-             << static_cast<int>(st.type_);
+  LOG(FATAL) << "NYI: ScheduleTree from type: " << static_cast<int>(st.type_);
   return nullptr;
 }
 } // namespace
@@ -159,17 +157,15 @@ ScheduleTree::ScheduleTree(isl::ctx ctx) : ctx_(ctx) {}
 ScheduleTree::~ScheduleTree() {}
 
 ScheduleTree::ScheduleTree(const ScheduleTree& st)
-    : ctx_(st.ctx_), children_(), type_(st.type_), elem_(nullptr) {
+    : ctx_(st.ctx_), children_(), type_(st.type_) {
   children_.reserve(st.children_.size());
   for (const auto& child : st.children()) {
     children_.push_back(ScheduleTree::makeScheduleTree(*child));
   }
 }
 
 ScheduleTreeUPtr ScheduleTree::makeScheduleTree(const ScheduleTree& tree) {
-  auto res = ScheduleTreeUPtr(new ScheduleTree(tree));
-  res->elem_ = makeElem(tree);
-  return res;
+  return makeElem(tree);
 }
 
 ScheduleTree* ScheduleTree::child(const vector<size_t>& positions) {
@@ -232,9 +228,7 @@ size_t ScheduleTree::scheduleDepth(const ScheduleTree* relativeRoot) const {
 std::unique_ptr<ScheduleTree> ScheduleTree::makeBand(
     isl::multi_union_pw_aff mupa,
     std::vector<ScheduleTreeUPtr>&& children) {
-  isl::ctx ctx = mupa.get_ctx();
-  std::unique_ptr<ScheduleTree> res(new ScheduleTree(ctx));
-  res->elem_ = ScheduleTreeElemBand::fromMultiUnionPwAff(mupa);
+  auto res = ScheduleTreeElemBand::fromMultiUnionPwAff(mupa);
   res->type_ = detail::ScheduleTreeType::Band;
   res->appendChildren(std::move(children));
   return res;
@@ -252,10 +246,7 @@ ScheduleTreeUPtr ScheduleTree::makeEmptyBand(const ScheduleTree* root) {
 std::unique_ptr<ScheduleTree> ScheduleTree::makeDomain(
     isl::union_set domain,
     std::vector<ScheduleTreeUPtr>&& children) {
-  isl::ctx ctx(domain.get_ctx());
-  std::unique_ptr<ScheduleTree> res(new ScheduleTree(ctx));
-  res->elem_ = std::unique_ptr<ScheduleTreeElemDomain>(
-      new ScheduleTreeElemDomain(domain));
+  auto res = std::unique_ptr<ScheduleTree>(new ScheduleTreeElemDomain(domain));
   res->type_ = detail::ScheduleTreeType::Domain;
   res->appendChildren(std::move(children));
   return res;
@@ -264,10 +255,8 @@ std::unique_ptr<ScheduleTree> ScheduleTree::makeDomain(
 std::unique_ptr<ScheduleTree> ScheduleTree::makeContext(
     isl::set context,
     std::vector<ScheduleTreeUPtr>&& children) {
-  isl::ctx ctx(context.get_ctx());
-  std::unique_ptr<ScheduleTree> res(new ScheduleTree(ctx));
-  res->elem_ = std::unique_ptr<ScheduleTreeElemContext>(
-      new ScheduleTreeElemContext(context));
+  auto res =
+      std::unique_ptr<ScheduleTree>(new ScheduleTreeElemContext(context));
   res->type_ = detail::ScheduleTreeType::Context;
   res->appendChildren(std::move(children));
   return res;
@@ -276,10 +265,7 @@ std::unique_ptr<ScheduleTree> ScheduleTree::makeContext(
 std::unique_ptr<ScheduleTree> ScheduleTree::makeFilter(
     isl::union_set filter,
     std::vector<ScheduleTreeUPtr>&& children) {
-  isl::ctx ctx(filter.get_ctx());
-  std::unique_ptr<ScheduleTree> res(new ScheduleTree(ctx));
-  res->elem_ = std::unique_ptr<ScheduleTreeElemFilter>(
-      new ScheduleTreeElemFilter(filter));
+  auto res = std::unique_ptr<ScheduleTree>(new ScheduleTreeElemFilter(filter));
   res->type_ = detail::ScheduleTreeType::Filter;
   res->appendChildren(std::move(children));
   return res;
@@ -299,22 +285,17 @@ std::unique_ptr<ScheduleTree> ScheduleTree::makeMappingUnsafe(
   TC_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,
-      std::move(children),
-      ScheduleTreeType::Mapping,
-      std::unique_ptr<ScheduleTreeElemMapping>(
-          new ScheduleTreeElemMapping(mapping))));
+  auto res =
+      std::unique_ptr<ScheduleTree>(new ScheduleTreeElemMapping(mapping));
+  res->appendChildren(std::move(children));
   return res;
 }
 
 std::unique_ptr<ScheduleTree> ScheduleTree::makeExtension(
     isl::union_map extension,
     std::vector<ScheduleTreeUPtr>&& children) {
-  isl::ctx ctx(extension.get_ctx());
-  ScheduleTreeUPtr res(new ScheduleTree(ctx));
-  res->elem_ = std::unique_ptr<ScheduleTreeElemExtension>(
-      new ScheduleTreeElemExtension(extension));
+  auto res =
+      std::unique_ptr<ScheduleTree>(new ScheduleTreeElemExtension(extension));
   res->type_ = detail::ScheduleTreeType::Extension;
   res->appendChildren(std::move(children));
   return res;
@@ -323,9 +304,8 @@ std::unique_ptr<ScheduleTree> ScheduleTree::makeExtension(
 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());
+  auto res =
+      std::unique_ptr<ScheduleTree>(new ScheduleTreeElemThreadSpecificMarker());
   res->type_ = detail::ScheduleTreeType::ThreadSpecificMarker;
   res->appendChildren(std::move(children));
   return res;
@@ -411,7 +391,7 @@ bool ScheduleTree::operator==(const ScheduleTree& other) const {
   if (children_.size() != other.children_.size()) {
     return false;
   }
-  if (!elemEquals(elem_.get(), other.elem_.get(), type_)) {
+  if (!elemEquals(this, &other, type_)) {
     return false;
   }
   TC_CHECK(!other.elemAs<ScheduleTreeElemSet>())

tc/core/polyhedral/schedule_tree.h

@@ -21,8 +21,8 @@
 #include <vector>
 
 #include "tc/core/check.h"
+#include "tc/core/polyhedral/mapping_types.h"
 #include "tc/core/polyhedral/options.h"
-#include "tc/core/polyhedral/schedule_tree_elem.h"
 #include "tc/core/utils/vararg.h"
 #include "tc/external/isl.h"
 
@@ -36,6 +36,23 @@ namespace detail {
 // ScheduleTree, convertible to and from isl::schedule.
 //
 struct ScheduleTree;
+struct ScheduleTreeElemSet;
+struct ScheduleTreeElemSequence;
+struct ScheduleTreeElemMapping;
+
+enum class ScheduleTreeType {
+  None,
+  Band,
+  Context,
+  Domain,
+  Extension,
+  Filter,
+  Sequence,
+  Set,
+  Mapping,
+  ThreadSpecificMarker,
+  Any,
+};
 
 } // namespace detail
 
@@ -116,19 +133,20 @@ struct ScheduleTree {
 
  private:
   ScheduleTree() = delete;
+
+ protected:
   ScheduleTree(
       isl::ctx ctx,
       std::vector<ScheduleTreeUPtr>&& children,
-      detail::ScheduleTreeType type,
-      std::unique_ptr<ScheduleTreeElemBase>&& elem)
-      : ctx_(ctx), type_(type), elem_(std::move(elem)) {
+      detail::ScheduleTreeType type)
+      : ctx_(ctx), type_(type) {
     appendChildren(std::move(children));
   }
 
   // Copy constructor for internal use only.
-  // Note that this does not copy the underlying elem_.
+  // Note that this does not account for a specific subclass of ScheduleTree.
   // All callers should use makeScheduleTree(const ScheduleTree&) instead,
-  // which copies the underlying elem_ as well as keeps the memory
+  // which dispatches the copying to subclasses as well as keeps the memory
   // management scheme consistent.
   ScheduleTree(const ScheduleTree& st);
 
@@ -394,23 +412,19 @@ struct ScheduleTree {
   static ScheduleTreeUPtr
   fromList(detail::ScheduleTreeType type, Arg&& arg, Args&&... args) {
     static_assert(
-        std::is_base_of<ScheduleTreeElemBase, T>::value,
-        "Can only construct elements derived from ScheduleTreeElemBase");
+        std::is_base_of<ScheduleTree, T>::value,
+        "Can only construct elements derived from ScheduleTree");
     static_assert(
         std::is_same<
             typename std::remove_reference<Arg>::type,
             ScheduleTreeUPtr>::value,
         "Arguments must be rvalue references to ScheduleTreeUPtr");
 
-    auto ctx = arg->ctx_;
-    std::vector<ScheduleTreeUPtr> children =
-        vectorFromArgs(std::forward<Arg>(arg), std::forward<Args>(args)...);
-
-    auto res = ScheduleTreeUPtr(new ScheduleTree(
-        ctx,
-        std::move(children),
-        type,
-        std::unique_ptr<ScheduleTreeElemBase>(new T)));
+    auto ctx = arg->ctx_; // FIXME: pass this to the constructor of T
+                          // when possible
+    auto res = ScheduleTreeUPtr(new T);
+    res->appendChildren(
+        vectorFromArgs(std::forward<Arg>(arg), std::forward<Args>(args)...));
 
     if (type == detail::ScheduleTreeType::Sequence ||
         type == detail::ScheduleTreeType::Set) {
@@ -453,7 +467,7 @@ struct ScheduleTree {
       const ScheduleTree* tree,
       detail::ScheduleTreeType type);
 
-  // View elem_ as the specified type.
+  // View this tree node as the specified type.
   // Returns nullptr if this is not the proper type.
   // Inline impl for now, does not justify an extra -inl.h file
   template <typename T>
@@ -464,13 +478,12 @@ struct ScheduleTree {
   template <typename T>
   const T* elemAs() const {
     static_assert(
-        std::is_base_of<ScheduleTreeElemBase, T>::value,
-        "Must call with a class derived from ScheduleTreeElemBase");
+        std::is_base_of<ScheduleTree, T>::value,
+        "Must call with a class derived from ScheduleTree");
     if (type_ != T::NodeType) {
       return nullptr;
     }
-    return static_cast<const T*>(
-        const_cast<const ScheduleTreeElemBase*>(elem_.get()));
+    return static_cast<const T*>(this);
   }
 
   virtual ScheduleTreeType type() const {
@@ -494,7 +507,6 @@ struct ScheduleTree {
 
  public:
   detail::ScheduleTreeType type_{detail::ScheduleTreeType::None};
-  std::unique_ptr<ScheduleTreeElemBase> elem_{nullptr};
 };
 
 } // namespace detail

tc/core/polyhedral/schedule_tree_elem.cc

@@ -40,7 +40,7 @@ namespace detail {
 std::unique_ptr<ScheduleTreeElemBand> ScheduleTreeElemBand::fromMultiUnionPwAff(
     isl::multi_union_pw_aff mupa) {
   isl::ctx ctx(mupa.get_ctx());
-  std::unique_ptr<ScheduleTreeElemBand> band(new ScheduleTreeElemBand);
+  std::unique_ptr<ScheduleTreeElemBand> band(new ScheduleTreeElemBand(ctx));
   band->mupa_ = mupa.floor();
   size_t n = band->mupa_.size();
   band->coincident_ = vector<bool>(n, false);
@@ -202,8 +202,8 @@ bool ScheduleTreeElemSet::operator==(const ScheduleTreeElemSet& other) const {
 }
 
 bool elemEquals(
-    const ScheduleTreeElemBase* e1,
-    const ScheduleTreeElemBase* e2,
+    const ScheduleTree* e1,
+    const ScheduleTree* e2,
     detail::ScheduleTreeType type) {
 #define ELEM_EQUALS_CASE(CLASS)                                              \
   else if (type == CLASS::NodeType) {                                        \
@@ -222,7 +222,7 @@ bool elemEquals(
   ELEM_EQUALS_CASE(ScheduleTreeElemSequence)
   ELEM_EQUALS_CASE(ScheduleTreeElemSet)
   else {
-    LOG(FATAL) << "NYI: ScheduleTreeElemBase::operator== for type: "
+    LOG(FATAL) << "NYI: ScheduleTree::operator== for type: "
                << static_cast<int>(type);
   }