[PASS] Export simplify and equal to python

Author: icemelon

include/tvm/ir_pass.h

@@ -9,6 +9,8 @@
 #ifndef TVM_IR_PASS_H_
 #define TVM_IR_PASS_H_
 
+#include <ir/IREquality.h>
+#include <pass/Simplify.h>
 #include <tvm/ir_functor.h>
 #include <unordered_map>
 #include <vector>
@@ -19,6 +21,8 @@
 namespace tvm {
 namespace ir {
 
+using Halide::Internal::equal;
+using Halide::Internal::simplify;
 
 /*!
  * \brief Schedule s' dependent operations.

src/c_api/c_api_pass.cc

@@ -13,6 +13,27 @@ namespace ir {
 using ArgStack = const std::vector<APIVariantValue>;
 using RetValue = APIVariantValue;
 
+TVM_REGISTER_API(_pass_Simplify)
+.set_body([](const ArgStack& args, RetValue *ret) {
+    CHECK(args.at(0).type_id == kNodeHandle);
+    if (dynamic_cast<Expr::ContainerType*>(args.at(0).sptr.get())) {
+      *ret = simplify(args.at(0).operator Expr());
+    } else {
+      *ret = simplify(args.at(0).operator Stmt());
+    }
+  });
+
+TVM_REGISTER_API(_pass_equal)
+.set_body([](const ArgStack& args, RetValue *ret) {
+    CHECK(args.at(0).type_id == kNodeHandle);
+    CHECK(args.at(1).type_id == kNodeHandle);
+    if (dynamic_cast<Expr::ContainerType*>(args.at(0).sptr.get())) {
+      *ret = equal(args.at(0).operator Expr(), args.at(1).operator Expr());
+    } else {
+      *ret = equal(args.at(0).operator Stmt(), args.at(1).operator Stmt());
+    }
+  });
+
 // make from two arguments
 #define REGISTER_PASS1(PassName)                                  \
   TVM_REGISTER_API(_pass_## PassName)                             \

tests/python/test_pass_basic.py

@@ -1,5 +1,15 @@
 import tvm
 
+def test_simplify():
+  x = tvm.Var('x')
+  e1 = tvm.ir_pass.Simplify(x + 2 + 1)
+  assert(tvm.ir_pass.equal(e1, x + 3))
+  e2 = tvm.ir_pass.Simplify(x * 3 + 5 * x)
+  assert(tvm.ir_pass.equal(e2, x * 8))
+  e3 = tvm.ir_pass.Simplify(x - x / 3 * 3)
+  assert(tvm.ir_pass.equal(e3, tvm.make.Mod(x, 3)))
+
+
 def test_verify_ssa():
     x = tvm.Var('x')
     y = tvm.Var()