[ARITH] DeduceBound

include/tvm/ir_visitor.h

@@ -46,6 +46,12 @@ class IRVisitor {
   virtual void Visit_(const Let* op);
   virtual void Visit_(const Free* op);
   virtual void Visit_(const Call* op);
+  virtual void Visit_(const Add* op);
+  virtual void Visit_(const Sub* op);
+  virtual void Visit_(const Mul* op);
+  virtual void Visit_(const Div* op);
+  virtual void Visit_(const Mod* op);
+  virtual void Visit_(const LT* op);
 };
 
 /*!

src/arithmetic/bound_deducer.cc

@@ -0,0 +1,133 @@
+/*!
+ *  Copyright (c) 2017 by Contributors
+ * \file bound_deducer.cc
+ */
+#include <tvm/expr.h>
+#include <tvm/ir_pass.h>
+#include <tvm/ir_visitor.h>
+#include <tvm/api_registry.h>
+#include <unordered_set>
+#include "./int_set.h"
+
+namespace tvm {
+namespace arith {
+
+using namespace ir;
+using Halide::Internal::Interval;
+
+// a visitor to find the path to the target variable
+// from a expression.
+class VariableFinder: public IRVisitor {
+ public:
+  explicit VariableFinder(Var target) : target_(target) {}
+
+  void Visit(const NodeRef& node) final {
+    if (finded_) return;
+    if (visited_.count(node.get()) != 0) return;
+    visited_.insert(node.get());
+
+    path_.push_back(node.get());
+    if (node.same_as(target_)) finded_ = true;
+    IRVisitor::Visit(node);
+    if (!finded_) path_.pop_back();
+  }
+
+  std::vector<const Node*> path_;
+
+ private:
+  bool finded_{false};
+  Var target_;
+  std::unordered_set<const Node*> visited_;
+};
+
+
+// get the path to the variable
+std::vector<const Node*> GetPath(Var target, Expr expr) {
+  VariableFinder v(target);
+  v.Visit(expr);
+  return v.path_;
+}
+
+
+// a visitor to deduce the bound of a variable from a expression
+class BoundDeducer: public IRVisitor {
+ public:
+  Expr Deduce(Var target, Expr expr) {
+    path_ = GetPath(target, expr);
+    target_ = target;
+    iter_ = 0;
+    result = make_zero(expr.type());
+
+    Visit(expr);
+    return result;
+  }
+
+  void Visit(const NodeRef& e) final {
+    if (e.get() == path_[iter_++]) {
+      IRVisitor::Visit(e);
+    } else {
+      LOG(FATAL) << "the current node is not match with the deduced path";
+    }
+  }
+
+  void Visit_(const Add* op) final {
+    bool left = op->a.get() == path_[iter_];
+    result -= left ? op->b : op->a;
+    Visit(left ? op->a : op->b);
+  }
+
+  void Visit_(const Sub* op) final {
+    bool left = op->a.get() == path_[iter_];
+    if (left) {
+      result += op->b;
+    } else {
+      result -= op->a;
+      result = -1 * result;
+      is_greater = !is_greater;
+    }
+    Visit(left ? op->a : op->b);
+  }
+
+  void Visit_(const Mul* op) final {
+    bool left = op->a.get() == path_[iter_];
+    Expr operand = left ? op->b : op->a;
+    if (is_negative_const(operand)) is_greater = !is_greater;
+    result /= operand;
+    Visit(left ? op->a : op->b);
+  }
+
+  void Visit_(const Div* op) final {
+    bool left = op->a.get() == path_[iter_];
+    Expr operand = left ? op->b : op->a;
+    if (is_negative_const(operand)) is_greater = !is_greater;
+    result = left ? result * operand : operand / result;
+    Visit(left ? op->a : op->b);
+  }
+
+  Expr result;
+  bool is_greater{true};
+
+ private:
+  Var  target_;
+  std::vector<const Node*> path_;
+  size_t iter_;
+};
+
+// Assuming e >= 0, deduce the bound of variable from it.
+IntSet DeduceBound(Var v, Expr e) {
+    BoundDeducer deducer;
+    deducer.Deduce(v, e);
+    Type t = deducer.result.type();
+    return deducer.is_greater ?
+      IntSet::range(Range(deducer.result, Cast::make(t, Interval::pos_inf))) :
+      IntSet::range(Range(Cast::make(t, Interval::neg_inf), deducer.result));
+}
+
+TVM_REGISTER_API(_pass_DeduceBound)
+.set_body([](TVMArgs args, TVMRetValue *ret) {
+    *ret = DeduceBound(args[0].operator Var(), args[1].operator Expr());
+  });
+
+
+} // namespace arith
+} // namespace tvm

