[PASS] Canonical form simplify (#34)

include/tvm/ir_pass.h

@@ -62,6 +62,13 @@ bool HasSideEffect(const Expr& e);
  */
 Stmt ConvertSSA(Stmt stmt);
 
+/*!
+ * \brief Simplify by applying canonical form.
+ * \param stmt The statement to be canonically simplifed.
+ * \return Canonicalized statement.
+ */
+Stmt CanonicalSimplify(Stmt stmt);
+
 /*!
  * \brief Substitute the var specified in key->var to be value.
  * \param stmt The source statement to be substituted

python/tvm/build.py

@@ -17,7 +17,8 @@ def build(sch,
           target,
           name="default_function",
           binds=None,
-          record_codes=None):
+          record_codes=None,
+          max_auto_unroll_step=8):
     """Build a function with arguments as signiture.
 
     Parameters
@@ -38,6 +39,9 @@ def build(sch,
         Dictionary that maps the binding of symbolic buffer to Tensor.
         By default, a new buffer is created for each tensor in the argument.
 
+    max_auto_unroll_step: int
+        Maximum step to perform automatic unrolling
+
     Returns
     -------
     f : Function, or pair of functions
@@ -64,6 +68,8 @@ def build(sch,
     bounds = schedule.InferBound(sch)
     stmt = schedule.ScheduleOps(sch, bounds)
     stmt = ir_pass.StorageFlatten(stmt, binds)
+    stmt = ir_pass.CanonicalSimplify(stmt)
+    stmt = ir_pass.UnrollLoop(stmt, max_auto_unroll_step)
     stmt = ir_pass.Simplify(stmt)
     fapi = ir_pass.MakeAPI(stmt, name, arg_list, len(arg_list))
     fsplits = ir_pass.SplitHostDevice(fapi)

src/api/api_pass.cc

@@ -59,6 +59,7 @@ TVM_REGISTER_API(_pass_PostOrderVisit)
 
 REGISTER_PASS1(ConvertSSA);
 REGISTER_PASS1(VerifySSA);
+REGISTER_PASS1(CanonicalSimplify);
 REGISTER_PASS4(Inline);
 REGISTER_PASS2(StorageFlatten);
 REGISTER_PASS2(UnrollLoop);