[RELAY][PASS] General OpFusion. (#2090)

include/tvm/relay/expr.h

@@ -429,6 +429,16 @@ inline const TTypeNode* ExprNode::type_as() const {
   return node;
 }
 
+/*!
+ * \brief Print node as text format.
+ * \param node The node to be printed.
+ * \param annotate An optional callback function for attaching
+ *        additional comment block to an expr.
+ * \return The text representation.
+ */
+std::string RelayPrint(
+    const NodeRef& node,
+    runtime::TypedPackedFunc<std::string(Expr)> annotate = nullptr);
 }  // namespace relay
 }  // namespace tvm
 #endif  // TVM_RELAY_EXPR_H_

include/tvm/runtime/packed_func.h

@@ -161,6 +161,8 @@ class TypedPackedFunc<R(Args...)> {
   using TSelf = TypedPackedFunc<R(Args...)>;
   /*! \brief default constructor */
   TypedPackedFunc() {}
+  /*! \brief constructor from null */
+  TypedPackedFunc(std::nullptr_t null) {}  // NOLINT(*)
   /*!
    * \brief construct by wrap a PackedFunc
    *

python/tvm/relay/base.py

@@ -22,15 +22,20 @@ def register_relay_node(type_key=None):
 
 
 class RelayNode(NodeBase):
-    def astext(self):
+    """Base class of all relay node."""
+    def astext(self, annotate=None):
         """Get the text format of the expression.
 
         Returns
         -------
         text : str
             The text format of the expression.
+
+        annotate: Optional[relay.Expr->str]
+            Optional annotate function to provide additional
+            information in the comment block.
         """
-        return _expr._text_print(self)
+        return _expr.RelayPrint(self, annotate)
 
 
 @register_relay_node

python/tvm/relay/build_module.py

@@ -173,11 +173,13 @@ def build(func,
     else:
         tophub_context = autotvm.util.EmptyContext()
 
+    cfg = BuildConfig.current
+
     with tophub_context:
         func = optimize(func)
         # Fuse ops before running code gen
         func = ir_pass.infer_type(func)
-        func = ir_pass.fuse_ops(func)
+        func = ir_pass.fuse_ops(func, cfg.opt_level)
         # Graph code generation
         func = ir_pass.infer_type(func)
         graph_gen = _graph_gen.GraphRuntimeCodegen(mod=None, target=target)

python/tvm/relay/expr.py

@@ -6,7 +6,6 @@ from numbers import Number as _Number
 import numpy as _np
 from .base import RelayNode, register_relay_node
 from . import _make
-from . import _expr
 from . import ty as _ty
 from .._ffi import base as _base
 from .. import nd as _nd
@@ -477,7 +476,7 @@ class TupleWrapper(object):
         text : str
             The text format of the tuple expression.
         """
-        return _expr._text_print(self.tuple_value)
+        return self.tuple_value.astext()
 
     def __getitem__(self, index):
         if index >= len(self):

python/tvm/relay/ir_pass.py

@@ -259,7 +259,7 @@ def structural_hash(value):
         raise TypeError(msg)
 
 
-def fuse_ops(expr):
+def fuse_ops(expr, opt_level=1):
     """Fuse operators in expr together.
 
     Parameters
@@ -267,9 +267,12 @@ def fuse_ops(expr):
     expr : tvm.relay.Expr
         The input expression.
 
+    opt_level : int
+        The level of fuse optimization.
+
     Returns
     -------
     transformed_expr : tvm.relay.Expr
         Transformed expression, containing fused result.
     """
-    return _ir_pass.FuseOps(expr)
+    return _ir_pass.FuseOps(expr, opt_level)

src/common/arena.h

@@ -38,11 +38,29 @@ class Arena {
   /*!
    * \brief Allocate a space from Arena for type T
    * \param T the data type to be allocated
+   * \note The space of T is not initialized.
    */
   template<typename T>
-  T* Alloc() {
+  T* allocate_() {
     return static_cast<T*>(Alloc(sizeof(T), alignof(T)));
   }
+  /*!
+   * \brief Create a new instance of type T.
+   * \param args The constructor argument.
+   * \tparam T the type to be created.
+   * \tparam Args Arguments to the constructor.
+   *
+   * \return The allocated object.
+   * \note The type T must be simple type, or only contain
+   *  memory allocated from the same arena.
+   *  Otherwise the destructor needs to be called explicitly.
+   */
+  template<typename T, typename... Args>
+  T* make(Args&&... args) {
+    T* ptr = allocate_<T>();
+    new (ptr) T(std::forward<Args>(args)...);
+    return ptr;
+  }
 
  private:
   // page size 16 KB
@@ -87,6 +105,44 @@ class Arena {
   }
 };
 
+/*!
+ * \brief Link list node
+ * \tparam T the content data type
+ */
+template<typename T>
+struct LinkNode {
+  /*! \brief The content value */
+  T value;
+  /*! \brief pointer to the next location */
+  LinkNode<T>* next{nullptr};
+};
+/*!
+ * \brief LinkedList structure
+ * \tparam T the content data type
+ * \note This is a simple data structure that can be used together with the arena.
+ * \sa LinkNode
+ */
+template<typename T>
+struct LinkedList {
+  /*! \brief Head pointer */
+  LinkNode<T>* head{nullptr};
+  /*! \brief Tail pointer */
+  LinkNode<T>* tail{nullptr};
+  /*!
+   * \brief Push a new node to the end of the linked list.
+   * \param node The node to be pushed.
+   */
+  void Push(LinkNode<T>* node) {
+    node->next = nullptr;
+    if (this->tail != nullptr) {
+      this->tail->next = node;
+      this->tail = node;
+    } else {
+      head = tail = node;
+    }
+  }
+};
+
 }  // namespace common
 }  // namespace tvm
 #endif  // TVM_COMMON_ARENA_H_

src/relay/backend/compile_engine.cc

@@ -109,6 +109,29 @@ class ScheduleGetter :
     return {};
   }
 