src/pass/ir_visitor.cc

@@ -67,7 +67,13 @@ TVM_STATIC_IR_FUNCTOR(IRVisitor, vtable)
 .DISPATCH_TO_VISIT(Store)
 .DISPATCH_TO_VISIT(Let)
 .DISPATCH_TO_VISIT(Call)
-.DISPATCH_TO_VISIT(Free);
+.DISPATCH_TO_VISIT(Free)
+.DISPATCH_TO_VISIT(Add)
+.DISPATCH_TO_VISIT(Sub)
+.DISPATCH_TO_VISIT(Mul)
+.DISPATCH_TO_VISIT(Div)
+.DISPATCH_TO_VISIT(Mod)
+.DISPATCH_TO_VISIT(LT);
 
 void IRVisitor::Visit_(const Variable* op) {}
 
@@ -128,6 +134,36 @@ void IRVisitor::Visit_(const Call *op) {
   VisitArray(op->args, this);
 }
 
+void IRVisitor::Visit_(const Add* op) {
+  this->Visit(op->a);
+  this->Visit(op->b);
+}
+
+void IRVisitor::Visit_(const Sub* op) {
+  this->Visit(op->a);
+  this->Visit(op->b);
+}
+
+void IRVisitor::Visit_(const Mul* op) {
+  this->Visit(op->a);
+  this->Visit(op->b);
+}
+
+void IRVisitor::Visit_(const Div* op) {
+  this->Visit(op->a);
+  this->Visit(op->b);
+}
+
+void IRVisitor::Visit_(const Mod* op) {
+  this->Visit(op->a);
+  this->Visit(op->b);
+}
+
+void IRVisitor::Visit_(const LT* op) {
+  this->Visit(op->a);
+  this->Visit(op->b);
+}
+
 TVM_STATIC_IR_FUNCTOR(IRVisitor, vtable)
 .set_dispatch<Reduce>([](const Reduce* op, IRVisitor* v) {
     VisitRDom(op->axis, v);
@@ -151,16 +187,16 @@ inline void Binary(const T* op, IRVisitor* v) {
 }
 
 TVM_STATIC_IR_FUNCTOR(IRVisitor, vtable)
-.set_dispatch<Add>(Binary<Add>)
-.set_dispatch<Sub>(Binary<Sub>)
-.set_dispatch<Mul>(Binary<Mul>)
-.set_dispatch<Div>(Binary<Div>)
-.set_dispatch<Mod>(Binary<Mod>)
+// .set_dispatch<Add>(Binary<Add>)
+// .set_dispatch<Sub>(Binary<Sub>)
+// .set_dispatch<Mul>(Binary<Mul>)
+// .set_dispatch<Div>(Binary<Div>)
+// .set_dispatch<Mod>(Binary<Mod>)
 .set_dispatch<Min>(Binary<Min>)
 .set_dispatch<Max>(Binary<Max>)
 .set_dispatch<EQ>(Binary<EQ>)
 .set_dispatch<NE>(Binary<NE>)
-.set_dispatch<LT>(Binary<LT>)
+// .set_dispatch<LT>(Binary<LT>)
 .set_dispatch<LE>(Binary<LE>)
 .set_dispatch<GT>(Binary<GT>)
 .set_dispatch<GE>(Binary<GE>)

tests/python/unittest/test_pass_deduce.py

@@ -0,0 +1,12 @@
+import tvm
+
+x = tvm.Var('x')
+y = tvm.Var('y')
+z = tvm.Var('z')
+a = tvm.Var('a')
+b = tvm.Var('b')
+
+e0 = (-a+x*y+z-b)
+print(type(e0))
+e1 = tvm.ir_pass.DeduceBound(x, e0)
+print(e1)