Rename bound to schedule, add graph related utils

include/tvm/tensor.h

@@ -239,4 +239,13 @@ DEFINE_OVERLOAD_SLICE_BINARY_OP(>);  // NOLINT(*)
 DEFINE_OVERLOAD_SLICE_BINARY_OP(<);  // NOLINT(*)
 
 }  // namespace tvm
+
+namespace std {
+template <>
+struct hash<::tvm::Operation> {
+  std::size_t operator()(const ::tvm::Operation& k) const {
+    return k.hash();
+  }
+};
+}
 #endif  // TVM_TENSOR_H_

src/README.md

@@ -2,5 +2,5 @@
 
 - c_api  C API related functions
 - lang The definition of DSL related data structure
+- schedule The operations on the schedule graph before converting to IR.
 - pass The optimization pass on the IR structure
-- bound Bound inference logics.

src/bound/bound.cc → src/schedule/bound.cc

@@ -8,7 +8,7 @@
 #include "./bound.h"
 
 namespace tvm {
-namespace bound {
+namespace schedule {
 
 // result = ceil((a / b)), both a and b are positive integer
 inline Expr DivCeil(Expr a, Expr b) {
@@ -89,5 +89,10 @@ void PassUp(const Schedule& s,
   }
 }
 
-}  // namespace bound
+
+std::unordered_map<IterVar, Range> InferBound(Schedule sch) {
+  return {};
+}
+
+}  // namespace schedule
 }  // namespace tvm

src/bound/bound.h → src/schedule/bound.h

@@ -3,15 +3,15 @@
  * \file bound.h
  * \brief The bound inference logics on the schedule.
  */
-#ifndef TVM_BOUND_BOUND_H_
-#define TVM_BOUND_BOUND_H_
+#ifndef TVM_SCHEDULE_BOUND_H_
+#define TVM_SCHEDULE_BOUND_H_
 
 #include <tvm/expr.h>
 #include <tvm/schedule.h>
 #include <unordered_map>
 
 namespace tvm {
-namespace bound {
+namespace schedule {
 
 /*!
  * \brief Infer the bound of all iteration variables relates to the schedule.
@@ -21,7 +21,7 @@ namespace bound {
  */
 std::unordered_map<IterVar, Range> InferBound(Schedule sch);
 
-}  // namespace bound
+}  // namespace schedule
 }  // namespace tvm
 
-#endif  // TVM_BOUND_BOUND_H_
+#endif  // TVM_SCHEDULE_BOUND_H_

src/schedule/graph.cc

+/*!
+ *  Copyright (c) 2016 by Contributors
+ * \file graph.cc
+ * \brief Utilities to get information about schedule graph.
+ */
+#include <tvm/ir.h>
+#include <tvm/ir_visitor.h>
+#include <unordered_set>
+#include "./int_set.h"
+#include "./graph.h"
+
+namespace tvm {
+namespace schedule {
+
+// construct a read graph that gives readers of each operation
+// that the root depend on
+ReadGraph CreateReadGraph(Operation root) {
+  std::unordered_map<Operation, std::vector<Tensor> > rmap;
+  rmap[root] = {};
+  std::vector<Operation> stack{root};
+  while (!stack.empty()) {
+    Operation r = stack.back();
+    stack.pop_back();
+    auto& vec = rmap.at(r);
+    if (r.as<ComputeOpNode>()) {
+      auto fvisit = [&vec, &rmap, &stack](const NodeRef& n) {
+        auto *call = n.as<ir::Call>();
+        if (call != nullptr && call->func.defined()) {
+          Tensor t(call->func.node_);
+          vec.push_back(t);
+          if (t->op.defined() && rmap.count(t->op) == 0) {
+            rmap[t->op] = {}; stack.push_back(t->op);
+          }
+        }
+      };
+      ir::PostOrderVisit(r.as<ComputeOpNode>()->body, fvisit);
+    } else {
+      LOG(FATAL) << "unknown operation mode";
+    }
+  }
+  return rmap;
+}
+
+
+void PostDFSOrder(const Operation& op,
+                    const ReadGraph& g,
+                    std::unordered_set<Operation>* visited,
+                    std::vector<Operation>* post_order) {
+  visited->insert(op);
+  for (const auto& t : g.at(op)) {
+    if (t->op.defined() && !visited->count(t->op)) {
+      PostDFSOrder(t->op, g, visited, post_order);
+    }
+  }
+  post_order->push_back(op);
+}
+
+std::vector<Operation> PostDFSOrder(
+    const Operation& root, const ReadGraph& g) {
+  std::unordered_set<Operation> visited;
+  std::vector<Operation> post_order;
+  PostDFSOrder(root, g, &visited, &post_order);
+  return post_order;
+}
+
+}  // namespace schedule
+}  // namespace tvm

