implement

include/tvm/relay/pass.h

@@ -188,7 +188,6 @@ Expr ForwardRewrite(const Expr& expr,
                     std::function<NodeRef(const Call&)> fcontext = nullptr,
                     std::function<Expr(const Expr&)> fmulti_ref_trigger = nullptr);
 
-
 /*! \brief A hashing structure in the style of std::hash. */
 struct StructuralHash {
   /*! \brief Hash a Relay type.
@@ -212,6 +211,29 @@ struct StructuralHash {
   size_t operator()(const Expr& expr) const;
 };
 
+/*! \brief turn a dataflow graph into A Normal Form.
+ *
+ * It will turn an expression that is in a graph form (with sharing implicit),
+ * to an expression with explicit sharing (A Normal Form).
+ *
+ * \param e the expression to observably share
+ * \param mod The module used for referencing global functions, can be
+ * None.
+ *
+ * \return expression in A Normal Form
+ */
+Expr ToANF(const Expr& e, const Module& mod);
+
+inline bool IsPrimitiveFunction(const Function& fn) {
+  NodeRef res = FunctionGetAttr(fn, "Primitive");
+  const ir::IntImm* pval = res.as<ir::IntImm>();
+  return pval && (pval->value != 0);
+}
+
+inline bool IsPrimitiveFunction(const Expr& e) {
+  return e.as<FunctionNode>() && IsPrimitiveFunction(Downcast<Function>(e));
+}
+
 }  // namespace relay
 }  // namespace tvm
 

python/tvm/relay/ir_pass.py

@@ -357,3 +357,20 @@ def alter_op_layout(expr):
         Transformed expression with alternated layout.
     """
     return _ir_pass.AlterOpLayout(expr)
+
+
+def to_anf(expr, mod=None):
+    """
+    Turn Graph Normal Form expression into A Normal Form Expression
+
+    Parameters
+    ----------
+    expr : tvm.relay.Expr
+        The input expression
+
+    Returns
+    -------
+    expr: tvm.relay.Expr
+      The output expression
+    """
+    return _ir_pass.to_anf(expr, mod)

src/relay/pass/let_list.h

@@ -36,6 +36,7 @@ class LetList {
    * \return a Var that hold the inserted expr.
    */
   Var Push(Var pv, Expr expr) {
+    CHECK(!used);
     lets_.emplace_back(std::make_pair(pv, expr));
     return pv;
   }
@@ -71,11 +72,13 @@ class LetList {
    *
    *  \return the wrapped expr.
    */
-  Expr Get(const Expr& body) const {
+  Expr Get(const Expr& body) {
+    CHECK(!used);
     Expr ret = body;
     for (auto rit = lets_.rbegin(); rit != lets_.rend(); ++rit) {
       ret = LetNode::make(std::get<0>(*rit), std::get<1>(*rit), ret);
     }
+    used = true;
     return ret;
   }
 
@@ -106,8 +109,13 @@ class LetList {
     return ll.Get(f(&ll));
   }
 
+  ~LetList() {
+    CHECK(used || (lets_.size() == 0));
+  }
+
  private:
   std::vector<std::pair<Var, Expr> > lets_;
+  bool used = false;
 };
 
 }  // namespace relay