[Arith] add SizeVar representing non-neg valued variable in a tensor shape (apache#4684)

* [arith] add ShapeVar representing non-neg valued variable in a tensor shape

* bounder remover; deal with div in int_set differently

* fix bounder_remover

* migrate unittest to use shape_var

* use tvm.shape_var in integration & relay tests

* add test case; fix Var register

* fix lint

* fix lint again

* add default ShapeVar visitor in Relay

* fix override

* fix ShapeVar visit bug

* revert IntervalSet for shape_var

* remove bound_remover

* remove is_var; use constructor for shapevar/var instead

* ShapeVar -> SizeVar; add constructor comments

* shape_var -> size_var in doc

* tindex -> size

Author: yzhliu

docs/api/python/tvm.rst

@@ -24,6 +24,7 @@ The user facing API for computation declaration.
    tvm.load_json
    tvm.save_json
    tvm.var
+   tvm.size_var
    tvm.const
    tvm.convert
    tvm.placeholder
@@ -49,6 +50,7 @@ The user facing API for computation declaration.
 .. autofunction:: tvm.load_json
 .. autofunction:: tvm.save_json
 .. autofunction:: tvm.var
+.. autofunction:: tvm.size_var
 .. autofunction:: tvm.const
 .. autofunction:: tvm.convert
 .. autofunction:: tvm.placeholder

include/tvm/expr.h

@@ -65,27 +65,33 @@ class Var;
  */
 class VarNode : public PrimExprNode {
  public:
+  /*! \brief constructor */
+  VarNode() {}
+  VarNode(DataType dtype, std::string name_hint);
+
   /*!
    * \brief The hint to the variable name.
    * \note Each variable is uniquely identified by its address.
    */
   std::string name_hint;
 
-  static Var make(DataType dtype, std::string name_hint);
-
   void VisitAttrs(AttrVisitor* v) {
     v->Visit("dtype", &dtype);
     v->Visit("name", &name_hint);
   }
 
   static constexpr const char* _type_key = "Variable";
-  TVM_DECLARE_FINAL_OBJECT_INFO(VarNode, PrimExprNode);
+  TVM_DECLARE_BASE_OBJECT_INFO(VarNode, PrimExprNode);
 };
 
 /*! \brief a named variable in TVM */
 class Var : public PrimExpr {
  public:
   explicit Var(ObjectPtr<Object> n) : PrimExpr(n) {}
+  /*! \brief constructor
+   * \param name_hint variable name
+   * \param t data type
+   */
   TVM_DLL explicit Var(std::string name_hint = "v",
                        DataType t = DataType::Int(32));
   /*!
@@ -114,6 +120,53 @@ class Var : public PrimExpr {
   using ContainerType = VarNode;
 };
 
+class SizeVar;
+/*!
+ * \brief A variable node represent a tensor index size,
+ * whose value must be non-negative.
+ */
+class SizeVarNode : public VarNode {
+ public:
+  /*! \brief constructor */
+  SizeVarNode() {}
+  /*! \brief constructor
+   * \param dtype data type
+   * \param name_hint variable name
+   */
+  SizeVarNode(DataType dtype, std::string name_hint);
+
+  static constexpr const char* _type_key = "SizeVar";
+  TVM_DECLARE_FINAL_OBJECT_INFO(SizeVarNode, VarNode);
+};
+
+/*! \brief a named variable represents a tensor index size */
+class SizeVar : public Var {
+ public:
+  explicit SizeVar(ObjectPtr<Object> n) : Var(n) {}
+  /*! \brief constructor
+   * \param name_hint variable name
+   * \param t data type
+   */
+  TVM_DLL explicit SizeVar(std::string name_hint = "s",
+                            DataType t = DataType::Int(32));
+  /*!
+   * \brief Get pointer to the internal value.
+   * \return the corresponding Variable.
+   */
+  const SizeVarNode* operator->() const {
+    return get();
+  }
+  /*!
+   * \brief Get pointer to the internal value.
+   * \return the corresponding Variable.
+   */
+  const SizeVarNode* get() const {
+    return static_cast<const SizeVarNode*>(data_.get());
+  }
+  /*! \brief type indicate the container type */
+  using ContainerType = SizeVarNode;
+};
+
 /*!
  * \brief Container of constant int that adds more constructors.
  *

include/tvm/ir.h

@@ -38,6 +38,7 @@ namespace ir {
 using IntImmNode = tvm::IntImmNode;
 using FloatImmNode = tvm::FloatImmNode;
 using VarNode = tvm::VarNode;
+using SizeVarNode = tvm::SizeVarNode;
 
 /*! \brief String constants, only used in asserts. */
 class StringImmNode : public PrimExprNode {
@@ -679,7 +680,7 @@ class AnyNode : public PrimExprNode {
   void VisitAttrs(AttrVisitor* v) {}
   /*! \brief Convert to var. */
   Var ToVar() const {
-    return VarNode::make(DataType::Int(32), "any_dim");
+    return Var("any_dim", DataType::Int(32));
   }
 
   TVM_DLL static PrimExpr make();

include/tvm/ir_functor_ext.h

@@ -133,6 +133,9 @@ class ExprFunctor<R(const PrimExpr& n, Args...)> {
   }
   // Functions that can be overriden by subclass
   virtual R VisitExpr_(const VarNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
+  virtual R VisitExpr_(const SizeVarNode* op, Args... args) {
+    return VisitExpr_(static_cast<const VarNode*>(op), std::forward<Args>(args)...);
+  }
   virtual R VisitExpr_(const LoadNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
   virtual R VisitExpr_(const LetNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
   virtual R VisitExpr_(const CallNode* op, Args... args) EXPR_FUNCTOR_DEFAULT;
@@ -174,6 +177,7 @@ class ExprFunctor<R(const PrimExpr& n, Args...)> {
     FType vtable;
     // Set dispatch
     IR_EXPR_FUNCTOR_DISPATCH(VarNode);
+    IR_EXPR_FUNCTOR_DISPATCH(SizeVarNode);
     IR_EXPR_FUNCTOR_DISPATCH(LoadNode);
     IR_EXPR_FUNCTOR_DISPATCH(LetNode);
     IR_EXPR_FUNCTOR_DISPATCH(CallNode);
@@ -297,6 +301,7 @@ class TVM_DLL ExprVisitor :
   using ExprFunctor::VisitExpr;
   // list of functions to override.
   void VisitExpr_(const VarNode* op) override;
+  void VisitExpr_(const SizeVarNode* op) override;
   void VisitExpr_(const LoadNode* op) override;
   void VisitExpr_(const LetNode* op) override;
   void VisitExpr_(const CallNode* op) override;
@@ -341,6 +346,7 @@ class TVM_DLL ExprMutator :
   using ExprFunctor::VisitExpr;
   // list of functions to override.
   PrimExpr VisitExpr_(const VarNode* op) override;
+  PrimExpr VisitExpr_(const SizeVarNode* op) override;
   PrimExpr VisitExpr_(const LoadNode* op) override;
   PrimExpr VisitExpr_(const LetNode* op) override;
   PrimExpr VisitExpr_(const CallNode* op) override;

python/tvm/api.py

@@ -192,6 +192,25 @@ def var(name="tindex", dtype=int32):
     return _api_internal._Var(name, dtype)
 
 
+def size_var(name="size", dtype=int32):
+    """Create a new variable represents a tensor shape size, which is non-negative.
+
+    Parameters
+    ----------
+    name : str
+        The name
+
+    dtype : str
+        The data type
+
+    Returns
+    -------
+    var : SizeVar
+        The result symbolic shape variable.
+    """
+    return _api_internal._SizeVar(name, dtype)
+
+
 def any(*args):
     """Create a new experssion of the union of all conditions in the arguments
 

python/tvm/expr.py

@@ -278,6 +278,25 @@ def __init__(self, name, dtype):
             _api_internal._Var, name, dtype)
 
 
+@register_object
+class SizeVar(Var):
+    """Symbolic variable to represent a tensor index size
+       which is greater or equal to zero
+
+    Parameters
+    ----------
+    name : str
+        The name
+
+    dtype : int
+        The data type
+    """
+    # pylint: disable=super-init-not-called
+    def __init__(self, name, dtype):
+        self.__init_handle_by_constructor__(
+            _api_internal._SizeVar, name, dtype)
+
+
 @register_object
 class Reduce(PrimExpr):
     """Reduce node.

python/tvm/hybrid/preprocessor.py

@@ -63,7 +63,7 @@ def visit_Call(self, node):
         _internal_assert(func_id in list(HYBRID_GLOBALS.keys()) + \
                          ['range', 'max', 'min', 'len'] + \
                          list(self.symbols.keys()), \
-                         "Function call id not in intrinsics' list")
+                         "Function call id " + func_id + " not in intrinsics' list")
         for elem in node.args:
             self.visit(elem)
 

src/api/api_ir.cc

@@ -33,7 +33,12 @@ namespace ir {
 
 TVM_REGISTER_GLOBAL("_Var")
 .set_body_typed([](std::string s, DataType t) {
-    return VarNode::make(t, s);
+    return Var(s, t);
+  });
+
+TVM_REGISTER_GLOBAL("_SizeVar")
+.set_body_typed([](std::string s, DataType t) {
+    return SizeVar(s, t);
   });
 
 TVM_REGISTER_GLOBAL("make.abs")

src/arithmetic/bound_deducer.cc

@@ -86,7 +86,7 @@ class BoundDeducer: public ExprVisitor {
 
   void VisitExpr(const PrimExpr& e) final {
     if (!success_) return;
-    if (e.get() == path_[iter_++]) {
+    if (iter_ < path_.size() && e.get() == path_[iter_++]) {
       ExprVisitor::VisitExpr(e);
     } else {
       success_ = false;
@@ -297,6 +297,7 @@ void BoundDeducer::Transform() {
 void BoundDeducer::Deduce() {
   Init();
   if (!success_) return;
+
   Relax();
   if (!success_) return;
   // get the path

src/arithmetic/const_int_bound.cc

@@ -284,6 +284,16 @@ class ConstIntBoundAnalyzer::Impl :
     }
   }
 
+  Entry VisitExpr_(const SizeVarNode* op) final {
+    SizeVar v = GetRef<SizeVar>(op);
+    auto it = var_map_.find(v);
+    if (it != var_map_.end()) {
+      return it->second;
+    } else {
+      return MakeBound(0, kPosInf);
+    }
+  }
+
   Entry VisitRightShift(const CallNode* op) {
     Entry a = VisitExpr(op->args[0]);
     Entry b = VisitExpr(op->args[1]);