src/schedule/graph.h

+/*!
+ *  Copyright (c) 2016 by Contributors
+ * \file graph.h
+ * \brief Utilities to get information about schedule graph.
+ */
+#ifndef TVM_SCHEDULE_GRAPH_H_
+#define TVM_SCHEDULE_GRAPH_H_
+
+#include <tvm/expr.h>
+#include <tvm/schedule.h>
+#include <unordered_map>
+#include <vector>
+
+namespace tvm {
+namespace schedule {
+
+/*!
+ * \brief data structure of Operation->Tensors it reads
+ */
+using ReadGraph = std::unordered_map<Operation, std::vector<Tensor> >;
+
+/*!
+ * \brief Get read graph of each operation to all the
+ *  Tensors that it directly depends on.
+ *
+ *  The result map contains Operations needed to finish root Operation.
+ * \param root The root operation.
+ * \return The result map.
+ */
+ReadGraph CreateReadGraph(const Operation& root);
+
+/*!
+ * \brief Get a post DFS ordered of operations in the graph.
+ * \param root The root of the graph.
+ * \param g The read graph.
+ * \return vector order of Operations in PostDFS order.
+ *
+ * \note PostDFSOrder is a special case of Topoligical order,
+ *   and can be used when topoligical order is needed.
+ */
+std::vector<Operation> PostDFSOrder(
+    const Operation& root, const ReadGraph& g);
+
+}  // namespace schedule
+}  // namespace tvm
+
+#endif  // TVM_SCHEDULE_GRAPH_H_

src/bound/int_set.cc → src/schedule/int_set.cc

@@ -7,7 +7,7 @@
 #include "./int_set.h"
 
 namespace tvm {
-namespace bound {
+namespace schedule {
 
 using namespace ir;
 
@@ -338,6 +338,5 @@ IntSet Eval(Expr e,
   return m.Eval(e);
 }
 
-}  // namespace bound
+}  // namespace schedule
 }  // namespace tvm
-

src/bound/int_set.h → src/schedule/int_set.h

@@ -3,14 +3,14 @@
  * \file int_set.h
  * \brief Abstraction for all integer set operations.
  */
-#ifndef TVM_BOUND_INT_SET_H_
-#define TVM_BOUND_INT_SET_H_
+#ifndef TVM_SCHEDULE_INT_SET_H_
+#define TVM_SCHEDULE_INT_SET_H_
 
 #include <tvm/expr.h>
 #include <tvm/schedule.h>
 
 namespace tvm {
-namespace bound {
+namespace schedule {
 
 // internal node container of int set.
 class IntSetNode;
@@ -97,7 +97,7 @@ void PassUp(const FuseNode* s,
  */
 IntSet Union(const Array<IntSet>& sets);
 
-}  // namespace bound
+}  // namespace schedule
 }  // namespace tvm
 
-#endif  // TVM_BOUND_INT_SET_H_
+#endif  // TVM_SCHEDULE_INT_SET_H_