Merge pull request #558 from facebookresearch/schedule-tree-evolution

ScheduleTree*: hide constructors

Author: ftynse

tc/core/polyhedral/schedule_isl_conversion.cc

@@ -236,12 +236,14 @@ namespace {
 
 std::unique_ptr<ScheduleTreeBand> fromIslScheduleNodeBand(
     isl::schedule_node_band b) {
-  auto res = ScheduleTreeBand::fromMultiUnionPwAff(b.get_partial_schedule());
-  res->permutable_ = b.get_permutable();
-  for (size_t i = 0; i < b.n_member(); ++i) {
-    res->coincident_[i] = b.member_get_coincident(i);
+  auto n = b.n_member();
+  std::vector<bool> coincident(n, false);
+  std::vector<bool> unroll(n, false);
+  for (size_t i = 0; i < n; ++i) {
+    coincident[i] = b.member_get_coincident(i);
   }
-  return res;
+  return ScheduleTreeBand::make(
+      b.get_partial_schedule(), b.get_permutable(), coincident, unroll);
 }
 
 std::unique_ptr<ScheduleTree> elemFromIslScheduleNode(isl::schedule_node node) {
@@ -250,19 +252,19 @@ std::unique_ptr<ScheduleTree> elemFromIslScheduleNode(isl::schedule_node node) {
     return fromIslScheduleNodeBand(band);
   } else if (auto context = node.as<isl::schedule_node_context>()) {
     auto c = context.get_context();
-    return std::unique_ptr<ScheduleTreeContext>(new ScheduleTreeContext(c));
+    return ScheduleTreeContext::make(c);
   } else if (auto domain = node.as<isl::schedule_node_domain>()) {
     auto c = domain.get_domain();
-    return std::unique_ptr<ScheduleTreeDomain>(new ScheduleTreeDomain(c));
+    return ScheduleTreeDomain::make(c);
   } else if (auto expansion = node.as<isl::schedule_node_expansion>()) {
     LOG(FATAL) << "expansion nodes not supported";
     return nullptr;
   } else if (auto extension = node.as<isl::schedule_node_extension>()) {
     auto e = extension.get_extension();
-    return std::unique_ptr<ScheduleTreeExtension>(new ScheduleTreeExtension(e));
+    return ScheduleTreeExtension::make(e);
   } else if (auto filter = node.as<isl::schedule_node_filter>()) {
     auto f = filter.get_filter();
-    return std::unique_ptr<ScheduleTreeFilter>(new ScheduleTreeFilter(f));
+    return ScheduleTreeFilter::make(f);
   } else if (auto guard = node.as<isl::schedule_node_guard>()) {
     LOG(FATAL) << "guard nodes not supported";
     return nullptr;
@@ -273,9 +275,9 @@ std::unique_ptr<ScheduleTree> elemFromIslScheduleNode(isl::schedule_node node) {
     LOG(FATAL) << "ScheduleTree::make called on explicit leaf";
     return nullptr;
   } else if (node.isa<isl::schedule_node_sequence>()) {
-    return std::unique_ptr<ScheduleTreeSequence>(new ScheduleTreeSequence(ctx));
+    return ScheduleTreeSequence::make(ctx);
   } else if (node.isa<isl::schedule_node_set>()) {
-    return std::unique_ptr<ScheduleTreeSet>(new ScheduleTreeSet(ctx));
+    return ScheduleTreeSet::make(ctx);
   }
   LOG(FATAL) << "NYI: ScheduleTree from type: "
              << isl_schedule_node_get_type(node.get());

tc/core/polyhedral/schedule_tree.cc

@@ -122,31 +122,6 @@ vector<ScheduleTree*> ancestorsInSubTree(
   }
   return res;
 }
-
-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) {
-    LOG(FATAL) << "Hit Error node!";
-  }
-  ELEM_MAKE_CASE(ScheduleTreeBand)
-  ELEM_MAKE_CASE(ScheduleTreeContext)
-  ELEM_MAKE_CASE(ScheduleTreeDomain)
-  ELEM_MAKE_CASE(ScheduleTreeExtension)
-  ELEM_MAKE_CASE(ScheduleTreeFilter)
-  ELEM_MAKE_CASE(ScheduleTreeMapping)
-  ELEM_MAKE_CASE(ScheduleTreeSequence)
-  ELEM_MAKE_CASE(ScheduleTreeSet)
-  ELEM_MAKE_CASE(ScheduleTreeThreadSpecificMarker)
-
-#undef ELEM_MAKE_CASE
-
-  LOG(FATAL) << "NYI: ScheduleTree from type: " << static_cast<int>(st.type_);
-  return nullptr;
-}
 } // namespace
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -163,7 +138,7 @@ ScheduleTree::ScheduleTree(const ScheduleTree& st)
 }
 
 ScheduleTreeUPtr ScheduleTree::makeScheduleTree(const ScheduleTree& tree) {
-  return makeElem(tree);
+  return tree.clone();
 }
 
 ScheduleTree* ScheduleTree::child(const vector<size_t>& positions) {
@@ -226,8 +201,10 @@ size_t ScheduleTree::scheduleDepth(const ScheduleTree* relativeRoot) const {
 std::unique_ptr<ScheduleTree> ScheduleTree::makeBand(
     isl::multi_union_pw_aff mupa,
     std::vector<ScheduleTreeUPtr>&& children) {
-  auto res = ScheduleTreeBand::fromMultiUnionPwAff(mupa);
-  res->appendChildren(std::move(children));
+  std::vector<bool> coincident(mupa.size(), false);
+  std::vector<bool> unroll(mupa.size(), false);
+  auto res = ScheduleTreeBand::make(
+      mupa, false, coincident, unroll, std::move(children));
   return res;
 }
 
@@ -243,25 +220,19 @@ ScheduleTreeUPtr ScheduleTree::makeEmptyBand(const ScheduleTree* root) {
 std::unique_ptr<ScheduleTree> ScheduleTree::makeDomain(
     isl::union_set domain,
     std::vector<ScheduleTreeUPtr>&& children) {
-  auto res = std::unique_ptr<ScheduleTree>(new ScheduleTreeDomain(domain));
-  res->appendChildren(std::move(children));
-  return res;
+  return ScheduleTreeDomain::make(domain, std::move(children));
 }
 
 std::unique_ptr<ScheduleTree> ScheduleTree::makeContext(
     isl::set context,
     std::vector<ScheduleTreeUPtr>&& children) {
-  auto res = std::unique_ptr<ScheduleTree>(new ScheduleTreeContext(context));
-  res->appendChildren(std::move(children));
-  return res;
+  return ScheduleTreeContext::make(context, std::move(children));
 }
 
 std::unique_ptr<ScheduleTree> ScheduleTree::makeFilter(
     isl::union_set filter,
     std::vector<ScheduleTreeUPtr>&& children) {
-  auto res = std::unique_ptr<ScheduleTree>(new ScheduleTreeFilter(filter));
-  res->appendChildren(std::move(children));
-  return res;
+  return ScheduleTreeFilter::make(filter, std::move(children));
 }
 
 std::unique_ptr<ScheduleTree> ScheduleTree::makeMappingUnsafe(
@@ -278,28 +249,19 @@ 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();
-  auto res =
-      std::unique_ptr<ScheduleTree>(new ScheduleTreeMapping(ctx, mapping));
-  res->appendChildren(std::move(children));
-  return res;
+  return ScheduleTreeMapping::make(ctx, mapping, std::move(children));
 }
 
 std::unique_ptr<ScheduleTree> ScheduleTree::makeExtension(
     isl::union_map extension,
     std::vector<ScheduleTreeUPtr>&& children) {
-  auto res =
-      std::unique_ptr<ScheduleTree>(new ScheduleTreeExtension(extension));
-  res->appendChildren(std::move(children));
-  return res;
+  return ScheduleTreeExtension::make(extension, std::move(children));
 }
 
 std::unique_ptr<ScheduleTree> ScheduleTree::makeThreadSpecificMarker(
     isl::ctx ctx,
     std::vector<ScheduleTreeUPtr>&& children) {
-  auto res =
-      std::unique_ptr<ScheduleTree>(new ScheduleTreeThreadSpecificMarker(ctx));
-  res->appendChildren(std::move(children));
-  return res;
+  return ScheduleTreeThreadSpecificMarker::make(ctx, std::move(children));
 }
 
 ////////////////////////////////////////////////////////////////////////////////

tc/core/polyhedral/schedule_tree.h

@@ -401,11 +401,11 @@ struct ScheduleTree {
         "Arguments must be rvalue references to ScheduleTreeUPtr");
 
     auto ctx = arg->ctx_;
-    auto res = new T(ctx);
-    flattenSequenceOrSet(res);
+    auto res = T::make(ctx);
+    flattenSequenceOrSet(res.get());
     res->appendChildren(
         vectorFromArgs(std::forward<Arg>(arg), std::forward<Args>(args)...));
-    return ScheduleTreeUPtr(res);
+    return res;
   }
 
   // Make a (deep) copy of "tree".
@@ -465,6 +465,10 @@ struct ScheduleTree {
   // Note that this function does _not_ output the child trees.
   virtual std::ostream& write(std::ostream& os) const = 0;
 
+  // Clone the current node.
+  // Note that this function does _not_ clone the child trees.
+  virtual ScheduleTreeUPtr clone() const = 0;
+
   //
   // Data members
   //