[Relay][RFC] Implement type checking for Any (apache#3221)

* Implement type checking for Any

Remove code generation related changes

Remove compile changes

Remove more

Remove unification hack

Add some code back that was needed, and clean up test

Refactor test cases

WIP

Implement TypeHint AST

Add test case which should fail

Remove unification changes, and fix bug with let rec

Restore unification for shapes

Improve error reporting while debugging

All examples type check

All examples type check

WIP

First version that works with hints, needs clean up

Remove dead code

Tweaks

Remove type hint

Remove unecessary type hint stuff

Remove more type hints

Clean up

Expose Any expression node

Address CR

Fix

Fix solver

Kill unecessary code

Fix

PyLint

Fix

Relocate loops

Fix license and test

Lint again

Lint again

Fix loops

Fix docstring

Fix template error

Fix compiler issue

Fix compile err

Remove more runtime changes

Restore buffer

Fix segfault

Fix

Fix arange

* Address feedback

* Fix typo

* Fix arange

* Fix op level3

* Fix issue with Python wrapper

include/tvm/ir.h

@@ -138,6 +138,15 @@ struct Reduce : public ExprNode<Reduce> {
   static constexpr const char* _type_key = "Reduce";
 };
 
+/*! \brief Any shape. */
+struct Any : public ExprNode<Any> {
+  TVM_DLL static Expr make();
+
+  void VisitAttrs(AttrVisitor* v) final {}
+  static const IRNodeType _type_info = IRNodeType::ExtensionExpr;
+  static constexpr const char* _type_key = "Any";
+};
+
 /*!
  * \brief Auxiliary data structure used in IR Pass to indicate a tensor.
  */

include/tvm/relay/attrs/transform.h

@@ -123,19 +123,19 @@ struct InitOpAttrs : public tvm::AttrsNode<InitOpAttrs> {
 
 /*! \brief Attributes used in arange operators */
 struct ArangeAttrs : public tvm::AttrsNode<ArangeAttrs> {
-  tvm::Expr start;
-  tvm::Expr stop;
-  tvm::Expr step;
+  Expr start;
+  Expr stop;
+  Expr step;
   DataType dtype;
 
   TVM_DECLARE_ATTRS(ArangeAttrs, "relay.attrs.ArangeAttrs") {
-    TVM_ATTR_FIELD(start).set_default(make_const(Float(32), 0))
+    TVM_ATTR_FIELD(start)
         .describe("Start of interval. The interval includes this value.");
     TVM_ATTR_FIELD(stop)
         .describe("Stop of interval. The interval does not include this value.");
-    TVM_ATTR_FIELD(step).set_default(make_const(Float(32), 1))
+    TVM_ATTR_FIELD(step)
         .describe("Spacing between values.");
-    TVM_ATTR_FIELD(dtype).set_default(NullValue<DataType>())
+    TVM_ATTR_FIELD(dtype)
         .describe("Target data type.");
   }
 };  // struct ArangeAttrs

include/tvm/relay/error.h

@@ -64,9 +64,10 @@ struct RelayErrorStream {
 
 struct Error : public dmlc::Error {
   Span sp;
-  explicit Error(const std::string& msg) : dmlc::Error(msg), sp() {}
-  Error(const std::stringstream& msg) : dmlc::Error(msg.str()), sp() {} // NOLINT(*)
-  Error(const RelayErrorStream& msg) : dmlc::Error(msg.str()), sp() {} // NOLINT(*)
+  explicit Error(const std::string& msg) : dmlc::Error(msg), sp(nullptr) {}
+  Error(const RelayErrorStream& msg) : dmlc::Error(msg.str()), sp(nullptr) {} // NOLINT(*)
+  Error(const Error& err) : dmlc::Error(err.what()), sp(nullptr) {}
+  Error() : dmlc::Error(""), sp(nullptr) {}
 };
 
 /*! \brief An abstraction around how errors are stored and reported.
@@ -118,7 +119,8 @@ class ErrorReporter {
    * \param err The error message to report.
    */
   inline void ReportAt(const GlobalVar& global, const NodeRef& node, std::stringstream& err) {
-    this->ReportAt(global, node, Error(err));
+    std::string err_msg = err.str();
+    this->ReportAt(global, node, Error(err_msg));
   }
 
   /*! \brief Report an error against a program, using the full program

include/tvm/relay/expr.h

@@ -561,6 +561,9 @@ inline const TTypeNode* ExprNode::type_as() const {
   return node;
 }
 
+/*! \brief Pretty print a Relay node, producing a fragment of the Relay text format. */
+std::string PrettyPrint(const NodeRef& node);
+
 /*!
  * \brief Render the node as a string in the Relay text format.
  * \param node The node to be rendered.

include/tvm/relay/op_attr_types.h

@@ -158,6 +158,22 @@ using FForwardRewrite = runtime::TypedPackedFunc<
 using FPrimalGradient = runtime::TypedPackedFunc<tvm::Array<Expr>(const Expr& orig_call,
                                                                   const Expr& output_grad)>;
 
+/*!
+ * \brief The codegeneration strategy for dynamic dimensions.
+ */
+enum AnyCodegenStrategy {
+  /*! \brief The default strategy of using completely variable dimensions. */
+  kVariableDimensions
+};
+
+/* \brief A runtime representation of shape. */
+using Shape = Array<IndexExpr>;
+
+using FShapeFunc = runtime::TypedPackedFunc<
+  Array<Tensor>(const Attrs& attrs,
+                const Array<Tensor>& inputs,
+                const Array<Shape>& out_shapes)>;
+
 }  // namespace relay
 }  // namespace tvm
 #endif  // TVM_RELAY_OP_ATTR_TYPES_H_

include/tvm/relay/type.h

@@ -35,6 +35,8 @@
 namespace tvm {
 namespace relay {
 
+using Any = tvm::ir::Any;
+
 /*! \brief Base type of the Relay type hiearchy. */
 class TypeNode : public RelayNode {
  public:
@@ -384,6 +386,7 @@ class TypeReporterNode : public Node {
    *  But it is possible for the solver to resolve src by dst as well.
    */
   TVM_DLL virtual void Assign(const Type& dst, const Type& src) = 0;
+
   /*!
    * \brief assert shape expression comparison.
    * \note Use assert only if any of the condition input is symbolic.

include/tvm/runtime/ndarray.h

@@ -190,6 +190,8 @@ class NDArray {
   TVM_DLL static void CopyFromTo(
       DLTensor* from, DLTensor* to, TVMStreamHandle stream = nullptr);
 
+  TVM_DLL std::vector<int64_t> Shape() const;
+
   // internal namespace
   struct Internal;
  protected:

python/tvm/_ffi/base.py

@@ -294,7 +294,7 @@ def get_last_ffi_error():
     """
     c_err_msg = py_str(_LIB.TVMGetLastError())
     py_err_msg, err_type = c2pyerror(c_err_msg)
-    if err_type.startswith("tvm.error."):
+    if err_type is not None and err_type.startswith("tvm.error."):
         err_type = err_type[10:]
     return ERROR_TYPE.get(err_type, TVMError)(py_err_msg)
 

python/tvm/api.py

@@ -479,7 +479,8 @@ def extern(shape,
             raise ValueError("nested tag is not allowed for now")
         tag = _tag.TagScope.get_current().tag
     shape = (shape,) if isinstance(shape, (_expr.Expr, _Integral)) else shape
-    shape = [shape] if isinstance(shape[0], (_expr.Expr, _Integral)) else shape
+    if shape == () or isinstance(shape[0], (_expr.Expr, _Integral)):
+        shape = [shape]
     if in_buffers is not None:
         in_buffers = [in_buffers] if not isinstance(in_buffers, list) else in_buffers
         if len(inputs) != len(in_buffers):

python/tvm/relay/__init__.py

@@ -63,6 +63,7 @@
 TensorType = ty.TensorType
 Kind = ty.Kind
 TypeVar = ty.TypeVar
+ShapeVar = ty.ShapeVar
 TypeConstraint = ty.TypeConstraint
 FuncType = ty.FuncType
 TypeRelation = ty.TypeRelation
@@ -71,6 +72,7 @@
 RefType = ty.RefType
 GlobalTypeVar = ty.GlobalTypeVar
 TypeCall = ty.TypeCall
+Any = ty.Any
 
 # Expr
 Expr = expr.Expr

python/tvm/relay/expr.py

@@ -570,6 +570,7 @@ def const(value, dtype=None):
     """
     if isinstance(value, (_base.numeric_types, (bool, list))):
         value = _np.array(value, dtype=dtype)
+
     if not dtype:
         # when dtype is None: int maps to "int32", float maps to "float32"
         map_dtype = {
@@ -578,6 +579,7 @@ def const(value, dtype=None):
             }.get(value.dtype, None)
         if map_dtype:
             value = value.astype(map_dtype)
+
     if isinstance(value, (_np.ndarray, _np.generic)):
         value = _nd.array(value)
 

python/tvm/relay/frontend/mxnet.py

@@ -491,9 +491,9 @@ def _mx_arange(inputs, attrs):
         raise tvm.error.OpAttributeUnimplemented(
             'Attribute "repeat" is not supported in operator arange.')
     new_attrs = {}
-    new_attrs["start"] = attrs.get_float("start", 0)
-    new_attrs["stop"] = attrs.get_float("stop")
-    new_attrs["step"] = attrs.get_float("step", 1)
+    new_attrs["start"] = _expr.const(attrs.get_float("start", 0.0))
+    new_attrs["stop"] = _expr.const(attrs.get_float("stop"))
+    new_attrs["step"] = _expr.const(attrs.get_float("step", 1.0))
     new_attrs["dtype"] = attrs.get_str("dtype", "float32")
     return _op.arange(**new_attrs)
 

python/tvm/relay/frontend/tensorflow.py

@@ -1059,9 +1059,9 @@ def _impl(inputs, attr, params):
         return AttrCvt(
             op_name="arange",
             ignores=['Tidx'],
-            extras={'start': start,
-                    "stop": limit,
-                    'step': delta,
+            extras={'start': _expr.const(start),
+                    "stop": _expr.const(limit),
+                    'step': _expr.const(delta),
                     'dtype': dtype})([], attr)
     return _impl
 
@@ -1269,8 +1269,8 @@ def _impl(inputs, attr, params):
             crop = crops[axis - 1]
             if crop != [0, 0]:
                 indices = tvm.relay.arange(
-                    crop[0],
-                    reshaped_permuted_shape[axis] - crop[1],
+                    _expr.const(crop[0]),
+                    _expr.const(reshaped_permuted_shape[axis] - crop[1]),
                     dtype='int32'
                 )
                 cropped = tvm.relay.take(cropped, indices=indices, axis=axis)

python/tvm/relay/loops.py

@@ -0,0 +1,65 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=invalid-name
+"""
+Utilities for building Relay loops.
+"""
+from .scope_builder import ScopeBuilder
+from . import expr as _expr
+
+def while_loop(cond, loop_vars, loop_bodies):
+    """
+    Construct a while loop.
+
+    Parameters
+    ----------
+
+    cond: Callable[Tuple[relay.Expr], relay.Expr]
+        The condition of the loop.
+
+    loop_vars:  Tuple[relay.Expr]
+        The variables being looped over.
+        The initial values of the loop, will be used to
+        construct the loop variables.
+
+    loop_bodies: Callable[Tuple[relay.Expr], Tuple[relay.Expr]]
+        The body of the loop, should be a function which
+        given loop variables produces the output result
+        also as a tuple
+
+    Returns
+    -------
+    loop: relay.Expr
+        The loop expression.
+    """
+    sb = ScopeBuilder()
+    loop = _expr.Var("while_loop")
+    fresh_vars = []
+
+    for i, loop_var in enumerate(loop_vars):
+        name = loop_var.name_hint if isinstance(loop_var, _expr.Var) else "arg{}".format(i)
+        new_var = _expr.var(name, type_annotation=sb.type_of(loop_var))
+        fresh_vars.append(new_var)
+
+    with sb.if_scope(cond(*fresh_vars)):
+        sb.ret(loop(*loop_bodies(*fresh_vars)))
+    with sb.else_scope():
+        sb.ret(_expr.Tuple(fresh_vars))
+
+    func = _expr.Function(fresh_vars, sb.get())
+    let = _expr.Let(loop, func, loop)
+    return let

python/tvm/relay/op/transform.py

@@ -17,7 +17,7 @@
 """Transform operators."""
 
 from . import _make
-from ..expr import TupleWrapper
+from ..expr import TupleWrapper, const
 
 
 def cast(data, dtype):
@@ -272,7 +272,7 @@ def full_like(data, fill_value):
     return _make.full_like(data, fill_value)
 
 
-def arange(start, stop=None, step=1, dtype="float32"):
+def arange(start, stop=None, step=None, dtype="float32"):
     """Return evenly spaced values within a given interval.
 
     .. note::
@@ -310,9 +310,13 @@ def arange(start, stop=None, step=1, dtype="float32"):
         relay.arange(1, 5) = [1, 2, 3, 4]
         relay.arange(1, 5, 1.5) = [1, 2.5, 4]
     """
+    if step is None:
+        step = const(1, dtype)
+
     if stop is None:
         stop = start
-        start = 0
+        start = const(0, dtype=dtype)
+
     return _make.arange(start, stop, step, dtype)
 
 

python/tvm/relay/scope_builder.py

@@ -42,7 +42,6 @@ def __exit__(self, ptype, value, trace):
         else:
             self._exit_cb()
 
-
 def _make_lets(bindings, ret_value):
     """Make a nested let expressions.
 
@@ -176,6 +175,24 @@ def _on_exit():
                 false_branch)
         return WithScope(None, _on_exit)
 
+
+    def type_of(self, expr):
+        """
+        Compute the type of an expression.
+
+        Parameters
+        ----------
+        expr: relay.Expr
+            The expression to compute the type of.
+        """
+        if isinstance(expr, _expr.Var):
+            return expr.type_annotation
+
+        ity = _ty.IncompleteType()
+        var = _expr.var("unify", ity)
+        self.let(var, expr)
+        return ity
+
     def ret(self, value):
         """Set the return value of this scope.
 

python/tvm/relay/ty.py

@@ -20,6 +20,7 @@
 from .base import RelayNode, register_relay_node
 from . import _make
 
+Any = _make.Any
 
 class Type(RelayNode):
     """The base type for all Relay types."""
@@ -137,6 +138,19 @@ def __init__(self, var, kind=Kind.Type):
         """
         self.__init_handle_by_constructor__(_make.TypeVar, var, kind)
 
+def ShapeVar(name):
+    """A helper which constructs a type var of which the shape kind.
+
+    Parameters
+    ----------
+    name : str
+
+    Returns
+    -------
+    type_var : tvm.relay.TypeVar
+        The shape variable.
+    """
+    return TypeVar(name, kind=Kind.ShapeVar)
 
 @register_relay_node
 class GlobalTypeVar(Type):

src/codegen/llvm/codegen_llvm.cc

@@ -970,7 +970,9 @@ llvm::Value* CodeGenLLVM::VisitExpr_(const Call* op) {
              op->call_type == Call::PureExtern) {
     return CreateCallExtern(op);
   } else {
-    LOG(FATAL) << "Unknown call type ";
+    LOG(FATAL) << "Unknown call type " <<
+      "name= " << op->name <<
+      " call_type= " << op->call_type;
     return nullptr;
   }
 }