+  Array<Tensor> VisitExpr_(const ConstantNode* op) final {
+    CHECK(op->is_scalar());
+    void* data = op->data->data;
+    DataType dtype = TVMType2Type(op->data->dtype);
+    Tensor value = tvm::compute({}, [&](const Array<tvm::Var>&) {
+        if (dtype == Int(32)) {
+          return make_const(dtype, static_cast<const int32_t*>(data)[0]);
+        } else if (dtype == Int(64)) {
+          return make_const(dtype, static_cast<const int64_t*>(data)[0]);
+        } else if (dtype == Float(32)) {
+          return make_const(dtype, static_cast<const float*>(data)[0]);
+        } else if (dtype == Float(64)) {
+          return make_const(dtype, static_cast<const double*>(data)[0]);
+        } else if (dtype == Bool()) {
+          return make_const(dtype, static_cast<const uint8_t*>(data)[0]);
+        } else {
+          LOG(FATAL) << "not handled";
+          return tvm::Expr();
+        }
+      });
+    return {value};
+  }
+
   Array<Tensor> VisitExpr_(const CallNode* call_node) final {
     static auto fcompute =
         Op::GetAttr<FTVMCompute>("FTVMCompute");

src/relay/ir/text_printer.cc

@@ -125,6 +125,8 @@ class TextPrinter :
     public TypeFunctor<void (const Type&, std::ostream& os)>,  // NOLINT(*)
     public AttrFunctor<void (const NodeRef&, std::ostream& os)> { // NOLINT(*)
  public:
+  explicit TextPrinter(runtime::TypedPackedFunc<std::string(Expr)> annotate)
+      : annotate_(annotate) {}
   /*!
    * \brief Print a node to string.
    * \param node.
@@ -279,11 +281,11 @@ class TextPrinter :
 
   TextValue VisitExpr_(const CallNode* op) final {
     // possibly through meta-data
-    TextValue call_op = GetValue(op->op);
     std::vector<TextValue> args;
     for (Expr arg : op->args) {
       args.emplace_back(GetValue(arg));
     }
+    TextValue call_op = GetValue(op->op);
     TextValue id = this->AllocTempVar();
     this->PrintIndent();
 
@@ -532,7 +534,9 @@ class TextPrinter :
    */
   void PrintOptionalInfo(const Expr& expr) {
     // additional information in comment.
-    if (expr->checked_type_.defined()) {
+    if (annotate_ != nullptr) {
+      stream_ << " # " << annotate_(expr);
+    } else if (expr->checked_type_.defined()) {
       stream_ << " # ty=";
       this->PrintType(expr->checked_type(), stream_);
     }
@@ -678,7 +682,10 @@ class TextPrinter :
       name = "%" + name;
     }
     TextValue val(GetUniqueName(name));
-    CHECK(!memo_.count(var)) << "Duplicated variable " << var;
+    // still print if ir is malformed, but show the error.
+    if (memo_.count(var)) {
+      memo_[var] = TextValue(val.name + "-malformed-ir");
+    }
     memo_[var] = val;
     return val;
   }
@@ -686,6 +693,8 @@ class TextPrinter :
  private:
   class AttrPrinter;
   friend class AttrPrinter;
+  /*! \brief additional comment function */
+  runtime::TypedPackedFunc<std::string(Expr)> annotate_;
   /*! \brief meta data context */
   TextMetaDataContext meta_;
   /*! \brief Check whether scope is still valid */
@@ -776,12 +785,15 @@ void TextPrinter::PrintCallAttrs(const Expr& op,
   os << ", " << meta_.GetMetaNode(attrs);
 }
 
-std::string RelayPrint(const NodeRef& node) {
-  return TextPrinter().Print(node);
+std::string RelayPrint(const NodeRef& node,
+                       runtime::TypedPackedFunc<std::string(Expr)> annotate) {
+  return TextPrinter(annotate).Print(node);
 }
 
-TVM_REGISTER_API("relay._expr._text_print")
-.set_body_typed<std::string(const NodeRef&)>(RelayPrint);
+TVM_REGISTER_API("relay._expr.RelayPrint")
+.set_body_typed<std::string(
+    const NodeRef&,
+    runtime::TypedPackedFunc<std::string(Expr)>)>(RelayPrint);
 
 }  // namespace relay
 }  // namespace tvm

src/relay/pass/fold_scale_axis.cc

@@ -10,6 +10,7 @@
 #include <tvm/relay/attrs/nn.h>
 #include <tvm/relay/expr_functor.h>
 #include "pattern_util.h"
+#include "pass_util.h"
 #include "../op/nn/layout.h"
 
 namespace tvm {
@@ -580,23 +581,6 @@ using FBackwardTransform = TypedPackedFunc<
 //----------------------------------------------
 // Generic Visitors for FScaleAxisBackward
 //----------------------------------------------
-/*!
- * \brief Get reference counter of each internal ExprNode in body.
- * \param body The body expression.
- * \return The reference count mapping.
- */
-std::unordered_map<const Node*, size_t>
-GetExprRefCount(const Expr& body) {
-  class ExprRefCounter : private ExprVisitor {
-   public:
-    std::unordered_map<const Node*, size_t>
-    Get(const Expr& body) {
-      this->VisitExpr(body);
-      return std::move(this->visit_counter_);
-    }
-  };
-  return ExprRefCounter().Get(body);
-}
 
 class BackwardPrep : private ExprVisitor {
  public:

src/relay/pass/fuse_ops.cc

@@ -9,13 +9,686 @@
 #include <tvm/ir_operator.h>
 #include <tvm/relay/pass.h>
 #include <tvm/relay/expr_functor.h>
+#include <tvm/relay/op_attr_types.h>
+#include "../../common/arena.h"
+
 
 namespace tvm {
 namespace relay {
 
-// Simple fuser that only makes each operator function as primitive.
-class SimpleFuser : public ExprMutator {
+/*
+  Note on Fusing algorithm:
+
+  The main challenge of genenral fusor is to handle possible diamond shape branches,
+  in the following graph, conv2d can be fused to elemwise add.
+
+            conv2d
+            /  |  \
+           /   |   \
+         op    op   op
+          \    |    /
+           \   |   /
+          elemwise add
+               |
+
+  However, at the point of conv2d we do not necessarily know that all its future path
+  will merge at the elemwise add. The new fusor algorithm applies post-dominator analysis.
+  The immediate post-dominator of a node defined by the closest node where all the future path goes into.
+  In the above case, the elemwise add is the post-dominator of conv2d. The general algorithm is as follows:
+
+  - Construct a DAG of dataflow graph for dominator analysis
+  - Construct a post-dominator tree which gives immediate post dominator of each node.
+  - Run fusion algorithm with the given post-dominator information.
+
+  Note that, because we run analysis on a DAG, we use a single pass post-dominator
+  tree construction algorithm via LCA, which is simpler than the full version that handles cycles.
+
+  The fusion algorithm traverses from each node and checks if it can be fused to its
+  immediate post dominator. It has to check the following things:
+
+  - CheckPath: check all the path between a node and its immediate post-dominator
+               satiesfies the fuse condition.
+  - Note that these intermediate node can already be fused with another nodes, the algorithm
+      will still run correctly.
+  - CommitFuse: mark all the nodes between source and post-dominator as the same group.
+  - We use an Union-Find data structure to manage the groups.
+*/
+using common::LinkNode;
+using common::LinkedList;
+
+/*!
+ * \brief Indexed data flow graph in forward direction.
+ *  This is a temporary data structure used for operator fusion analysis.
+ *
+ *  This data structure only captures the dataflow fragement and
+ *  could ignore blocks like let by simply ordering each dataflow block
+ *  and mark the output node as extern_ref;
+ */
+class IndexedForwardGraph {
+ public:
+  struct Node;
+  /*!
+   * The forward edge in the dataflow graph.
+   */
+  struct Edge {
+    /*! \brief The corresponding node */
+    Node* node{nullptr};
+    /*! \brief The respective pattern of this op */
+    OpPatternKind pattern{kOpaque};
+  };
+  /*! \brief A node in the graph. */
+  struct Node {
+    /*! \brief weak reference to the corresponding edge. */
+    const tvm::Node* ref{nullptr};
+    /*! \brief The index of the node in topological order. */
+    size_t index{0};
+    /*! \brief Whether this node is referenced by external source */
+    bool extern_ref{false};
+    /*! \brief The general pattern in the node */
+    OpPatternKind pattern{kOpaque};
+    /*! \brief The outputs of the node. */
+    LinkedList<Edge> outputs;
+  };
+  /*! \brief The node map that maps node to graph */
+  std::unordered_map<const tvm::Node*, Node*> node_map;
+  /*! \brief All the nodes in post DFS order */
+  std::vector<Node*> post_dfs_order;
+
+  /*! \brief Dump the graph into string. */
+  void DebugDump() {
+    std::ostringstream os;
+    for (size_t i = 0; i < post_dfs_order.size(); ++i) {
+      Node* node = post_dfs_order[i];
+      os << "node[" << i << "], "
+         << GetRef<NodeRef>(node->ref)
+         << " outputs=[";
+      for (auto* link = node->outputs.head; link != nullptr; link = link->next) {
+        os << link->value.node->index << ", ";
+      }
+      os << "]\n";
+    }
+    LOG(INFO) << os.str();
+  }
+  /*!
+   * \brief create a indexed forward graph.
+   * \param arena The arena used for data allocation.
+   * \param body The body of the expression to create a graph.
+   */
+  static IndexedForwardGraph Create(common::Arena* arena, const Expr& body);
+
+ private:
+  class Creator;
+};
+
+// Creator of post dominator tree of the dataflow
+class IndexedForwardGraph::Creator : private ExprVisitor {
+ public:
+  explicit Creator(common::Arena* arena)
+      : arena_(arena) {}
+
+  IndexedForwardGraph Prepare(const Expr& body) {
+    this->Update(body, nullptr, kOpaque);
+    this->VisitExpr(body);
+    return std::move(graph_);
+  }
+
+ private:
+  /*! \brief allocator of all the internal node object */
+  common::Arena* arena_;
+  // The output.
+  IndexedForwardGraph graph_;
+  // attribute equal comparator
+  AttrsEqual attr_equal_;
+  // Update the message stored at the node.
+  void Update(const Expr& node,
+              IndexedForwardGraph::Node* parent,
+              OpPatternKind pattern) {
+    const tvm::Node* key = node.get();
+    IndexedForwardGraph::Node* current;
+    auto it = graph_.node_map.find(key);
+    if (it != graph_.node_map.end()) {
+      current = it->second;
+    } else {
+      current = arena_->make<IndexedForwardGraph::Node>();
+      graph_.node_map[key] = current;
+    }
+    if (parent != nullptr) {
+      auto* link = arena_->make<LinkNode<IndexedForwardGraph::Edge> >();
+      link->value.node = parent;
+      link->value.pattern = pattern;
+      current->outputs.Push(link);
+    } else {
+      current->extern_ref = true;
+    }
+  }
+  void AddNode(const tvm::Node* key) {
+    auto it = graph_.node_map.find(key);
+    CHECK(it != graph_.node_map.end())
+        << "Cannot find node " << GetRef<NodeRef>(key);
+    IndexedForwardGraph::Node* node = it->second;
+    CHECK(node->ref == nullptr);
+    node->ref = key;
+    node->index = graph_.post_dfs_order.size();
+    graph_.post_dfs_order.push_back(node);
+  }
+
+  // Post order tree
+  void VisitExpr_(const FunctionNode* op) {
+    for (auto param : op->params) {
+      this->Update(param, nullptr, kOpaque);
+    }
+    this->Update(op->body, nullptr, kOpaque);
+    ExprVisitor::VisitExpr_(op);
+  }
+
+  void VisitExpr_(const ConstantNode* op) {
+    this->AddNode(op);
+    Node* node = graph_.node_map.at(op);
+    DataType dtype = TVMType2Type(op->data->dtype);
+    // This rule must be consistent with code generator.
+    bool is_simple_const = (
+        dtype == Int(32) ||
+        dtype == Int(64) ||
+        dtype == Float(32) ||
+        dtype == Float(64) ||
+        dtype == Bool());
+    if (op->is_scalar() && is_simple_const) {
+      node->pattern = kElemWise;
+    } else {
+      // for now, mark non-scalar constant
+      // as opaque, we will not choose to fuse it.
+      node->pattern = kOpaque;
+    }
+  }
+
+  void VisitExpr_(const CallNode* call) {
+    CHECK(graph_.node_map.count(call));
+    Node* node = graph_.node_map.at(call);
+    static auto fpattern =
+        Op::GetAttr<TOpPattern>("TOpPattern");
+    // setup pattern.
+    OpPatternKind op_pattern = kOpaque;
+    if (const OpNode* opnode = call->op.as<OpNode>()) {
+      op_pattern = static_cast<OpPatternKind>(fpattern[GetRef<Op>(opnode)]);
+    }
+    node->pattern = op_pattern;
+    const auto* rtype = call->checked_type().as<TensorTypeNode>();
+    // pass the message back to all the children it references.
+    for (size_t i = 0; i < call->args.size(); ++i) {
+      const auto* arg_type =
+          call->args[i]->checked_type().as<TensorTypeNode>();
+      // specifically check if result type
+      OpPatternKind edge_pattern = op_pattern;
+      if (edge_pattern == kBroadcast &&
+          arg_type != nullptr &&
+          rtype != nullptr &&
+          attr_equal_(rtype->shape, arg_type->shape)) {
+        edge_pattern = kElemWise;
+      }
+      this->Update(call->args[i], node, edge_pattern);
+    }
+    ExprVisitor::VisitExpr_(call);
+    this->AddNode(call);
+  }
+
+  void VisitExpr_(const TupleNode* op) {
+    for (const Expr& field : op->fields) {
+      this->Update(field, nullptr, kOpaque);
+    }
+    ExprVisitor::VisitExpr_(op);
+    this->AddNode(op);
+  }
+
+  void VisitExpr_(const TupleGetItemNode* op) {
+    CHECK(graph_.node_map.count(op));
+    Node* node = graph_.node_map.at(op);
+    this->Update(op->tuple, node, kOpaque);
+    ExprVisitor::VisitExpr_(op);
+    this->AddNode(op);
+  }
+
+  void VisitExpr_(const VarNode* op) {
+    this->AddNode(op);
+  }
+
+  void VisitExpr_(const LetNode* op) {
+    // do not fuse through let.
+    this->Update(op->var, nullptr, kOpaque);
+    this->Update(op->value, nullptr, kOpaque);
+    this->Update(op->body, nullptr, kOpaque);
+    ExprVisitor::VisitExpr_(op);
+    this->AddNode(op);
+  }
+
+  void VisitExpr_(const IfNode* op) {
+    // do not fuse through if.
+    this->Update(op->cond, nullptr, kOpaque);
+    this->Update(op->true_branch, nullptr, kOpaque);
+    this->Update(op->false_branch, nullptr, kOpaque);
+    ExprVisitor::VisitExpr_(op);
+    this->AddNode(op);
+  }
+};
+
+IndexedForwardGraph IndexedForwardGraph::Create(
+    common::Arena* arena, const Expr& body) {
+  return Creator(arena).Prepare(body);
+}
+
+/*!
+ * \brief Dominator tree that represent domination or
+ *  post domination relation of the node.
+ */
+class DominatorTree {
  public:
+  /*!
+   * \brief A node in the dominator tree.
+   */
+  struct Node {
+    /*! \brief The node in the tree */
+    IndexedForwardGraph::Node* gnode{nullptr};
+    /*! \brief parent of the tree */
+    Node* parent{nullptr};
+    /*! \brief current depth*/
+    int depth{0};
+    /*! \brief aggregated pattern to parent */
+    OpPatternKind pattern{kOpaque};
+  };
+  // index -> node.
+  std::vector<Node*> nodes;
+  /*!
+   * \brief compute a post dominator relation for a given dataflow graph.
+   * \param arena The arena used for node allocation.
+   * \param graph The graph to be analyze.
+   * \return The dominator tree of the graph.
+   * \note This algorithm makes use of the fact that graph is DAG,
+   *       and runs a single pass algorithm via LCA.
+   */
+  static DominatorTree PostDom(common::Arena* arena,
+                               const IndexedForwardGraph& graph);
+
+ private:
+  // Combine pattern together.
+  static OpPatternKind CombinePattern(
+      OpPatternKind lhs, OpPatternKind rhs) {
+    if (lhs > rhs) return lhs;
+    return rhs;
+  }
+  /*!
+   * \brief Find the least common acenstor of the two nodes.
+   * \param lhs The left node.
+   * \param rhs The right node.
+   * \param edge_pattern
+   *        The combined edge pattern across all the parents.
+   * \return The least common acenstor of thw two.
+   */
+  static Node* LeastCommonAcenstor(
+      Node* lhs,
+      Node* rhs,
+      OpPatternKind* edge_pattern) {
+    while (lhs != rhs) {
+      if (lhs == nullptr) return nullptr;
+      if (rhs == nullptr) return nullptr;
+      if (lhs->depth < rhs->depth) {
+        edge_pattern[0] = CombinePattern(
+            edge_pattern[0], rhs->pattern);
+        rhs = rhs->parent;
+      } else if (rhs->depth < lhs->depth) {
+        edge_pattern[0] = CombinePattern(
+            edge_pattern[0], lhs->pattern);
+        lhs = lhs->parent;
+      } else {
+        lhs = lhs->parent;
+        rhs = rhs->parent;
+        edge_pattern[0] = CombinePattern(
+            edge_pattern[0], lhs->pattern);
+        edge_pattern[0] = CombinePattern(
+            edge_pattern[0], rhs->pattern);
+      }
+    }
+    return lhs;
+  }
+};
+
+DominatorTree DominatorTree::PostDom(common::Arena* arena,
+                                     const IndexedForwardGraph& graph) {
+  DominatorTree tree;
+  tree.nodes.resize(graph.post_dfs_order.size(), nullptr);
+  // reverse topo order
+  for (size_t i = graph.post_dfs_order.size(); i != 0; --i) {
+    size_t index = i - 1;
+    Node* tnode = arena->make<Node>();
+    auto* gnode = graph.post_dfs_order[index];
+    tnode->gnode = gnode;
+    if (gnode->extern_ref) {
+      tnode->depth = 1;
+      tnode->parent = nullptr;
+      tnode->pattern = kOpaque;
+    } else {
+      // find the LCAs of all outputs.
+      OpPatternKind pattern = kElemWise;
+      Node* parent = nullptr;
+      for (auto link = gnode->outputs.head; link != nullptr; link= link->next) {
+        size_t oindex = link->value.node->index;
+        CHECK_LT(oindex, tree.nodes.size());
+        Node* onode = tree.nodes[oindex];
+        CHECK(onode != nullptr);
+        if (parent != nullptr) {
+          parent = LeastCommonAcenstor(parent, onode, &pattern);
+        } else {
+          parent = onode;
+        }
+        pattern = CombinePattern(pattern, link->value.pattern);
+      }
+      CHECK(parent != nullptr);
+      tnode->depth = parent->depth + 1;
+      tnode->parent = parent;
+      tnode->pattern = pattern;
+    }
+    tree.nodes[index] = tnode;
+  }
+  return tree;
+}
+
+/*!
+ * \brief A partition of the graph marked by union find data structure.
+ */
+class GraphPartitioner {
+ public:
+  explicit GraphPartitioner(common::Arena* arena, int opt_level)
+      : arena_(arena), opt_level_(opt_level) {}
+  /*!
+   * \brief Group as a union find data structure.
+   */
+  struct Group {
+    /*! \brief The parent in the union find data structure. */
+    Group* parent{nullptr};
+    /*! \brief The pattern of the group */
+    OpPatternKind pattern;
+    /*! \brief reference to the root node. */
+    const tvm::Node* root_ref{nullptr};
+    /*!
+     * \brief Reference to the master node,
+     * this field is not nullptr only if pattern is kOutEWiseFusable.
+     */
+    const tvm::Node* master_ref{nullptr};
+    /*!
+     * \brief Find the group root, perform path compression
+     * \return The root type node.
+     */
+    Group* FindRoot() {
+      // fast path
+      if (this->parent == nullptr) return this;
+      // slow path with path compression.
+      Group* root = this;
+      while (root->parent != nullptr) {
+        root = root->parent;
+      }
+      for (Group* p = this; p != root;) {
+        Group* parent = p->parent;
+        p->parent = root;
+        p = parent;
+      }
+      return root;
+    }
+  };
+  /*!
+   * \brief Partition a graph.
+   * \return group assignments of each node.
+   */
+  std::vector<Group*> Partition(const IndexedForwardGraph& graph);
+
+ private:
+  /*! \brief The internal arena for temporary space. */
+  common::Arena* arena_;
+  /*! \brief optimization level for fuse operation. */
+  int opt_level_;
+  /*! \brief The internal groups. */
+  std::vector<Group*> groups_;
+  /*! \brief internal field used for deduplication */
+  std::unordered_set<IndexedForwardGraph::Node*> visited_;
+  // Internal implelementation of CheckPath
+  template<typename F>
+  bool CheckPath_(IndexedForwardGraph::Node* src,
+                  IndexedForwardGraph::Node* sink,
+                  F fcond) {
+    if (visited_.count(src)) return true;
+    visited_.insert(src);
+    Group* gnode =  groups_[src->index];
+    CHECK(gnode != nullptr);
+    gnode = gnode->FindRoot();
+    if (!fcond(gnode->pattern, src == sink)) return false;
+    if (src == sink) return true;
+    for (auto link = src->outputs.head; link != nullptr; link = link->next) {
+      if (!CheckPath_(link->value.node, sink, fcond)) return false;
+    }
+    return true;
+  }
+  /*!
+   * \brief Check all the node between src and sink satisfies fcond.
+   *
+   * src and sink are not checked.
+   *
+   * \param src The source node.
+   * \param sink The termination node.
+   * \param fcond The condition to be checked.
+   * \tparam F the condition function.
+   * \note sink must be a post-dominator of src.
+   */
+  template<typename F>
+  bool CheckPath(IndexedForwardGraph::Node* src,
+                 IndexedForwardGraph::Node* sink,
+                 F fcond) {
+    CHECK(!src->extern_ref);
+    visited_.clear();
+    CHECK(src != sink);
+    for (auto link = src->outputs.head; link != nullptr; link = link->next) {
+      if (!CheckPath_(link->value.node, sink, fcond)) return false;
+    }
+    return true;
+  }
+  // Combine two patterns together.
+  static OpPatternKind CombinePattern(
+      OpPatternKind lhs, OpPatternKind rhs) {
+    if (lhs > kBroadcast && rhs > kBroadcast) {
+      LOG(FATAL) << "Cannot merge two complex group together";
+    }
+    if (lhs > rhs) return lhs;
+    return rhs;
+  }
+  /*!
+   * \brief Merge the child group to the parent.
+   * \param child The child group.
+   * \param parent The parent group.
+   */
+  void MergeFromTo(Group* child, Group* parent) {
+    child = child->FindRoot();
+    parent = parent->FindRoot();
+    if (child == parent) return;
+    child->parent = parent;
+    // update master ref and pattern
+    if (child->master_ref != nullptr) {
+      CHECK(parent->master_ref == nullptr);
+      parent->master_ref = child->master_ref;
+      parent->pattern = CombinePattern(
+          child->pattern, parent->pattern);
+    }
+  }
+  // Internal implelementation of CommitFuse
+  void CommitFuse_(IndexedForwardGraph::Node* src,
+                   IndexedForwardGraph::Node* sink,
+                   Group* target) {
+    if (src == sink) return;
+    if (visited_.count(src)) return;
+    visited_.insert(src);
+    Group* gnode = groups_[src->index];
+    CHECK(gnode != nullptr);
+    // merge the current group to the parent if possible.
+    MergeFromTo(gnode, target);
+    for (auto link = src->outputs.head; link != nullptr; link = link->next) {
+      CommitFuse_(link->value.node, sink, target);;
+    }
+  }
+  /*!
+   * \brief Commit fusion operation.
+   * \param src The source node.
+   * \param sink The termination node.
+   * \tparam group the group to be committed.
+   * \note sink must be a post-dominator of src.
+   */
+  void CommitFuse(IndexedForwardGraph::Node* src,
+                  IndexedForwardGraph::Node* sink) {
+    Group* target = groups_[sink->index];
+    visited_.clear();
+    CHECK(src != sink);
+    CommitFuse_(src, sink, target);
+  }
+
+  // Initialize the groups.
+  void InitGroups(const IndexedForwardGraph& graph) {
+    groups_.resize(graph.post_dfs_order.size());
+    for (size_t nid = 0; nid < groups_.size(); ++nid) {
+      const auto* graph_node = graph.post_dfs_order[nid];
+      auto* group_node = arena_->make<Group>();
+      group_node->pattern = graph_node->pattern;
+      group_node->root_ref = graph_node->ref;
+      // set master ref if necessary.
+      if (group_node->pattern == kOutEWiseFusable) {
+        group_node->master_ref = graph_node->ref;
+      }
+      groups_[nid] = group_node;
+    }
+  }
+
+  // execute the fusion algorithm.
+  void RunFuse(const IndexedForwardGraph& graph,
+               const DominatorTree& post_dom_tree,
+               int phase) {
+    for (size_t nid = 0; nid < groups_.size(); ++nid) {
+      // the group of current node has been specified already.
+      auto* graph_node = graph.post_dfs_order[nid];
+      auto* dom_node = post_dom_tree.nodes[nid];
+      Group* group_node = groups_[nid];
+      CHECK(group_node != nullptr);
+      // no actions for opaque nodes
+      if (group_node->pattern == kOpaque) continue;
+      // no actions needed if the current node have no dominator
+      if (dom_node->parent == nullptr) continue;
+      CHECK(!graph_node->extern_ref);
+      // Skip if current node is already fused to the parent.
+      size_t dom_parent_gindex = dom_node->parent->gnode->index;
+      if (groups_[dom_parent_gindex] != nullptr &&
+          group_node->FindRoot() == groups_[dom_parent_gindex]->FindRoot()) {
+        continue;
+      }
+      // Try to fuse current node to its post-dominator.
+      if (group_node->pattern == kOutEWiseFusable) {
+        if (phase != 0) continue;
+        // Path for OutEWiseFusable: conv2d
+        // Check if the dominator relation is elemwise.
+        if (dom_node->parent != nullptr && dom_node->pattern == kElemWise) {
+          CHECK(dom_node->parent->gnode != nullptr);
+          // The fuse can be executed if all the intermediate ops are still broadcast.
+          auto fcond = [](OpPatternKind kind, bool is_sink) {
+            return kind <= kBroadcast;
+          };
+          if (CheckPath(graph_node, dom_node->parent->gnode, fcond)) {
+            CommitFuse(graph_node, dom_node->parent->gnode);
+          }
+        }
+      } else if (group_node->pattern <= kBroadcast) {
+        // The fuse can be executed if all the intermediate ops are still broadcast.
+        auto fcond = [](OpPatternKind kind, bool is_sink) {
+          if (!is_sink) {
+            return kind <= kBroadcast;
+          } else {
+            return (kind <= kBroadcast ||
+                    kind == kCommReduce ||
+                    kind == kOutEWiseFusable);
+          }
+        };
+        if (CheckPath(graph_node, dom_node->parent->gnode, fcond)) {
+          CommitFuse(graph_node, dom_node->parent->gnode);
+        }
+      } else if (group_node->pattern == kInjective) {
+        // defer injective fusion to second phase.
+        // so conv2d always finishes fusing.
+        if (phase != 1) continue;
+        // Check if all path are injective.
+        auto fcond = [](OpPatternKind kind, bool is_sink) {
+          return kind <= kInjective;
+        };
+        if (CheckPath(graph_node, dom_node->parent->gnode, fcond)) {
+          CommitFuse(graph_node, dom_node->parent->gnode);
+        }
+      } else {
+        // do nothing.
+        CHECK(group_node->pattern == kCommReduce);
+      }
+    }
+  }
+};
+
+std::vector<GraphPartitioner::Group*>
+GraphPartitioner::Partition(const IndexedForwardGraph& graph) {
+  this->InitGroups(graph);
+  if (opt_level_ == 0) return std::move(groups_);
+  // get post dominator tree
+  auto post_dom_tree = DominatorTree::PostDom(arena_, graph);
+  // run fusion algorithm.
+  for (int phase = 0; phase < 2; ++phase) {
+    this->RunFuse(graph, post_dom_tree, phase);
+  }
+  return std::move(groups_);
+}
+
+class FuseMutator : private ExprMutator {
+ public:
+  // Run the transform
+  Expr Transform(const Expr& body, int fuse_opt_level) {
+    // setup the group map.
+    auto graph = IndexedForwardGraph::Create(&arena_, body);
+    auto groups = GraphPartitioner(&arena_, fuse_opt_level).Partition(
+        graph);
+    for (size_t nid = 0; nid < graph.post_dfs_order.size(); ++nid) {
+      CHECK(graph.post_dfs_order[nid]->ref != nullptr);
+      gmap_[graph.post_dfs_order[nid]->ref] = groups[nid];
+    }
+    // The following line can be used for debug.
+    // this->DebugDumpGroup(body);
+    return this->Mutate(body);
+  }
+
+
+ private:
+  /*! \brief Temporary information from each group. */
+  struct GroupInfo {
+   public:
+    // The parameters of the function.
+    Array<Var> params;
+    // The arguments to call the functions.
+    Array<Expr> arguments;
+    // Get a new parameter or allocate an old one
+    Var GetOrAllocParam(const Expr& expr, const Type& type) {
+      // run linear scan as most fused groups contain only a few inputs.
+      for (size_t i = 0; i < arguments.size(); ++i) {
+        if (expr.same_as(arguments[i])) return params[i];
+      }
+      // create a new parameter.
+      std::ostringstream os;
+      os << "p" << params.size();
+      auto var = VarNode::make(os.str(), type);
+      params.push_back(var);
+      arguments.push_back(expr);
+      return var;
+    }
+  };
+  /*! \brief Internal arena. */
+  common::Arena arena_;
+  /*! \brief The group assignment map. */
+  std::unordered_map<const Node*, GraphPartitioner::Group*> gmap_;
+  /* \brief Internal group information map. */
+  std::unordered_map<GraphPartitioner::Group*, GroupInfo> ginfo_;
   // Skip primitive function.
   Expr VisitExpr_(const FunctionNode* fn_node) {
     NodeRef res = FunctionGetAttr(GetRef<Function>(fn_node), "Primitive");
@@ -26,48 +699,74 @@ class SimpleFuser : public ExprMutator {
       return ExprMutator::VisitExpr_(fn_node);
     }
   }
-
+  // Transform calls.
   Expr VisitExpr_(const CallNode* call) {
     if (call->op.as<OpNode>()) {
-      // Placeholder fusion algorithm which abstracts
-      // single definitions into functions only.
-      Array<Var> params;
-      Array<Expr> inner_args;
-      Array<Expr> args;
-
-      int param_number = 0;
+      // If it is a primitive op call
+      // then we must have a group assignment for it already.
+      CHECK(gmap_.count(call));
+      auto* ret_group = gmap_.at(call)->FindRoot();
+      Array<Expr> new_args;
       for (auto arg : call->args) {
-        std::ostringstream os;
-        os << "p" << param_number++;
         auto type = arg->checked_type();
-        auto var = VarNode::make(os.str(), type);
-        params.push_back(var);
-        inner_args.push_back(var);
-        args.push_back(this->Mutate(arg));
+        CHECK(gmap_.count(arg.get()))
+            << "cannot find group of " << arg;
+        auto* arg_group = gmap_.at(arg.get())->FindRoot();
+        Expr new_arg = this->Mutate(arg);
+
+        if (ret_group != arg_group) {
+          Var param = ginfo_[ret_group].GetOrAllocParam(new_arg, type);
+          new_args.push_back(param);
+        } else {
+          new_args.push_back(new_arg);
+        }
+      }
+      auto new_call = CallNode::make(
+          call->op, new_args, call->attrs, call->type_args);
+
+      if (ret_group->root_ref == call) {
+        // This is the root of the group
+        // create the new call node.
+        const GroupInfo& ginfo = ginfo_[ret_group];
+        auto func = FunctionNode::make(
+            ginfo.params, new_call, call->checked_type(), {});
+        func = FunctionSetAttr(func, "Primitive", tvm::Integer(1));
+        return CallNode::make(func, ginfo.arguments, Attrs());
+      } else {
+        // This is an intermediate node of a fused function
+        // simply return the new call.
+        return new_call;
       }
-      auto body = CallNode::make(call->op, inner_args, call->attrs);
-      auto func = FunctionNode::make(
-          params, body, call->checked_type(), {});
-      func = FunctionSetAttr(func, "Primitive", tvm::Integer(1));
-      return CallNode::make(func, args, Attrs());
     } else {
       return ExprMutator::VisitExpr_(call);
     }
   }
+  // Debug function, dump the group assignment in text.
+  void DebugDumpGroup(const Expr& body) {
+    std::string text = RelayPrint(body, [this](const Expr& expr) -> std::string {
+        auto it = gmap_.find(expr.get());
+        if (it == gmap_.end()) return "";
+        std::ostringstream os;
+        auto *group = it->second->FindRoot();
+        os << "group=" << group;
+        return os.str();
+      });
+    LOG(INFO) << "Dump of group info:\n" << text;
+  }
 };
 
 
-Expr FuseOps(const Expr& expr) {
+Expr FuseOps(const Expr& expr, int fuse_opt_level) {
   // First we convert all chains of fusable ops into
   // abstracted functions which we mark as primtive
   // then we convert these primtive functions into
   // new operators.
-  return SimpleFuser().Mutate(expr);
+  return FuseMutator().Transform(expr, fuse_opt_level);
 }
 
 TVM_REGISTER_API("relay._ir_pass.FuseOps")
 .set_body([](TVMArgs args, TVMRetValue *ret) {
-    *ret = FuseOps(args[0]);
+    *ret = FuseOps(args[0], args[1]);
 });
 }  // namespace relay
 }  // namespace tvm

src/relay/pass/pass_util.h

+/*!
+ *  Copyright (c) 2018 by Contributors.
+ *
+ * \file tvm/relay/pass/pass_util.h
+ * \brief Utilities for writing
+ */
+#ifndef TVM_RELAY_PASS_PASS_UTIL_H_
+#define TVM_RELAY_PASS_PASS_UTIL_H_
+
+#include <tvm/relay/op.h>
+#include <tvm/relay/expr.h>
+#include <tvm/relay/attrs/transform.h>
+
+namespace tvm {
+namespace relay {
+
+/*!
+ * \brief Get reference counter of each internal ExprNode in body.
+ * \param body The body expression.
+ * \return The reference count mapping.
+ */
+std::unordered_map<const Node*, size_t>
+GetExprRefCount(const Expr& body);
+
+}  // namespace relay
+}  // namespace tvm
+#endif  // TVM_RELAY_PASS_PASS_UTIL_H_

src/relay/pass/type_infer.cc

@@ -442,6 +442,9 @@ class TypeInferencer::Resolver : public ExprMutator {
     VarNode* new_var =(
         std::is_base_of<VarNode, T>::value ?
         static_cast<VarNode*>(new_e.node_.get()) : nullptr);
+    FunctionNode* new_fn =(
+        std::is_base_of<FunctionNode, T>::value ?
+        static_cast<FunctionNode*>(new_e.node_.get()) : nullptr);
 
     // check if we need update the new_e
     bool need_update_type = !checked_type.same_as(new_e->checked_type_);
@@ -454,7 +457,17 @@ class TypeInferencer::Resolver : public ExprMutator {
         update_missing_type_annotation_ &&
         !new_var->type_annotation.defined());
 
-    if (!need_update_type && !need_update_var && !need_update_call) return new_e;
+    bool need_update_fn = (
+        std::is_base_of<FunctionNode, T>::value &&
+        update_missing_type_annotation_ &&
+        !new_fn->ret_type.defined());
+
+    if (!need_update_type &&
+        !need_update_var &&
+        !need_update_call &&
+        !need_update_fn) {
+      return new_e;
+    }
 
     if (!new_e.node_.unique()) {
       // Copy on write optimization
@@ -467,6 +480,9 @@ class TypeInferencer::Resolver : public ExprMutator {
       new_var = (
           std::is_base_of<VarNode, T>::value ?
           static_cast<VarNode*>(new_e.node_.get()) : nullptr);
+      new_fn = (
+          std::is_base_of<FunctionNode, T>::value ?
+          static_cast<FunctionNode*>(new_e.node_.get()) : nullptr);
     }
 
     // attach the information.
@@ -483,6 +499,11 @@ class TypeInferencer::Resolver : public ExprMutator {
     if (need_update_var) {
       new_var->type_annotation = checked_type;
     }
+    if (need_update_fn) {
+      auto* fn_type = checked_type.as<FuncTypeNode>();
+      CHECK(fn_type != nullptr);
+      new_fn->ret_type = fn_type->ret_type;
+    }
     return new_e;
   }
 

src/relay/pass/type_solver.cc

@@ -85,18 +85,18 @@ Type TypeSolver::Unify(const Type& dst, const Type& src) {
 void TypeSolver::AddConstraint(const TypeConstraint& constraint) {
   if (auto *op = constraint.as<TypeRelationNode>()) {
     // create a new relation node.
-    RelationNode* rnode = make<RelationNode>();
+    RelationNode* rnode = arena_.make<RelationNode>();
     rnode->rel = GetRef<TypeRelation>(op);
     rel_nodes_.push_back(rnode);
     // populate the type information.
     for (size_t i = 0; i < op->args.size(); ++i) {
       // insert link to the type list
-      LinkNode<TypeNode*>* tlink = make<LinkNode<TypeNode*> >();
+      LinkNode<TypeNode*>* tlink = arena_.make<LinkNode<TypeNode*> >();
       TypeNode* tnode = GetTypeNode(op->args[i]);
       tlink->value = tnode;
       rnode->type_list.Push(tlink);
       // insert type->relation node
-      LinkNode<RelationNode*>* rlink = make<LinkNode<RelationNode*> >();
+      LinkNode<RelationNode*>* rlink = arena_.make<LinkNode<RelationNode*> >();
       rlink->value = rnode;
       tnode->rel_list.Push(rlink);
     }

src/relay/pass/type_solver.h

@@ -16,6 +16,8 @@
 namespace tvm {
 namespace relay {
 
+using common::LinkNode;
+using common::LinkedList;
 /*!
  * \brief Interface of type solver used in type inference.
  *
@@ -70,41 +72,6 @@ class TypeSolver {
   // All the object in the structure is managed by a arena allocator
   // which releases the memory upon distruction of the type solver.
   /*!
-   * \brief Link list node
-   * \tparam T the content data type
-   */
-  template<typename T>
-  struct LinkNode {
-    /*! \brief The content value */
-    T value;
-    /*! \brief pointer to the next location */
-    LinkNode<T>* next{nullptr};
-  };
-  /*!
-   * \brief LinkedList structure
-   * \tparam T the content data type
-   */
-  template<typename T>
-  struct LinkedList {
-    /*! \brief Head pointer */
-    LinkNode<T>* head{nullptr};
-    /*! \brief Tail pointer */
-    LinkNode<T>* tail{nullptr};
-    /*!
-     * \brief Push a new node to the end of the linked list.
-     * \param node The node to be pushed.
-     */
-    void Push(LinkNode<T>* node) {
-      node->next = nullptr;
-      if (this->tail != nullptr) {
-        this->tail->next = node;
-        this->tail = node;
-      } else {
-        head = tail = node;
-      }
-    }
-  };
-  /*!
    * \brief type node struct
    *  TypeNode implements a union-find data structure(via parent)
    *  that can unifies the same types to the name resolved_type.
@@ -165,18 +132,6 @@ class TypeSolver {
   /*! \brief Reporter that reports back to self */
   TypeReporter reporter_;
   /*!
-   * \brief Create function to create a new node ptr via arena
-   * \tparam The type parameter
-   * \return The node pointer.
-   */
-  template<typename T>
-  T* make() {
-    T* ptr = arena_.Alloc<T>();
-    // call constructor
-    new (ptr) T();
-    return ptr;
-  }
-  /*!
    * \brief GetTypeNode that is corresponds to t.
    * if it do not exist, create a new one.
    * \return The type node.
@@ -186,7 +141,7 @@ class TypeSolver {
     if (it != tmap_.end()) {
       return it->second->FindRoot();
     } else {
-      TypeNode* n = make<TypeNode>();
+      TypeNode* n = arena_.make<TypeNode>();
       type_nodes_.push_back(n);
       n->resolved_type = t;
       tmap_[t] = n;

src/relay/pass/util.cc

@@ -129,5 +129,23 @@ TVM_REGISTER_API("relay._ir_pass.free_type_vars")
     }
   });
 
+/*!
+ * \brief Get reference counter of each internal ExprNode in body.
+ * \param body The body expression.
+ * \return The reference count mapping.
+ */
+std::unordered_map<const Node*, size_t>
+GetExprRefCount(const Expr& body) {
+  class ExprRefCounter : private ExprVisitor {
+   public:
+    std::unordered_map<const Node*, size_t>
+    Get(const Expr& body) {
+      this->VisitExpr(body);
+      return std::move(this->visit_counter_);
+    }
+  };
+  return ExprRefCounter().Get(body);
+}
+
 }  // namespace relay
 }  // namespace tvm

tests/python/relay/test_ir_text_printer.py

@@ -33,6 +33,7 @@ def test_env():
     text = env.astext()
     assert "def @myf" in text
     assert "%1 = add(%0, %0) # ty=float32" in text
+    show(env.astext(annotate=lambda x: str(x.checked_type.dtype)))
     show(text)
 
 

tests/python/relay/test_pass_fold_scale_axis.py

@@ -46,6 +46,8 @@ def test_fold_fwd_simple():
         weight = relay.var("weight", type_dict["weight"])
         y1_folded = relay.ir_pass.forward_fold_scale_axis(y1)
         y1_expected = expected(x, weight, in_bias, in_scale, channels)
+        y1_folded = relay.ir_pass.infer_type(y1_folded)
+        y1_expected = relay.ir_pass.infer_type(y1_expected)
         assert relay.ir_pass.alpha_equal(y1_folded, y1_expected)
 
     check((2, 4, 10, 10), 2)
@@ -113,6 +115,8 @@ def test_fold_fwd_dual_path():
         type_dict = {x.name_hint:x.checked_type for x in y1.params}
         weight = relay.var("weight", type_dict["weight"])
         y1_expected = expected(x, weight, in_bias, in_scale, channels)
+        y1_folded = relay.ir_pass.infer_type(y1_folded)
+        y1_expected = relay.ir_pass.infer_type(y1_expected)
         assert relay.ir_pass.alpha_equal(y1_folded, y1_expected)
 
     check((2, 4, 10, 3), 3)
@@ -194,6 +198,8 @@ def test_fold_bwd_simple():
         weight = relay.var("weight", type_dict["weight"])
         y1_folded = relay.ir_pass.backward_fold_scale_axis(y1)
         y1_expected = expected(x, weight, out_bias, out_scale, channels)
+        y1_folded = relay.ir_pass.infer_type(y1_folded)
+        y1_expected = relay.ir_pass.infer_type(y1_expected)
         assert relay.ir_pass.alpha_equal(y1_folded, y1_expected)
 
     check((2, 4, 10, 10), 8)
@@ -255,6 +261,8 @@ def test_fold_bwd_dual_path():
         weight = relay.var("weight", type_dict["weight"])
         y1_folded = relay.ir_pass.backward_fold_scale_axis(y1)
         y1_expected = expected(x, weight, out_bias, out_scale, channels)
+        y1_folded = relay.ir_pass.infer_type(y1_folded)
+        y1_expected = relay.ir_pass.infer_type(y1_expected)
         assert relay.ir_pass.alpha_equal(y1_folded, y1_expected)
 
     check((2, 4, 10, 10), 8)

tests/python/relay/test_pass_fuse_ops.py

@@ -3,15 +3,103 @@ from tvm import relay
 
 def test_fuse_simple():
     """Simple testcase."""
-    x = relay.var("x", shape=(10, 20))
-    y = relay.add(x, x)
-    z = relay.exp(y)
+    def before():
+        x = relay.var("x", shape=(10, 20))
+        y = relay.add(x, relay.const(1, "float32"))
+        z = relay.exp(y)
+        return relay.Function([x], z)
+
+    def expected():
+        x = relay.var("p", shape=(10, 20))
+        y = relay.add(x, relay.const(1, "float32"))
+        z = relay.exp(y)
+        f1 = relay.Function([x], z)
+        x = relay.var("x", shape=(10, 20))
+        y = relay.Call(f1, [x])
+        return relay.Function([x], y)
+
+    z = before()
     z = relay.ir_pass.infer_type(z)
-    zz = relay.ir_pass.fuse_ops(z)
+    zz = relay.ir_pass.fuse_ops(z, opt_level=2)
+    zz = relay.ir_pass.infer_type(zz)
     zz = relay.ir_pass.fuse_ops(zz)
     zz = relay.ir_pass.infer_type(zz)
-    zz.astext()
+    after = relay.ir_pass.infer_type(expected())
+    assert relay.ir_pass.alpha_equal(zz, after)
+
+
+
+def test_conv2d_fuse():
+    """Test fusion case of conv2d"""
+    def before(dshape):
+        x = relay.var("x", shape=dshape)
+        y = relay.nn.conv2d(x, relay.var("w1"),
+                            kernel_size=(3, 3),
+                            padding=(1, 1),
+                            channels=16)
+        # this is the next dominator.
+        y1 = relay.add(relay.const(1, "float32"), y)
+        y = relay.add(y, y1)
+        # second path
+        z2 = relay.nn.conv2d(y, relay.var("w2"),
+                             kernel_size=(1, 1),
+                             padding=(0,0),
+                             channels=16)
+        z3 = relay.nn.conv2d(y, relay.var("w3"),
+                             kernel_size=(3, 3),
+                             padding=(1,1),
+                             channels=16)
+        # add can only be fused to z1
+        z = relay.add(z2, z3)
+        return relay.Function(relay.ir_pass.free_vars(z), z)
+
+    def expected(dshape):
+        # segment 1
+        x = relay.var("p0", shape=dshape)
+        w = relay.var("p1")
+        y = relay.nn.conv2d(x, w,
+                            kernel_size=(3, 3),
+                            padding=(1, 1),
+                            channels=16)
+        y1 = relay.add(relay.const(1, "float32"), y)
+        y = relay.add(y, y1)
+        f1 = relay.Function([x, w], y)
+        # segment 2
+        x = relay.var("p0", shape=dshape)
+        w = relay.var("p1")
+        z2 = relay.nn.conv2d(x, w,
+                             kernel_size=(3, 3),
+                             padding=(1,1),
+                             channels=16)
+        f2 = relay.Function([x, w], z2)
+        # segment 3
+        x = relay.var("p0", shape=dshape)
+        w = relay.var("p1")
+        offset = relay.var("p2", shape=dshape)
+        z3 = relay.nn.conv2d(x, w,
+                             kernel_size=(1, 1),
+                             padding=(0, 0),
+                             channels=16)
+        z3 = relay.add(z3, offset)
+        f3 = relay.Function([x, w, offset], z3)
+        # compose
+        x = relay.var("x", shape=dshape)
+        y = relay.Call(f1, [x, relay.var("w1")])
+        z2 = relay.Call(f2, [y, relay.var("w3")])
+        z3 = relay.Call(f3, [y, relay.var("w2"), z2])
+        z = z3
+        return relay.Function(relay.ir_pass.free_vars(z), z)
+
+    dshape = (1, 16, 64, 64)
+    z = before(dshape)
+    z = relay.ir_pass.infer_type(z)
+    zz = relay.ir_pass.fuse_ops(z, opt_level=2)
+    zz = relay.ir_pass.infer_type(zz)
+    after = relay.ir_pass.infer_type(expected(dshape))
+    assert relay.ir_pass.alpha_equal(zz, after)
+
 
 
 if __name__ == "__main__":
     test_fuse_simple()
+    test_conv2d_fuse()