[RELAY] Refactor type inference to use type solver (#1779)

.gitignore

@@ -190,6 +190,13 @@ build*
 
 # Jetbrain
 .idea
+.ipython
+.jupyter
+.nv
+.pylint.d
+.python_history
+.pytest_cache
+.local
 
 # tmp file
 .nfs*

include/tvm/relay/op.h

@@ -13,10 +13,9 @@
 #include <utility>
 #include <vector>
 
-#include "../attrs.h"
-#include "./base.h"
-#include "./expr.h"
-#include "./type.h"
+#include "base.h"
+#include "expr.h"
+#include "type.h"
 
 namespace tvm {
 namespace relay {
@@ -45,7 +44,7 @@ class OpNode : public relay::ExprNode {
   Array<AttrFieldInfo> arguments;
   /*!
    * \brief The type key of the attribute field
-   *  This can be empty, in which case it defaults to
+   *  This can be empty, in which case it defaults to anything.
    */
   std::string attrs_type_key;
   /*!
@@ -156,11 +155,13 @@ class OpRegistry {
    */
   inline OpRegistry& add_type_rel(
       const std::string& rel_name,
-      std::function<Array<Type>(const Array<Type>&, int)> type_rel_func);
-
+      runtime::TypedPackedFunc<bool(const Array<Type>&,
+                                    int,
+                                    const Attrs&,
+                                    const TypeReporter&)> type_rel_func);
   /*!
    * \brief Set the type key of attributes.
-   * \param type_key The type of of the attrs field.x
+   * \param type_key The type of of the attrs field.
    * \return reference to self.
    */
   inline OpRegistry& set_attrs_type_key(const std::string& type_key);
@@ -348,23 +349,25 @@ inline OpRegistry& OpRegistry::add_argument(const std::string& name,
 
 inline OpRegistry& OpRegistry::add_type_rel(
     const std::string& rel_name,
-    std::function<Array<Type>(const Array<Type>&, int)> type_rel_func) {
+    runtime::TypedPackedFunc<bool(const Array<Type>&,
+                                  int,
+                                  const Attrs&,
+                                  const TypeReporter&)> type_rel_func) {
   auto func_name = std::string("tvm.relay.type_relation.") + rel_name;
-
-  TypedEnvFunc<Array<Type>(const Array<Type>&, int)> env_type_rel_func;
+  TypeRelationFn env_type_rel_func;
 
   if (runtime::Registry::Get(func_name)) {
     auto env_func = EnvFunc::Get(func_name);
     env_type_rel_func = env_func;
   } else {
     runtime::Registry::Register(func_name)
-        .set_body_typed<Array<Type>(const Array<Type>&, int)>(type_rel_func);
+        .set_body(type_rel_func.packed());
     auto env_func = EnvFunc::Get(func_name);
     env_type_rel_func = env_func;
   }
 
-  std::vector<TypeParam> type_params;
-  std::vector<Type> arg_types;
+  Array<TypeParam> type_params;
+  Array<Type> arg_types;
 
   // Add inputs.
   std::string input_name_prefix = "in";
@@ -375,15 +378,27 @@ inline OpRegistry& OpRegistry::add_type_rel(
     arg_types.push_back(param);
   }
 
-  auto ty_call_args = Array<Type>(arg_types);
+  Array<Type> ty_call_args = arg_types;
 
   // Add output type.
   auto out_param = TypeParamNode::make("out", TypeParamNode::Kind::kType);
   type_params.push_back(out_param);
+  // this will trigger copy on write.
   ty_call_args.push_back(out_param);
 
+  // The attributes of primitive op is nullptr
+  //
+  // The attributes of primitive operator can vary at the call site.
+  // The type of sum is also dependent on Attrs being passed.
+  // So puting nullptr in the Attrs means that the operator is polymorphic on Attrs.
+  //
+  // A common example is sum(x, axis), where the choice of axis
+  // can affect the type of the function.
   TypeConstraint type_rel =
-      TypeRelationNode::make(rel_name, env_type_rel_func, ty_call_args);
+      TypeRelationNode::make(env_type_rel_func,
+                             ty_call_args,
+                             arg_types.size(),
+                             Attrs());
 
   auto func_type =
       FuncTypeNode::make(arg_types, out_param, type_params, {type_rel});

include/tvm/relay/pass.h

@@ -26,7 +26,7 @@ namespace relay {
  * \return A type checked expression with its checked_type field populated.
  */
 Expr InferType(const Environment& env, const Expr& e);
-Expr InferType(const Environment& env, const GlobalVar& v, const Function& e);
+Expr InferType(const Environment& env, const GlobalVar& var, const Function& f);
 
 /*!
  * \brief Check that types are well formed by applying "kinding rules".
@@ -69,7 +69,7 @@ bool AlphaEqual(const Expr& e1, const Expr& e2);
  *
  * For example: `forall s, Tensor[f32, s]` is equal to
  * `forall w, Tensor[f32, w]`.
- * 
+ *
  * See https://en.wikipedia.org/wiki/Lambda_calculus#Alpha_equivalence
  * for more details.
  *

include/tvm/relay/type.h

@@ -11,7 +11,8 @@
 #include <tvm/node/node.h>
 #include <string>
 
-#include "./base.h"
+#include "base.h"
+#include "../attrs.h"
 
 namespace tvm {
 namespace relay {
@@ -116,10 +117,10 @@ class TypeParamNode : public TypeNode {
   /*! \brief possible kinds of TypeParam */
   enum Kind : int {
     /*! \brief template variable in shape expression */
-    kShapeVar = 0,
-    kShape = 1,
+    kType = 0,
+    kShapeVar = 1,
     kBaseType = 2,
-    kType = 3
+    kShape = 3
   };
   /*!
    * \brief The variable itself is only meaningful when
@@ -144,6 +145,33 @@ class TypeParamNode : public TypeNode {
 RELAY_DEFINE_NODE_REF(TypeParam, TypeParamNode, Type);
 
 /*!
+ * \brief IncompleteType.
+ * This is intermediate values that is used during type inference.
+ *
+ * If we view the type relations as "computational graph of types",
+ * then IncompleteType represents intermediate values of the graph,
+ * TypeParam represents the input to the graph.
+ */
+class IncompleteType;
+
+/*! \brief IncompleteType container node */
+class IncompleteTypeNode : public TypeNode {
+ public:
+  TypeParamNode::Kind kind;
+
+  void VisitAttrs(tvm::AttrVisitor* v) final {
+    v->Visit("kind", &kind);
+  }
+
+  TVM_DLL static IncompleteType make(TypeParamNode::Kind kind);
+
+  static constexpr const char* _type_key = "relay.IncompleteType";
+  TVM_DECLARE_NODE_TYPE_INFO(IncompleteTypeNode, TypeNode);
+};
+
+RELAY_DEFINE_NODE_REF(IncompleteType, IncompleteTypeNode, Type);
+
+/*!
  * \brief Potential Constraints in the type.
  * \note This is reserved for future use.
  */
@@ -190,7 +218,8 @@ class FuncTypeNode : public TypeNode {
     v->Visit("span", &span);
   }
 
-  TVM_DLL static FuncType make(tvm::Array<Type> arg_types, Type ret_type,
+  TVM_DLL static FuncType make(tvm::Array<Type> arg_types,
+                               Type ret_type,
                                tvm::Array<TypeParam> type_params,
                                tvm::Array<TypeConstraint> type_constraints);
 
@@ -200,11 +229,102 @@ class FuncTypeNode : public TypeNode {
 
 RELAY_DEFINE_NODE_REF(FuncType, FuncTypeNode, Type);
 
+/*!
+ * \brief The type of tuple values.
+ */
+class TupleType;
+/*!
+ * \brief TupleType container.
+ */
+class TupleTypeNode : public TypeNode {
+ public:
+  /*! \brief The type of each field in the tuple. */
+  tvm::Array<Type> fields;
+
+  TupleTypeNode() {}
+
+  void VisitAttrs(tvm::AttrVisitor* v) final { v->Visit("fields", &fields); }
+
+  TVM_DLL static TupleType make(tvm::Array<Type> fields);
+
+  static constexpr const char* _type_key = "relay.TypeTuple";
+  TVM_DECLARE_NODE_TYPE_INFO(TupleTypeNode, TypeNode);
+};
+
+RELAY_DEFINE_NODE_REF(TupleType, TupleTypeNode, Type);
+
+class TypeReporter;
+
+/*!
+ * \brief reporter that reports back to the
+ *  type resolution information.
+ */
+class TypeReporterNode : public Node {
+ public:
+  /*!
+   * \brief Create a type equality constraint.
+   *
+   *  The "assign direction" acts as a hint to the solver
+   *  showing that it is more likely to resolve dst by src.
+   *  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 equals each other.
+   * \param lhs The left operand.
+   * \param rhs The right operand.
+   */
+  TVM_DLL virtual void AssertEQ(const ShapeExpr& lhs, const ShapeExpr& rhs) = 0;
+
+  // solver is not serializable.
+  void VisitAttrs(tvm::AttrVisitor* v) final {}
+
+  static constexpr const char* _type_key = "relay.TypeReporter";
+  TVM_DECLARE_NODE_TYPE_INFO(TypeReporterNode, Node);
+};
+
+/*!
+ * \brief Container class of TypeReporter.
+ * \sa TypeReporterNode
+ */
+class TypeReporter : public NodeRef {
+ public:
+  TypeReporter() {}
+  explicit TypeReporter(::tvm::NodePtr<::tvm::Node> n) : NodeRef(n) {
+  }
+  TypeReporterNode* operator->() const {
+    return static_cast<TypeReporterNode*>(node_.get());
+  }
+  using ContainerType = TypeReporterNode;
+};
+
+/*!
+ * \brief User defined type constraint function.
+ *
+ * If the input type information can be used to fully decide
+ * the IncompleteTypes, then the function should call
+ * reporter.Assign to report the new types, and return true.
+ * Otherwise, the function should return false.
+ *
+ * \param args The arguments to the relation.
+ *   The types are stored in the form of
+ *   [input_type_0, input_type_1, ... input_type_n,
+ *    output_type_0, output_type_1, ... output_type_m]
+ *
+ * \param num_inputs Number of input types in the args.
+ * \param attrs The additional attributes of the operator.
+ * \param reporter The reporter to report solution to.
+ * \return false if This relation cannot be resolved.
+ *   true if this relation has been resolved.
+ */
 using TypeRelationFn =
-    TypedEnvFunc<Array<Type>(const Array<Type>&, int)>;
+    TypedEnvFunc<bool(const Array<Type>& args,
+                      int num_inputs,
+                      const Attrs& attrs,
+                      const TypeReporter& reporter)>;
 
 /*!
- * \brief Opaque type relation, is an input-output relation on types.
+ * \brief User defined type relation, is an input-output relation on types.
  */
 class TypeRelation;
 /*!
@@ -214,24 +334,30 @@ class TypeRelation;
  */
 class TypeRelationNode : public TypeConstraintNode {
  public:
-  /*! \brief The name of the function */
-  std::string name;
-
   /*!
    * \brief The function on input and output variables which
    *  this is not directly serializable,
    *  need to be looked-up in the environment.
    */
-  TypeRelationFn func_;
-
+  TypeRelationFn func;
   /*! \brief The type arguments to the type function. */
   tvm::Array<Type> args;
+  /*! \brief Number of inputs arguments */
+  int num_inputs;
+  /*! \brief Attributes to the relation function */
+  Attrs attrs;
 
   void VisitAttrs(tvm::AttrVisitor* v) final {
-    v->Visit("name", &name);
+    v->Visit("func", &func);
+    v->Visit("args", &args);
+    v->Visit("num_inputs", &num_inputs);
+    v->Visit("attrs", &attrs);
   }
 
-  TVM_DLL static TypeRelation make(std::string name, TypeRelationFn func_, Array<Type> args);
+  TVM_DLL static TypeRelation make(TypeRelationFn func,
+                                   Array<Type> args,
+                                   int num_args,
+                                   Attrs attrs);
 
   static constexpr const char* _type_key = "relay.TypeRelation";
   TVM_DECLARE_NODE_TYPE_INFO(TypeRelationNode, TypeConstraintNode);
@@ -239,30 +365,6 @@ class TypeRelationNode : public TypeConstraintNode {
 
 RELAY_DEFINE_NODE_REF(TypeRelation, TypeRelationNode, TypeConstraint);
 
-/*!
- * \brief The type of tuple values.
- */
-class TupleType;
-/*!
- * \brief TupleType container.
- */
-class TupleTypeNode : public TypeNode {
- public:
-  /*! \brief The type of each field in the tuple. */
-  tvm::Array<Type> fields;
-
-  TupleTypeNode() {}
-
-  void VisitAttrs(tvm::AttrVisitor* v) final { v->Visit("fields", &fields); }
-
-  TVM_DLL static TupleType make(tvm::Array<Type> fields);
-
-  static constexpr const char* _type_key = "relay.TypeTuple";
-  TVM_DECLARE_NODE_TYPE_INFO(TupleTypeNode, TypeNode);
-};
-
-RELAY_DEFINE_NODE_REF(TupleType, TupleTypeNode, Type);
-
 // The following fields contains advanced typing
 // Only keep the class name and reserved for future usage.
 class GenericTensorType;

include/tvm/runtime/packed_func.h

@@ -175,7 +175,17 @@ class TypedPackedFunc<R(Args...)> {
    *
    * \param packed The packed function
    */
-  inline explicit TypedPackedFunc(PackedFunc packed);
+  inline TypedPackedFunc(PackedFunc packed);  // NOLINT(*)
+  /*!
+   * \brief constructor from TVMRetValue
+   * \param value The TVMRetValue
+   */
+  inline TypedPackedFunc(const TVMRetValue& value);  // NOLINT(*)
+  /*!
+   * \brief constructor from TVMArgValue
+   * \param value The TVMArgValue
+   */
+  inline TypedPackedFunc(const TVMArgValue& value);  // NOLINT(*)
   /*!
    * \brief construct from a lambda function with the same signature.
    *
@@ -196,7 +206,7 @@ class TypedPackedFunc<R(Args...)> {
              std::is_convertible<FLambda,
                                  std::function<R(Args...)>
                                  >::value>::type>
-  explicit TypedPackedFunc(const FLambda& typed_lambda) {
+  TypedPackedFunc(const FLambda& typed_lambda) {  // NOLINT(*)
     this->AssignTypedLambda(typed_lambda);
   }
   /*!
@@ -1144,6 +1154,14 @@ TypedPackedFunc<R(Args...)>::TypedPackedFunc(PackedFunc packed)
   : packed_(packed) {}
 
 template<typename R, typename ...Args>
+TypedPackedFunc<R(Args...)>::TypedPackedFunc(const TVMRetValue& value)
+    : packed_(value.operator PackedFunc()) {}
+
+template<typename R, typename ...Args>
+TypedPackedFunc<R(Args...)>::TypedPackedFunc(const TVMArgValue& value)
+    : packed_(value.operator PackedFunc()) {}
+
+template<typename R, typename ...Args>
 template<typename FType>
 inline void TypedPackedFunc<R(Args...)>::AssignTypedLambda(FType flambda) {
   packed_ = PackedFunc([flambda](const TVMArgs& args, TVMRetValue* rv) {

python/tvm/relay/expr.py

@@ -2,16 +2,18 @@
 """The expression nodes of Relay."""
 from __future__ import absolute_import
 from .base import NodeBase, register_relay_node
-from ._ir_pass import _get_checked_type
 from . import _make
 from .. import convert
 
 
 class Expr(NodeBase):
     """The base type for all Relay expressions."""
-
     def checked_type(self):
-        return _get_checked_type(self)
+        ret = self._checked_type_
+        if ret is None:
+            raise ValueError("The type checker has not populated"
+                             " the checked_type for this node")
+        return ret
 
     def __call__(self, *args):
         converted_args = []

python/tvm/relay/ty.py

@@ -52,11 +52,10 @@ class Kind(IntEnum):
        with. For example one's of kind BaseType can only be `float32`, `int32`,
        and so on.
     """
-    ShapeVar = 0
-    Shape = 1
+    Type = 0
+    ShapeVar = 1
     BaseType = 2
-    Type = 3
-
+    Shape = 3
 
 @register_relay_node
 class TypeParam(Type):
@@ -68,7 +67,7 @@ class TypeParam(Type):
     functions which are generic over types.
     """
 
-    def __init__(self, var, kind):
+    def __init__(self, var, kind=Kind.Type):
         """Construct a TypeParam.
 
         Parameters
@@ -76,7 +75,7 @@ class TypeParam(Type):
         var: tvm.expr.Var
             The tvm.Var which backs the type parameter.
 
-        kind: Kind
+        kind: Kind, optional
             The kind of the type parameter.
 
         Returns
@@ -130,8 +129,7 @@ class FuncType(Type):
                  arg_types,
                  ret_type,
                  type_params,
-                 type_constraints
-                ):
+                 type_constraints):
         """Construct a function type.
 
         Parameters
@@ -153,6 +151,29 @@ class FuncType(Type):
 @register_relay_node
 class IncompleteType(Type):
     """An incomplete type."""
-
-    def __init__(self, kind):
+    def __init__(self, kind=Kind.Type):
         self.__init_handle_by_constructor__(_make.IncompleteType, kind)
+
+
+@register_relay_node
+class TypeRelation(TypeConstraint):
+    """Type relation in relay.
+
+    Parameters
+    ----------
+    func : EnvFunc
+        User defined relation function.
+
+    args : list of types
+        List of types to the func.
+
+    num_inputs: int
+        Number of input arguments in args,
+        this act as a hint for type inference.
+
+    attrs : Attrs
+        The attribute attached to the relation information
+    """
+    def __init__(self, func, args, num_inputs, attrs):
+        self.__init_handle_by_constructor__(_make.TypeRelation,
+                                            func, args, num_inputs, attrs)

src/common/arena.h

+/*!
+ * Copyright 2018 by Contributors
+ *
+ * \file arena.h
+ * \brief Arena allocator that allocates
+ *  memory chunks and frees them all during destruction time.
+ */
+#ifndef TVM_COMMON_ARENA_H_
+#define TVM_COMMON_ARENA_H_
+
+#include <type_traits>
+
+namespace tvm {
+namespace common {
+
+const constexpr int kArenaPageSize = 16 << 10;
+
+/*!
+ * \brief Arena allocator that allocates memory from continuous
+ *  chunk and frees them all only during destruction.
+ */
+class Arena {
+ public:
+  Arena() {
+    // eagerly allocate the first page.
+    head_ = reinterpret_cast<PageHeader*>(new Page());
+    head_->next = nullptr;
+    head_->ptr = sizeof(PageHeader);
+  }
+  ~Arena() {
+    // delete all the allocated pages.
+    while (head_ != nullptr) {
+      Page* page = reinterpret_cast<Page*>(head_);
+      head_ = head_->next;
+      delete page;
+    }
+  }
+  /*!
+   * \brief Allocate a space from Arena for type T
+   * \param T the data type to be allocated
+   */
+  template<typename T>
+  T* Alloc() {
+    return static_cast<T*>(Alloc(sizeof(T), alignof(T)));
+  }
+
+ private:
+  // page size 16 KB
+  // The page data type;
+  using Page = std::aligned_storage<kArenaPageSize, 1024>::type;
+  /*! \brief Page header */
+  struct PageHeader {
+    /*! \brief points to the next page */
+    PageHeader* next;
+    /*! \brief memory allocator ptr inside page */
+    size_t ptr;
+  };
+  /* \brief The page header */
+  PageHeader* head_{nullptr};
+  /*!
+   * \brief Align ptr by upper bound.
+   * \param ptr The pointer value.
+   * \param align The alignment requirement.
+   */
+  size_t UpperAlign(size_t ptr, size_t align) {
+    return ptr + (align - (ptr % align)) % align;
+  }
+  /*!
+   * \brief Internal aligned alloc function.
+   * \param size The size of the memory.
+   * \param align The alignment requirement.
+   */
+  void* Alloc(size_t size, size_t align) {
+    size_t ptr = UpperAlign(head_->ptr, align);
+    if (ptr + size <= kArenaPageSize) {
+      head_->ptr = ptr + size;
+      return reinterpret_cast<char*>(head_) + ptr;
+    } else {
+      PageHeader* new_head = reinterpret_cast<PageHeader*>(new Page());
+      new_head->next = head_;
+      ptr = UpperAlign(sizeof(PageHeader), align);
+      CHECK_LE(ptr + size, kArenaPageSize);
+      new_head->ptr = ptr + size;
+      head_ = new_head;
+      return reinterpret_cast<char*>(head_) + ptr;
+    }
+  }
+};
+
+}  // namespace common
+}  // namespace tvm
+#endif  // TVM_COMMON_ARENA_H_

src/relay/ir/environment.cc

@@ -6,7 +6,6 @@
 #include <tvm/relay/environment.h>
 #include <tvm/relay/pass.h>
 #include <sstream>
-#include "./../pass/resolve.h"
 
 namespace tvm {
 namespace relay {
@@ -49,7 +48,7 @@ void EnvironmentNode::Add(const GlobalVar &var,
     auto checked_func = GetRef<Function>(func_node);
     auto type = checked_func->checked_type();
 
-    CHECK(IsFullyResolved(type));
+    CHECK(type.as<IncompleteTypeNode>() == nullptr);
 
     if (functions.find(var) != functions.end()) {
       if (!update) {
@@ -68,7 +67,7 @@ void EnvironmentNode::Add(const GlobalVar &var,
       this->functions.Set(var, checked_func);
     }
   } else {
-    throw Error("internal error: unknown item type, unreachable code");
+    LOG(FATAL) << "internal error: unknown item type, unreachable code";
   }
 }
 

src/relay/ir/type.cc

@@ -55,7 +55,27 @@ TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable)
     << node->kind << ")";
 });
 
-FuncType FuncTypeNode::make(tvm::Array<Type> arg_types, Type ret_type,
+IncompleteType IncompleteTypeNode::make(TypeParamNode::Kind kind) {
+  auto n = make_node<IncompleteTypeNode>();
+  n->kind = std::move(kind);
+  return IncompleteType(n);
+}
+
+TVM_REGISTER_API("relay._make.IncompleteType")
+.set_body([](TVMArgs args, TVMRetValue* ret) {
+    int kind = args[0];
+    *ret = IncompleteTypeNode::make(static_cast<TypeParamNode::Kind>(kind));
+  });
+
+TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable)
+.set_dispatch<IncompleteTypeNode>(
+    [](const IncompleteTypeNode* node,
+       tvm::IRPrinter* p) {
+      p->stream << "IncompleteTypeNode(" << node->kind << ", " << node << ")";
+    });
+
+FuncType FuncTypeNode::make(tvm::Array<Type> arg_types,
+                            Type ret_type,
                             tvm::Array<TypeParam> type_params,
                             tvm::Array<TypeConstraint> type_constraints) {
   NodePtr<FuncTypeNode> n = make_node<FuncTypeNode>();
@@ -79,24 +99,28 @@ TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable)
             << node->type_constraints << ")";
 });
 
-TypeRelation TypeRelationNode::make(std::string name, TypeRelationFn func, Array<Type> args) {
+TypeRelation TypeRelationNode::make(TypeRelationFn func,
+                                    Array<Type> args,
+                                    int num_inputs,
+                                    Attrs attrs) {
   NodePtr<TypeRelationNode> n = make_node<TypeRelationNode>();
-  n->name = std::move(name);
-  n->func_ = std::move(func);
+  n->func = std::move(func);
   n->args = std::move(args);
+  n->num_inputs = num_inputs;
+  n->attrs = std::move(attrs);
   return TypeRelation(n);
 }
 
 TVM_REGISTER_API("relay._make.TypeRelation")
 .set_body([](TVMArgs args, TVMRetValue *ret) {
-  *ret = TypeRelationNode::make(args[0], args[1], args[2]);
+    *ret = TypeRelationNode::make(args[0], args[1], args[2], args[3]);
 });
 
 TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable)
-.set_dispatch<TypeRelationNode>([](const TypeRelationNode *node,
-                                       tvm::IRPrinter *p) {
-  p->stream << "TypeRelationNode(" << node->name << ", " << node->args
-    << ")";
+.set_dispatch<TypeRelationNode>([](const TypeRelationNode *node, tvm::IRPrinter *p) {
+    p->stream << "TypeRelationNode("
+              << node->func->name
+              << ", " << node->args << ")";
 });
 
 TupleType TupleTypeNode::make(Array<Type> fields) {

src/relay/op/type_relations.cc

@@ -8,7 +8,6 @@
 #include <tvm/relay/logging.h>
 #include <tvm/relay/op.h>
 #include <numeric>
-#include "../pass/incomplete_type.h"
 #include "./type_relations.h"
 
 namespace tvm {
@@ -30,18 +29,19 @@ int ToInt(const tvm::Expr& e) {
   return imm->value;
 }
 
-Array<Type> IdentityRel(const Array<Type>& types, int num_args) {
-  CHECK_EQ(types.size(), 2);
-  auto t1 = ToTensorType(types[0]);
-  if (t1 && types[1].as<IncompleteTypeNode>()) {
-    return {t1, t1};
-  } else {
-    return types;
+bool IdentityRel(const Array<Type>& types,
+                 int num_inputs,
+                 const Attrs& attrs,
+                 const TypeReporter& reporter) {
+  for (size_t i = 1; i < types.size(); ++i) {
+    reporter->Assign(types[i], types[0]);
   }
+  return true;
 }
 
-static Type ConcreteBroadcast(const TensorType& t1, const TensorType& t2,
-                              DataType output_dtype) {
+Type ConcreteBroadcast(const TensorType& t1,
+                       const TensorType& t2,
+                       DataType output_dtype) {
   RELAY_LOG(INFO) << "ConcreteBroadcast: t1=" << t1 << " t2=" << t2
                   << std::endl;
   auto sh1 = t1->shape;
@@ -73,7 +73,7 @@ static Type ConcreteBroadcast(const TensorType& t1, const TensorType& t2,
     Array<ShapeExpr> smaller;
 
     for (int i = 0; i < (full_len - suffix_len); i++) {
-      smaller.push_back(tvm::ir::IntImm::make(HalideIR::Int(64), 1));
+      smaller.push_back(make_const(tvm::Int(64), 1));
     }
 
     if (sh1.size() < sh2.size()) {
@@ -93,46 +93,52 @@ static Type ConcreteBroadcast(const TensorType& t1, const TensorType& t2,
 
     CHECK_EQ(larger.size(), smaller.size());
 
-    Array<HalideIR::Expr> out_shape;
+    Array<ShapeExpr> out_shape;
     for (size_t i = 0; i < smaller.size(); i++) {
       auto left = smaller[i].as<tvm::ir::IntImm>();
       auto right = larger[i].as<tvm::ir::IntImm>();
       CHECK(left);
       CHECK(right);
       int64_t dim = std::max(left->value, right->value);
-      out_shape.push_back(tvm::ir::IntImm::make(HalideIR::Int(64), dim));
+      out_shape.push_back(make_const(tvm::Int(64), dim));
     }
 
     return TensorTypeNode::make(out_shape, output_dtype);
   }
 }
 
-Array<Type> BroadcastRel(const Array<Type>& types, int num_args) {
+bool BroadcastRel(const Array<Type>& types,
+                  int num_inputs,
+                  const Attrs& attrs,
+                  const TypeReporter& reporter) {
   CHECK_EQ(types.size(), 3);
   RELAY_LOG(INFO) << "In1: " << types[0] << "In2: " << types[1]
                   << "Out: " << types[2] << std::endl;
-  if (auto t1 = ToTensorType(types[0])) {
-    if (auto t2 = ToTensorType(types[1])) {
-      CHECK_EQ(t1->dtype, t2->dtype);
-      return {t1, t2, ConcreteBroadcast(t1, t2, t1->dtype)};
+  if (auto t0 = ToTensorType(types[0])) {
+    if (auto t1 = ToTensorType(types[1])) {
+      CHECK_EQ(t0->dtype, t1->dtype);
+      reporter->Assign(types[2], ConcreteBroadcast(t0, t1, t0->dtype));
+      return true;
     }
   }
-
-  return types;
+  return false;
 }
 
-/* A relation which specifies broadcasting rules for operations which
-   compute boolean results.
-*/
-Array<Type> BroadcastCompRel(const Array<Type>& types, int num_args) {
+bool BroadcastCompRel(const Array<Type>& types,
+                      int num_inputs,
+                      const Attrs& attrs,
+                      const TypeReporter& reporter) {
   CHECK_EQ(types.size(), 3);
-  if (auto t1 = ToTensorType(types[0])) {
-    if (auto t2 = ToTensorType(types[1])) {
-      return {t1, t2, ConcreteBroadcast(t1, t2, HalideIR::Bool())};
+  RELAY_LOG(INFO) << "In1: " << types[0] << "In2: " << types[1]
+                  << "Out: " << types[2] << std::endl;
+  if (auto t0 = ToTensorType(types[0])) {
+    if (auto t1 = ToTensorType(types[1])) {
+      CHECK_EQ(t0->dtype, t1->dtype);
+      reporter->Assign(types[2], ConcreteBroadcast(t0, t1, ::tvm::Bool()));
+      return true;
     }
   }
-
-  return types;
+  return false;
 }
 
 /*! \brief Handle concrete concat case from known input to output. */
@@ -175,10 +181,10 @@ inline Type ConcreteConcatRel(const Type& input_type) {
 
     auto out_axis_dim = std::accumulate(axis_dims.begin(), axis_dims.end(), 0);
 
-    Array<tvm::Expr> out_shape = { tvm::ir::IntImm::make(HalideIR::Int(64), out_axis_dim) };
+    Array<tvm::Expr> out_shape = { make_const(Int(64), out_axis_dim) };
 
     for (auto dim : dims) {
-      out_shape.push_back(tvm::ir::IntImm::make(HalideIR::Int(64), dim));
+      out_shape.push_back(make_const(Int(64), dim));
     }
 
     return TensorTypeNode::make(out_shape, dtype);
@@ -188,19 +194,18 @@ inline Type ConcreteConcatRel(const Type& input_type) {
   }
 }
 
-Array<Type> ConcatRel(const Array<Type>& types, int num_args) {
+bool ConcatRel(const Array<Type>& types,
+               int num_inputs,
+               const Attrs& attrs,
+               const TypeReporter& reporter) {
   CHECK_EQ(types.size(), 2);
-
-  if (types[0].as<IncompleteTypeNode>() && types[1].as<IncompleteTypeNode>()) {
-    return types;
-  } else if (types[1].as<IncompleteTypeNode>()) {
-    return { types[0], ConcreteConcatRel(types[0]) };
-  } else {
-    throw TypeRelationError(
-      "can not deduce relationship between the " \
-      "type of concat's input and output");
+  if (types[0].as<TupleTypeNode>()) {
+    reporter->Assign(types[1], ConcreteConcatRel(types[0]));
+    return true;
   }
+  return false;
 }
 
+
 }  // namespace relay
 }  // namespace tvm

src/relay/op/type_relations.h

@@ -24,42 +24,72 @@ struct TypeRelationError : Error {
       : Error(msg) {}
 };
 
-/*! \brief The identity type relation maps a single input variable
- * to the output variable.
+/*!
+ * \brief The identity type relation, all the types are equal.
  *
  * \param types The input and output types to the relation.
- * \param num_args The number of input arguments.
- * \return The (potentially partial) solution to the relation.
+ * \param num_inputs The number of input arguments.
+ * \param attrs The attributes
+ * \param reporter The reporter.
+ * \return true whether relation has been resolved.
  */
-Array<Type> IdentityRel(const Array<Type>& types, int num_args);
-/*! \brief The broadcast type relation, implements the broadcasting
+bool IdentityRel(const Array<Type>& types,
+                 int num_inputs,
+                 const Attrs& attrs,
+                 const TypeReporter& reporter);
+
+/*!
+ * \brief The broadcast type relation, implements the broadcasting
  * rule over the two input types producing the broadcasted type.
  *
  * \param types The input and output types to the relation.
- * \param num_args The number of input arguments.
- * \return The (potentially partial) solution to the relation.
+ * \param num_inputs The number of input arguments.
+ * \param attrs The attributes
+ * \param reporter The reporter.
+ * \return true whether relation has been resolved.
  */
-Array<Type> BroadcastRel(const Array<Type>& types, int num_args);
-/*! \brief The broadcast type relation, implements the broadcasting
- * rule over the two input types producing the broadcasted type.
+bool BroadcastRel(const Array<Type>& types,
+                  int num_inputs,
+                  const Attrs& attrs,
+                  const TypeReporter& reporter);
+
+/*!
+ * \brief The broadcast type relation, implements the broadcasting
+ *  rule over the two input types producing the broadcasted type.
  *
  * This differs from BroadcastRel in the return dtype,
- * it instead returns bool, for use in comparsion operators
+ * it instead returns bool(uint8), for use in comparsion operators
  * such as equal, not_equal, lt, and so on.
  *
  * \param types The input and output types to the relation.
- * \param num_args The number of input arguments.
- * \return The (potentially partial) solution to the relation.
+ * \param num_inputs The number of input arguments.
+ * \param attrs The attributes
+ * \param reporter The reporter.
+ * \return true whether relation has been resolved.
  */
-Array<Type> BroadcastCompRel(const Array<Type>& types, int num_args);
+bool BroadcastCompRel(const Array<Type>& types,
+                      int num_inputs,
+                      const Attrs& attrs,
+                      const TypeReporter& reporter);
 
-/*! \brief The concat relation.
+/*!
+ * \brief The The concat relation, implements the broadcasting
+ *  rule over the two input types producing the broadcasted type.
  *
- * This relation takes a single input which must be a single tensor
- * or an arbitrary sized tuple. It combines these input dimensions
- * together to produce the output example.
+ * This differs from BroadcastRel in the return dtype,
+ * it instead returns bool(uint8), for use in comparsion operators
+ * such as equal, not_equal, lt, and so on.
+ *
+ * \param types The input and output types to the relation.
+ * \param num_inputs The number of input arguments.
+ * \param attrs The attributes
+ * \param reporter The reporter.
+ * \return true whether relation has been resolved.
  */
-Array<Type> ConcatRel(const Array<Type>& types, int num_args);
+bool ConcatRel(const Array<Type>& types,
+               int num_inputs,
+               const Attrs& attrs,
+               const TypeReporter& reporter);
 
 }  // namespace relay
 }  // namespace tvm

src/relay/pass/incomplete_type.h

-/*!
- *  Copyright (c) 2018 by Contributors
- * \file incomplete_type.h
- * \brief A way to defined arbitrary function signature with dispatch on types.
- */
-
-#ifndef TVM_RELAY_PASS_INCOMPLETE_TYPE_H_
-#define TVM_RELAY_PASS_INCOMPLETE_TYPE_H_
-
-#include <tvm/relay/expr.h>
-
-namespace tvm {
-namespace relay {
-
-/*!
- * \brief Represents a portion of an incomplete type.
- */
-class IncompleteType;
-
-/*! \brief IncompleteType container node */
-class IncompleteTypeNode : public TypeNode {
- public:
-  TypeParamNode::Kind kind;
-
-  void VisitAttrs(tvm::AttrVisitor* v) final { v->Visit("kind", &kind); }
-
-  TVM_DLL static IncompleteType make(TypeParamNode::Kind kind);
-
-  static constexpr const char* _type_key = "relay.IncompleteType";
-  TVM_DECLARE_NODE_TYPE_INFO(IncompleteTypeNode, TypeNode);
-};
-
-RELAY_DEFINE_NODE_REF(IncompleteType, IncompleteTypeNode, Type);
-
-}  // namespace relay
-}  // namespace tvm
-
-#endif  // TVM_RELAY_PASS_INCOMPLETE_TYPE_H_

src/relay/pass/resolve.cc

-/*!
- *  Copyright (c) 2018 by Contributors
- * \file resolve.cc
- * \brief Resolve incomplete types to complete types.
- */
-
-#include <tvm/relay/expr.h>
-#include <tvm/relay/expr_functor.h>
-#include "./resolve.h"
-#include "./type_visitor.h"
-
-namespace tvm {
-namespace relay {
-
-struct ResolveTypeType : TypeMutator {
-  const TypeUnifier &unifier;
-
-  explicit ResolveTypeType(const TypeUnifier &unifier) : unifier(unifier) {}
-
-  Type VisitType(const Type &t) override {
-    if (!t.defined()) {
-      auto inc_ty = IncompleteTypeNode::make(TypeParamNode::Kind::kType);
-      unifier->Insert(inc_ty);
-      return inc_ty;
-    } else {
-      return TypeMutator::VisitType(t);
-    }
-  }
-
-  Type VisitType_(const IncompleteTypeNode *op) override {
-    return unifier->Subst(GetRef<IncompleteType>(op));
-  }
-};
-
-struct ResolveTypeExpr : ExprMutator {
-  const TypeUnifier &unifier;
-
-  explicit ResolveTypeExpr(const TypeUnifier &unifier) : unifier(unifier) {}
-
-  Expr Mutate(const Expr &e) {
-    // NB: a bit tricky here.
-    //
-    // We want to store resolved type without having
-    // to re-typecheck the entire term.
-    //
-    // Since we know that e : T[...] under some holes
-    // then it is the case that if we resolve types
-    // present in e, then we can type it under T
-    // with the wholes filled in.
-    //
-    // We will visit e like normal building a new
-    // term, then resolve e's old type and write
-    // it back into the new node.
-    auto new_e = ExprMutator::Mutate(e);
-    CHECK(e->checked_type_.defined());
-    auto resolved_cty = VisitType(e->checked_type_);
-    new_e->checked_type_ = resolved_cty;
-    return new_e;
-  }
-
-  Type VisitType(const Type &t) {
-    return ResolveTypeType(unifier).VisitType(t);
-  }
-};
-
-Type Resolve(const TypeUnifier &unifier, const Type &ty) {
-  CHECK(ty.defined());
-  return ResolveTypeType(unifier).VisitType(ty);
-}
-
-Expr Resolve(const TypeUnifier &unifier, const Expr &expr) {
-  return ResolveTypeExpr(unifier).Mutate(expr);
-}
-
-struct FullyResolved : TypeVisitor<> {
-  bool incomplete;
-
-  FullyResolved() : incomplete(true) {}
-
-  void VisitType(const Type &t) override {
-    if (!t.defined()) {
-      incomplete = true;
-    } else {
-      return TypeVisitor<>::VisitType(t);
-    }
-  }
-
-  void VisitType_(const IncompleteTypeNode *ty_var) override {
-    incomplete = false;
-  }
-};
-
-bool IsFullyResolved(const Type &t) {
-  auto fr = FullyResolved();
-  fr.VisitType(t);
-  return fr.incomplete;
-}
-
-}  // namespace relay
-}  // namespace tvm

src/relay/pass/resolve.h

-/*!
- *  Copyright (c) 2018 by Contributors
- * \file tvm/relay/resolve.h
- * \brief Resolve incomplete types to complete types.
- */
-#ifndef TVM_RELAY_PASS_RESOLVE_H_
-#define TVM_RELAY_PASS_RESOLVE_H_
-
-#include <tvm/relay/expr.h>
-#include <string>
-#include "./unifier.h"
-
-namespace tvm {
-namespace relay {
-
-/*! \brief Resolve a type containing incomplete types.
- *
- * This pass replaces incomplete types with their representative, and
- * converts types which are not defined into fresh variables.
- *
- * \param unifier The unifier containing the unification data.
- * \param ty The type to resolve.
- * \returns The resolved type.
- */
-Type Resolve(const TypeUnifier& unifier, const Type& ty);
-
-/*! \brief Resolve an expression containing incomplete types.
- *
- * This pass replaces incomplete types with their representative, and
- * converts types which are not defined into fresh variables.
- *
- * \param unifier The unifier containing the unification data.
- * \param ty The expression to resolve.
- * \returns The resolved expression.
- */
-Expr Resolve(const TypeUnifier& unifier, const Expr& expr);
-
-/*! \brief Check if all types have been filled in.
- *   \param t The type.
- *   \returns True if the type is resolved, false otherwise.
- */
-bool IsFullyResolved(const Type& t);
-
-}  // namespace relay
-}  // namespace tvm
-
-#endif  // TVM_RELAY_PASS_RESOLVE_H_

src/relay/pass/type_functor.h

@@ -8,7 +8,6 @@
 
 #include <tvm/node/ir_functor.h>
 #include <tvm/relay/expr.h>
-#include "./incomplete_type.h"
 
 namespace tvm {
 namespace relay {

src/relay/pass/type_infer.cc

@@ -22,607 +22,360 @@
 
 #include <tvm/relay/error.h>
 #include <tvm/relay/expr_functor.h>
-#include <tvm/relay/logging.h>
 #include <tvm/relay/pass.h>
-#include "./incomplete_type.h"
-#include "./resolve.h"
-#include "./type_subst.h"
-#include "./type_visitor.h"
-#include "./unifier.h"
+#include "type_solver.h"
+#include "type_subst.h"
 
 namespace tvm {
 namespace relay {
-
-using namespace tvm::runtime;
-
-// // We declare this for forward compatibility.
-struct ConstraintData {};
-
-/*! \brief A more efficient representation of the type relation
- * data needed for type checking.
- */
-struct TypeRelationData : ConstraintData {
-  std::string name;
-  std::vector<Type> args;
-  TypeRelationFn func;
-  Span span;
-
-  explicit TypeRelationData(const TypeRelation& ty_rel)
-      : TypeRelationData(ty_rel->args, ty_rel->func_, ty_rel->span) {}
-
-  TypeRelationData(const Array<Type>& args, const TypeRelationFn& func, const Span& sp)
-      : func(func), span(sp) {
-    for (auto arg : args) {
-      this->args.push_back(arg);
-    }
+//
+// The inference algorithm can roughly be devided into three stages:
+// - Populate the constraints by visiting the expression (TypeInferencer.GetType)
+//   - solver.AddConstraint and solver.Unify are called to populate the necessary constraints
+// - Solve the constraints (solver_.Solve)
+// - Recreate expression with the resolved checked_type (Resolver.VisitExpr)
+//
+class TypeInferencer : private ExprFunctor<Type(const Expr&)> {
+ public:
+  // constructors
+  TypeInferencer()
+      : env_(EnvironmentNode::make({})) {
   }
-
-  TypeRelation ToTypeRel() const {
-    Array<Type> args = Array<Type>(this->args.begin(), this->args.end());
-    return TypeRelationNode::make(
-        this->name, this->func, args);
+  explicit TypeInferencer(Environment env)
+      : env_(env) {
   }
-};
-
-struct TypeContext {
-  std::unordered_map<Var, Type, NodeHash> var_map;
-  std::vector<std::vector<TypeRelationData> > constraints;
 
-  TypeContext() { constraints.push_back({}); }
+  // inference the type of expr.
+  Expr Infer(Expr expr);
 
-  void Insert(const Var& id, const Type& t) { var_map[id] = t; }
-
-  void AddConstraint(const TypeConstraint& constraint) {
-      constraints.back().push_back(TypeRelationData(Downcast<TypeRelation>(constraint)));
+ private:
+  // type resolver that maps back to type
+  class Resolver;
+  // internal environment
+  Environment env_;
+  // map from expression to checked type
+  // type inferencer will populate it up
+  std::unordered_map<Expr, Type, NodeHash, NodeEqual> type_map_;
+  // The solver used by the inferencer.
+  TypeSolver solver_;
+  // Unify two types
+  Type Unify(const Type& t1, const Type& t2, const Span& span) {
+    // TODO(tqchen, jroesch): propagate span to solver
+    try {
+      return solver_.Unify(t1, t2);
+    } catch (const dmlc::Error &e) {
+      LOG(FATAL)
+          << "Error unifying `"
+          << t1
+          << "` and `"
+          << t2
+          << "`: " << e.what();
+      return Type();
+    }
   }
-
-  Type Lookup(const Var& var) {
-    auto type = var_map.find(var);
-    if (type != var_map.end()) {
-      return (*type).second;
-    } else {
-      throw FatalTypeError(std::string("undeclared local variable: ") + var->name_hint);
+  // Lazily get type for expr
+  // will call visit to deduce it if it is not in the type_map_
+  Type GetType(const Expr &expr) {
+    auto it = type_map_.find(expr);
+    if (it != type_map_.end()) {
+      return it->second;
     }
+    Type ret = this->VisitExpr(expr);
+    type_map_[expr] = ret;
+    return ret;
   }
 
-  struct Scope {
-    TypeContext& tc;
-    explicit Scope(TypeContext& tc) : tc(tc) { tc.constraints.push_back({}); }
-    ~Scope() { tc.constraints.pop_back(); }
-  };
-};
-
-struct CheckedExpr {
-  Expr expr;
-  Type type;
-  CheckedExpr(Expr e, Type t) : expr(e), type(t) {}
-  CheckedExpr() {}
-};
-
-enum SolverResult : int;
-
-class TypeInferencer : private ExprFunctor<CheckedExpr(const Expr&)> {
- private:
-  TypeContext context;
-
- public:
-  Environment env;
-  TypeUnifier unifier;
-
-  template <typename T>
-  T WithScope(const std::function<T()>& f) {
-    TypeContext::Scope fr(context);
-    return f();
+  // Visitor logics
+  Type VisitExpr_(const VarNode* op) final {
+    // The type of Var can already been lookedup in type_map_;
+    LOG(FATAL) << "Cannot find binding for var " << GetRef<Var>(op);
+    return Type();
   }
 
-  TypeInferencer();
-  TypeInferencer(Environment env, TypeUnifier unifier)
-      : env(env), unifier(unifier) {}
-  explicit TypeInferencer(Environment env);
-
-  CheckedExpr Infer(const Expr &expr);
-
-  FuncType Instantiate(FuncType fn_ty, tvm::Array<Type> &ty_args);
-
-  Type Normalize(const Type& t);
-
-  void ReportError(const std::string& msg, Span sp);
-  [[noreturn]] void FatalError(const std::string& msg, Span sp);
-
-  Type Unify(const Type &t1, const Type& t2, Span sp);
-  Type Resolve(const Type &t);
-  Expr Resolve(const Expr &e);
-
-  /*! \brief Attempt to solve a single relation. */
-  void Solve(TypeRelationData& ty_rel);
-
-  /*! \brief Attempt to solve all pending relations.
-   *
-   * If the solver
-   */
-  SolverResult Solve(std::vector<TypeRelationData>& rels);
-
-  /*! \brief Check that all relations hold. */
-  bool RelationsHold(bool scope_only = false);
-
-  /*! \brief Visit a function node, extra flag controls behavior. */
-  CheckedExpr VisitFunction(const Function& f, bool generalize);
-
- private:
-  CheckedExpr VisitExpr_(const VarNode* op) override;
-  CheckedExpr VisitExpr_(const GlobalVarNode* op) override;
-  CheckedExpr VisitExpr_(const ConstantNode* op) override;
-  CheckedExpr VisitExpr_(const TupleNode* op) override;
-  CheckedExpr VisitExpr_(const ParamNode* op) override;
-  CheckedExpr VisitExpr_(const FunctionNode* op) override;
-  CheckedExpr VisitExpr_(const CallNode* op) override;
-  CheckedExpr VisitExpr_(const LetNode* op) override;
-  CheckedExpr VisitExpr_(const IfNode* op) override;
-  CheckedExpr VisitExpr_(const OpNode* op) override;
-};
-
-TypeInferencer::TypeInferencer() {
-  this->env = EnvironmentNode::make({});
-  this->unifier = TypeUnifierNode::make(UnionFindNode::make({}));
-}
-
-TypeInferencer::TypeInferencer(Environment env) : env(env) {
-  this->unifier = TypeUnifierNode::make(UnionFindNode::make({}));
-}
-
-CheckedExpr TypeInferencer::Infer(const Expr& expr) {
-  RELAY_LOG(INFO) << "TypeInferencer::Check expr=" << expr << std::endl;
-  CheckedExpr checked_expr = this->VisitExpr(expr);
-  RELAY_LOG(INFO) << "TypeInferencer::Check type=" << checked_expr.type
-                  << std::endl;
-  Type final_type = checked_expr.type;
-  RELAY_LOG(INFO) << "TypeInferencer::Check type_after_subst=" << final_type
-                  << std::endl;
-  checked_expr.expr->checked_type_ = final_type;
-  return checked_expr;
-}
-
-CheckedExpr TypeInferencer::VisitExpr_(const VarNode* op) {
-  auto var = GetRef<Var>(op);
-  return {var, this->context.Lookup(var)};
-}
-
-CheckedExpr TypeInferencer::VisitExpr_(const GlobalVarNode* op) {
-  GlobalVar var = GetRef<GlobalVar>(op);
-  Expr e = this->env->Lookup(var);
-  return {var, e->checked_type()};
-}
-
-CheckedExpr TypeInferencer::VisitExpr_(const ConstantNode* const_node) {
-  return {GetRef<Constant>(const_node), const_node->tensor_type()};
-}
-
-CheckedExpr TypeInferencer::VisitExpr_(const TupleNode* op) {
-  Tuple pl = GetRef<Tuple>(op);
-
-  std::vector<Expr> field_exprs;
-  std::vector<Type> field_types;
-  for (auto field = pl->fields.begin(); field != pl->fields.end(); field++) {
-    auto checked_field = Infer(*field);
-    field_exprs.push_back(checked_field.expr);
-    field_types.push_back(checked_field.type);
+  Type VisitExpr_(const ParamNode* op) final {
+    // directly handled by Funtion
+    LOG(FATAL) << "not reached";
+    return Type();
   }
 
-  return {TupleNode::make(field_exprs), TupleTypeNode::make(field_types)};
-}
-
-CheckedExpr TypeInferencer::VisitExpr_(const ParamNode* param) {
-  // We should trigger error here and move param code direclty into function
-  // checking.
-  auto rtype = this->Resolve(param->type);
-  // This is a special case ... not sure if there is a better way
-  // to handle this.
-  param->var->checked_type_ = rtype;
-  return {ParamNode::make(param->var, rtype), rtype};
-}
-
-CheckedExpr TypeInferencer::VisitFunction(const Function& f, bool generalize) {
-  // First we add the parameters to the context allowing us to check their
-  // types.
-
-  // TODO(@jroesch): support polymorphism
+  Type VisitExpr_(const GlobalVarNode* op) final {
+    GlobalVar var = GetRef<GlobalVar>(op);
+    Expr e = env_->Lookup(var);
+    return e->checked_type();
+  }
 
-  std::vector<Type> param_types;
-  std::vector<Param> params;
+  Type VisitExpr_(const ConstantNode* op) final {
+    return op->tensor_type();
+  }
 
-  return this->WithScope<CheckedExpr>([&]() -> CheckedExpr {
-    for (auto param : f->params) {
-      CheckedExpr checked_param = this->Infer(param);
-      Type arg_type;
-      param_types.push_back(checked_param.type);
-      params.push_back(GetRef<Param>(checked_param.expr.as<ParamNode>()));
-      this->context.Insert(param->var, checked_param.type);
+  Type VisitExpr_(const TupleNode* op) final {
+    // TODO(tqchen, jroesch)
+    // tuple should be a constraint in the type solver
+    // to handle cases where the field type is not known.
+    Array<Type> fields;
+    for (Expr field : op->fields) {
+      fields.push_back(GetType(field));
     }
+    return TupleTypeNode::make(fields);
+  }
 
-    auto checked_body = this->Infer(f->body);
-    auto inferred_rtype = checked_body.type;
-    auto annotated_rtype = Resolve(f->ret_type);
-
-    auto unified_rtype = this->Unify(inferred_rtype, annotated_rtype, f->span);
-
-    CHECK(RelationsHold(true));
-
-    Array<TypeConstraint> cs;
+  Type VisitExpr_(const OpNode* op) final {
+    return op->op_type;
+  }
 
-    for (auto cons : this->context.constraints.back()) {
-      cs.push_back(cons.ToTypeRel());
+  Type VisitExpr_(const LetNode* op) final {
+    Type vtype = GetType(op->value);
+    if (op->value_type.defined()) {
+      vtype = Unify(vtype, op->value_type, op->span);
     }
-
-    return {FunctionNode::make(params, unified_rtype, checked_body.expr, {}),
-            FuncTypeNode::make(param_types, unified_rtype, {}, cs)};
-  });
-}
-
-CheckedExpr TypeInferencer::VisitExpr_(const FunctionNode* op) {
-  return this->VisitFunction(GetRef<Function>(op), false);
-}
-
-FuncType TypeInferencer::Instantiate(FuncType fn_ty,
-                                     tvm::Array<Type>& ty_args) {
-  tvm::Map<TypeParam, Type> subst_map;
-
-  // Build a subsitituion map up from the function type and type arguments.
-  // Eventually allow the type vars to be passed in.
-  for (auto ty_param : fn_ty->type_params) {
-    IncompleteType fresh = IncompleteTypeNode::make(ty_param->kind);
-    this->unifier->Insert(fresh);
-    ty_args.push_back(fresh);
-    subst_map.Set(ty_param, fresh);
+    CHECK(!type_map_.count(op->var));
+    // NOTE: no scoping is necessary becase var are unique in program
+    type_map_[op->var] = vtype;
+    return GetType(op->body);
   }
 
-  Type inst_ty = FuncTypeNode::make(fn_ty->arg_types, fn_ty->ret_type, {},
-                                    fn_ty->type_constraints);
-  inst_ty = TypeSubst(inst_ty, subst_map);
-
-  CHECK(KindCheck(this->env, inst_ty));
-
-  return GetRef<FuncType>(inst_ty.as<FuncTypeNode>());
-}
-
-CheckedExpr TypeInferencer::VisitExpr_(const CallNode* op) {
-  Call c = GetRef<Call>(op);
-
-  auto checked_op = this->Infer(c->op);
-
-  RELAY_LOG(INFO) << "TypeInferencer::VisitExpr_ op=" << c << std::endl
-                  << "fn_ty=" << checked_op.type << std::endl;
-
-  auto fn_ty_node = checked_op.type.as<FuncTypeNode>();
-
-  if (!fn_ty_node) {
-    this->FatalError("only expressions with function types can be called",
-                     c->op->span);
+  Type VisitExpr_(const IfNode* op) final {
+    // Ensure the type of the guard is of Tensor[Bool, ()],
+    // that is a rank-0 boolean tensor.
+    Type cond_type = this->GetType(op->cond);
+    this->Unify(cond_type,
+                TensorTypeNode::Scalar(tvm::Bool()),
+                op->cond->span);
+    Type checked_true = this->GetType(op->true_branch);
+    Type checked_false = this->GetType(op->false_branch);
+    return this->Unify(checked_true, checked_false, op->span);
   }
 
-  // We now have a function type.
-  FuncType fn_ty = GetRef<FuncType>(fn_ty_node);
-
-  tvm::Array<Type> ty_args;
-  if (ty_args.size() != 0) {
-    throw Error("found manually suplied type args, not supported");
+  // Handle special case basic primitive operator,
+  // if successful return the return type
+  Type PrimitiveCall(const FuncTypeNode* op,
+                     Array<Type> arg_types,
+                     const Attrs& attrs) {
+    if (op->type_params.size() != arg_types.size() + 1) return Type();
+    if (op->type_constraints.size() != 1) return Type();
+    const TypeRelationNode* rel = op->type_constraints[0].as<TypeRelationNode>();
+    if (rel == nullptr) return Type();
+    // validate if the type parameter matches up
+    for (size_t i = 0; i < op->type_params.size(); ++i) {
+      if (!op->type_params[i].same_as(rel->args[i])) return Type();
+    }
+    Type rtype = IncompleteTypeNode::make(TypeParamNode::Kind::kType);
+    arg_types.push_back(rtype);
+    // we can do simple replacement here
+    solver_.AddConstraint(TypeRelationNode::make(
+        rel->func, arg_types, arg_types.size() - 1, attrs));
+    return rtype;
   }
 
-  fn_ty = Instantiate(fn_ty, ty_args);
+  // instantiate the function type with fresh
+  FuncType Instantiate(const FuncTypeNode* fn_ty, Array<Type>* ty_args) {
+    tvm::Map<TypeParam, Type> subst_map;
 
-  std::vector<Type> arg_types;
-  std::vector<Expr> checked_args;
-
-  for (auto arg : c->args) {
-    auto checked_arg = this->Infer(arg);
-    arg_types.push_back(checked_arg.type);
-    checked_args.push_back(checked_arg.expr);
+    // Build a subsitituion map up from the function type and type arguments.
+    // Eventually allow the type vars to be passed in.
+    for (auto ty_param : fn_ty->type_params) {
+      IncompleteType fresh = IncompleteTypeNode::make(ty_param->kind);
+      subst_map.Set(ty_param, fresh);
+      ty_args->push_back(fresh);
+    }
+    Type ret_type = fn_ty->ret_type;
+
+    // If the function type is incomplete, place a new IncompleteType
+    // This relax the fn_ty to inputs -> Any
+    // The type checking can still pass when there are additional constraints on the type
+    // This is a temporary work around to check recursive functions whose
+    // return type is not yet known.
+    if (!ret_type.defined()) {
+      ret_type = IncompleteTypeNode::make(TypeParamNode::Kind::kType);
+    }
+    Type inst_ty = FuncTypeNode::make(fn_ty->arg_types,
+                                      ret_type, {},
+                                      fn_ty->type_constraints);
+    inst_ty = TypeSubst(inst_ty, subst_map);
+    return Downcast<FuncType>(inst_ty);
   }
 
-  auto type_arity = fn_ty->arg_types.size();
-  auto number_of_args = arg_types.size();
+  // Handle general call node.
+  Type GeneralCall(const CallNode* op, Array<Type> arg_types) {
+    Type ftype = GetType(op->op);
+    auto* fn_ty_node = ftype.as<FuncTypeNode>();
+    CHECK(fn_ty_node != nullptr)
+        << "only expressions with function types can be called, at "
+        << op->span;
+
+    Array<Type> type_args;
+    FuncType fn_ty = Instantiate(fn_ty_node, &type_args);
+    size_t type_arity = fn_ty->arg_types.size();
+    size_t number_of_args = arg_types.size();
+
+    if (type_arity != number_of_args) {
+      if (type_arity < number_of_args) {
+        LOG(FATAL) << "the function is provided too many arguments " << op->span;
+      } else {
+        LOG(FATAL) << "the function is provided too few arguments" << op->span;
+      }
+    }
+    for (size_t i = 0; i < fn_ty->arg_types.size(); i++) {
+      this->Unify(fn_ty->arg_types[i], arg_types[i], op->args[i]->span);
+    }
 
-  if (type_arity != number_of_args) {
-    if (type_arity < number_of_args) {
-      this->FatalError("the function is provided too many arguments", c->span);
-    } else {
-      this->FatalError("the function is provided too few arguments", c->span);
+    for (auto cs : fn_ty->type_constraints) {
+      solver_.AddConstraint(cs);
     }
+    return fn_ty->ret_type;
   }
 
-  for (size_t i = 0; i < fn_ty->arg_types.size(); i++) {
-    this->Unify(fn_ty->arg_types[i], arg_types[i], c->args[i]->span);
+  Type VisitExpr_(const CallNode* op) final {
+    // Fast path: well-formed primitive op
+    Array<Type> arg_types;
+    for (Expr arg : op->args) {
+      arg_types.push_back(GetType(arg));
+    }
+    if (const OpNode* opnode = op->op.as<OpNode>()) {
+      Type rtype = PrimitiveCall(opnode->op_type.as<FuncTypeNode>(),
+                                 arg_types,
+                                 op->attrs);
+      if (rtype.defined()) return rtype;
+    }
+    return GeneralCall(op, arg_types);
   }
 
-  // After we unify the arguments we should know more about the type
-  // arguments, let's run a quick pass over them to find new
-  // representatives.
-
-  for (size_t i = 0; i < ty_args.size(); i++) {
-    ty_args.Set(i, this->unifier->Subst(ty_args[i]));
+  Type VisitExpr_(const FunctionNode* f) final {
+    for (auto param : f->params) {
+      type_map_[param->var] = param->type;
+      type_map_[param] = param->type;
+    }
+    Type rtype = GetType(f->body);
+    // Run solver using the currently known information
+    solver_.Solve();
+    // Trying to resolve
+    Array<Type> arg_types;
+    for (size_t i = 0; i < f->params.size(); ++i) {
+      Param param = f->params[i];
+      Type atype = solver_.Resolve(param->type);
+      CHECK(atype.as<IncompleteTypeNode>() == nullptr)
+          << "Cannot resolve type of " << i
+          << "-th parameter of function at" << f->span;
+      arg_types.push_back(atype);
+    }
+    rtype = solver_.Resolve(rtype);
+    CHECK(rtype.as<IncompleteTypeNode>() == nullptr)
+        << "Cannot resolve return type of function at" << f->span;
+    // do not support constraint lifting for now.
+    return FuncTypeNode::make(arg_types, rtype, f->type_params, {});
   }
+};
 
-  // Add type constraints from the function types.
-  for (auto cs : fn_ty->type_constraints) {
-    context.AddConstraint(cs);
+class TypeInferencer::Resolver : public ExprMutator {
+ public:
+  Resolver(const std::unordered_map<Expr, Type, NodeHash, NodeEqual>& tmap,
+           TypeSolver* solver)
+      : tmap_(tmap), solver_(solver) {
   }
 
-  auto new_call =
-      CallNode::make(checked_op.expr, checked_args, c->attrs, ty_args);
-
-  return {new_call, fn_ty->ret_type};
-}
-
-CheckedExpr TypeInferencer::VisitExpr_(const LetNode* op) {
-  Let let = GetRef<Let>(op);
-
-  CheckedExpr checked_value;
-  Type annotated_ty = Resolve(let->value_type);
-
-  // If we are let-defining a function, we want to be able to
-  // recursively name the function in order to support recursive
-  // local definitions.
-  if (let->value.as<FunctionNode>()) {
-    context.Insert(let->var, annotated_ty);
-    checked_value = Infer(let->value);
-  } else {
-    checked_value = Infer(let->value);
+  Expr VisitExpr_(const VarNode* op) final {
+    return AttachCheckedType(op);
   }
 
-  Type unified_ty = this->Unify(checked_value.type, annotated_ty, let->span);
-
-  // Update type context with unified type now that we have
-  // solved this equation.
-  context.Insert(let->var, unified_ty);
-
-  auto checked_body = Infer(let->body);
-
-  auto checked_let = LetNode::make(let->var, checked_value.expr,
-                                   checked_body.expr, let->value_type);
-
-  return {checked_let, checked_body.type};
-}
-
-CheckedExpr TypeInferencer::VisitExpr_(const IfNode* op) {
-  If ifn = GetRef<If>(op);
-
-  // Ensure the type of the guard is of Tensor[Bool, ()],
-  // that is a rank-0 boolean tensor.
-  auto checked_cond = this->Infer(ifn->cond);
-  auto cond_type = checked_cond.type;
-
-  this->Unify(cond_type, TensorTypeNode::make({}, HalideIR::Bool()),
-              ifn->cond->span);
-  auto checked_true = this->Infer(ifn->true_branch);
-  auto checked_false = this->Infer(ifn->false_branch);
-  auto unified_type =
-      this->Unify(checked_true.type, checked_false.type, ifn->span);
-  auto checked_if =
-      IfNode::make(checked_cond.expr, checked_true.expr, checked_false.expr);
-  return {checked_if, unified_type};
-}
-
-CheckedExpr TypeInferencer::VisitExpr_(const OpNode* op_node) {
-  auto op = GetRef<Op>(op_node);
-  return {op, op->op_type};
-}
-
-Type TypeInferencer::Resolve(const Type &t) {
-  if (t.defined()) {
-    return ::tvm::relay::Resolve(this->unifier, t);
-  } else {
-    return IncompleteTypeNode::make(TypeParamNode::Kind::kType);
+  Expr VisitExpr_(const ConstantNode* op) final {
+    return AttachCheckedType(op);
   }
-}
 
-Expr TypeInferencer::Resolve(const Expr &e) {
-  CHECK(e.defined());
-  return ::tvm::relay::Resolve(this->unifier, e);
-}
+  Expr VisitExpr_(const GlobalVarNode* op) final {
+    return AttachCheckedType(op);
+  }
 
-void TypeInferencer::Solve(TypeRelationData & ty_rel) {
-  Array<Type> normalized_args;
+  Expr VisitExpr_(const OpNode* op) final {
+    return ExprMutator::VisitExpr_(op);
+  }
 
-  for (auto arg : ty_rel.args) {
-    normalized_args.push_back(Resolve(arg));
+  Expr VisitExpr_(const TupleNode* op) final {
+    return AttachCheckedType(op);
   }
 
-  auto new_args = ty_rel.func(normalized_args, ty_rel.args.size());
+  Expr VisitExpr_(const ParamNode* op) final {
+    return ExprMutator::VisitExpr_(op);
+  }
 
-  CHECK(new_args.size() == normalized_args.size());
-  tvm::Array<Type> final_args;
+  Expr VisitExpr_(const FunctionNode* op) final {
+    return AttachCheckedType(op);
+  }
 
-  for (size_t i = 0; i < new_args.size(); i++) {
-    ty_rel.args[i] = Unify(normalized_args[i], new_args[i], ty_rel.span);
+  Expr VisitExpr_(const CallNode* op) final {
+    return AttachCheckedType(op);
   }
-}
 
-int NumSolvedVars(const Array<Type>& vars) {
-  int num = 0;
-  for (auto var : vars) {
-    if (!var.as<IncompleteTypeNode>()) {
-      num += 1;
-    }
+  Expr VisitExpr_(const LetNode* op) final {
+    return AttachCheckedType(op);
   }
-  return num;
-}
 
-enum SolverResult : int {
-  Failed = -1,
-  Progress = 0,
-  Done = 1,
-};
+  Expr VisitExpr_(const IfNode* op) final {
+    return AttachCheckedType(op);
+  }
 
-SolverResult TypeInferencer::Solve(std::vector<TypeRelationData>& rels) {
-  // We start in the done state with zero progress.
-  SolverResult status = SolverResult::Done;
-  int progress = 0;
-
-  do {
-    // Upon rentering the loop we reset the state.
-    status = SolverResult::Done;
-    progress = 0;
-
-    std::vector<int> complete;
-
-    int i = 0;
-    // We will now process each relation in order.
-    for (TypeRelationData& ty_rel : rels) {
-      int arity = ty_rel.args.size();
-      int pre_solved = NumSolvedVars(ty_rel.args);
-      RELAY_LOG(INFO) << "TypeInferencer::Solve: "
-                      << "TypeRelation= "
-                      << ", Arity=" << arity << ", Solved=" << pre_solved
-                      << std::endl;
-      // If the relation is already solved then we will make no progress but try
-      // to set the status to done.
-      if (pre_solved == arity) {
-        status = static_cast<SolverResult>((status && SolverResult::Done));
-        complete.push_back(i);
-      // If there are unsolved variables we will try to solve some.
-      } else if (pre_solved < arity) {
-        Solve(ty_rel);
-        int post_solved = NumSolvedVars(ty_rel.args);
-
-        // If we solved any variables we will try to downgrade status to
-        // progress update the type relation, and then bump the progress counter
-        // by one.
-        if (post_solved > pre_solved) {
-          status =
-              static_cast<SolverResult>((status && SolverResult::Progress));
-          progress += 1;
-        }
+  // attach checked type to the mutated node.
+  template<typename T>
+  Expr AttachCheckedType(const T* op) {
+    auto it = tmap_.find(GetRef<Expr>(op));
+    CHECK(it != tmap_.end());
+    Type checked_type = solver_->Resolve(it->second);
+    CHECK(checked_type.as<IncompleteTypeNode>() == nullptr)
+        << "Cannot resolve type of " << GetRef<Expr>(op)
+        << " at " << op->span;
+    Expr new_e = ExprMutator::VisitExpr_(op);
+    if (!checked_type.same_as(new_e->checked_type_)) {
+      // Copy on write optimization
+      // If new_e is an old expression,
+      // we make a copy mutating an existing reference.
+      if (!new_e.node_.unique()) {
+        new_e = Expr(make_node<T>(*new_e.as<T>()));
       }
-      i++;
+      new_e->checked_type_ = checked_type;
     }
-
-    // If we made no progress and we aren't finished, then the state should be
-    // downgraded to fail, then we should exit the loop.
-    if (progress == 0 && status != SolverResult::Done) {
-      status = SolverResult::Failed;
-      break;
-    }
-
-    // Remove the satisfied relations.
-    for (auto i : complete) {
-      if (rels.size() > 1) {
-        rels[i] = rels.back();
-        rels.pop_back();
-      } else {
-        rels.pop_back();
-      }
-    }
-
-    std::reverse(rels.begin(), rels.end());
-  } while (status == SolverResult::Progress);
-  return status;
-}
-
-bool TypeInferencer::RelationsHold(bool scope_only) {
-  // If we are only checking the top scope,
-  // slice out the constraints.
-  //
-  // Otherwise we use all of them.
-  std::vector<std::vector<TypeRelationData> > constraints;
-
-  if (scope_only) {
-    constraints = {context.constraints[0]};
-  } else {
-    constraints = context.constraints;
+    return new_e;
   }
 
-  RELAY_LOG(INFO) << "TypeInferencer::RelationsHold: scope_only= " << scope_only
-                  << std::endl;
-  bool all_hold = true;
-  for (auto ty_rels : context.constraints) {
-    auto status = Solve(ty_rels);
-    RELAY_LOG(INFO) << "status= " << status << std::endl;
-    if (status == SolverResult::Failed || status == SolverResult::Progress) {
-      all_hold = false;
-    } else if (status == SolverResult::Done) {
-      continue;
-    } else {
-      throw InternalError("found invalid value for SolverResult");
-    }
+  Type VisitType(const Type &t) final {
+    return solver_->Resolve(t);
   }
 
-  return all_hold;
+ private:
+  const std::unordered_map<Expr, Type, NodeHash, NodeEqual>& tmap_;
+  TypeSolver* solver_;
+};
+
+
+Expr TypeInferencer::Infer(Expr expr) {
+  // step 0: populate the constraints
+  GetType(expr);
+  // step 1: solve the constraints
+  solver_.Solve();
+  // step 2: attach resolved types to checked_type field
+  return Resolver(type_map_, &solver_).VisitExpr(expr);
 }
 
 Expr InferType(const Environment& env, const Expr& e) {
-  TypeInferencer ti(env);
-  auto checked_expr = ti.Infer(e);
-  CHECK(ti.RelationsHold());
-  return ti.Resolve(checked_expr.expr);
+  return TypeInferencer(env).Infer(e);
 }
 
-Expr InferType(const Environment& env, const GlobalVar& var,
+Expr InferType(const Environment& env,
+               const GlobalVar& var,
                const Function& func) {
-  TypeInferencer ti(env);
-  auto func_copy = FunctionNode::make(func->params, func->ret_type, func->body,
-                                      func->type_params);
-  func_copy->checked_type_ = ti.Resolve(func_copy->fn_type());
+  Function func_copy = Function(make_node<FunctionNode>(*func.operator->()));
+  func_copy->checked_type_ = func_copy->fn_type();
   env->functions.Set(var, func_copy);
-  auto checked_expr = ti.Infer(func);
-  CHECK(ti.RelationsHold());
+  Expr func_ret = TypeInferencer(env).Infer(func_copy);
   auto map_node = env->functions.CopyOnWrite();
   map_node->data.erase(var.node_);
-  return ti.Resolve(checked_expr.expr);
-}
-
-void TypeInferencer::FatalError(const std::string& msg, Span sp) {
-  throw FatalTypeError(
-      "internal error: this exception should"
-      "be handled and errors reported with Environment::display_errors\n" +
-      msg);
-}
-
-Type TypeInferencer::Unify(const Type& t1, const Type& t2, Span sp) {
-  try {
-    return this->unifier->Unify(t1, t2);
-  } catch (const dmlc::Error &e) {
-    std::stringstream ss;
-    ss << "Error unifying `";
-    ss << t1;
-    ss << "` and `";
-    ss << t2;
-    ss << "`: " << e.what();
-    this->FatalError(ss.str(), sp);
-  }
+  return func_ret;
 }
 
 TVM_REGISTER_API("relay._ir_pass.check_expr")
-    .set_body([](TVMArgs args, TVMRetValue* ret) {
-      Environment env = args[0];
-      Expr e = args[1];
-      *ret = InferType(env, e);
-    });
-
-// TODO(@jroesch): put in a better namespace.
-TVM_REGISTER_API("relay._ir_pass._get_checked_type")
-    .set_body([](TVMArgs args, TVMRetValue* ret) {
-      Expr e = args[0];
-      *ret = e->checked_type();
-    });
-
-/* Incomplete Type */
-
-IncompleteType IncompleteTypeNode::make(TypeParamNode::Kind kind) {
-  auto n = make_node<IncompleteTypeNode>();
-  n->kind = std::move(kind);
-  return IncompleteType(n);
-}
-
-TVM_REGISTER_API("relay._make.IncompleteType")
-    .set_body([](TVMArgs args, TVMRetValue* ret) {
-      int kind = args[0];
-      *ret = IncompleteTypeNode::make(static_cast<TypeParamNode::Kind>(kind));
-    });
-
-TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable)
-    .set_dispatch<IncompleteTypeNode>([](const IncompleteTypeNode* node,
-                                         tvm::IRPrinter* p) {
-      p->stream << "IncompleteTypeNode(" << node->kind << ", " << node << ")";
-    });
+.set_body([](TVMArgs args, TVMRetValue* ret) {
+    Environment env = args[0];
+    Expr e = args[1];
+    *ret = InferType(env, e);
+  });
 
 }  // namespace relay
 }  // namespace tvm

src/relay/pass/type_solver.cc

+/*!
+ *  Copyright (c) 2018 by Contributors
+ * \file type_solver.cc
+ * \brief Type solver implementations.
+ */
+#include <string>
+#include "type_solver.h"
+
+namespace tvm {
+namespace relay {
+
+class TypeSolver::Reporter : public TypeReporterNode {
+ public:
+  explicit Reporter(TypeSolver* solver)
+      : solver_(solver) {}
+
+  void Assign(const Type& dst, const Type& src) final {
+    solver_->Unify(dst, src);
+  }
+
+  void AssertEQ(const ShapeExpr& lhs, const ShapeExpr& rhs) final {
+    // TODO(tqchen)
+  }
+
+ private:
+  TypeSolver* solver_;
+};
+
+// constructor
+TypeSolver::TypeSolver()
+    : reporter_(make_node<Reporter>(this)) {
+}
+
+// destructor
+TypeSolver::~TypeSolver() {
+  // call destructor of all non-POD arena object
+  for (TypeNode* ptr : type_nodes_) {
+    ptr->~TypeNode();
+  }
+  for (RelationNode* ptr : rel_nodes_) {
+    ptr->~RelationNode();
+  }
+}
+
+// Add equality constraint
+Type TypeSolver::Unify(const Type& dst, const Type& src) {
+  // Known limitation
+  // - handle composite types whose component can be unknown.
+  // - handle shape pattern matching
+  TypeNode* lhs = GetTypeNode(dst);
+  TypeNode* rhs = GetTypeNode(src);
+  if (lhs->resolved_type.as<IncompleteTypeNode>()) {
+    MergeFromTo(lhs, rhs);
+    return rhs->resolved_type;
+  } else if (rhs->resolved_type.as<IncompleteTypeNode>()) {
+    MergeFromTo(rhs, lhs);
+    return lhs->resolved_type;
+  } else {
+    lhs->parent = rhs;
+    CHECK(AlphaEqual(lhs->resolved_type, rhs->resolved_type))
+        << "Incompatible parent types in UF:"
+        << lhs->resolved_type << " and " << rhs->resolved_type;
+    return rhs->resolved_type;
+  }
+}
+
+// Add type constraint to the solver.
+void TypeSolver::AddConstraint(const TypeConstraint& constraint) {
+  if (auto *op = constraint.as<TypeRelationNode>()) {
+    // create a new relation node.
+    RelationNode* rnode = make<RelationNode>();
+    rnode->rel = GetRef<TypeRelation>(op);
+    rel_nodes_.push_back(rnode);
+    // populate the type information.
+    for (size_t i = 0; i < op->args.size(); ++i) {
+      // insert link to the type list
+      LinkNode<TypeNode*>* tlink = make<LinkNode<TypeNode*> >();
+      TypeNode* tnode = GetTypeNode(op->args[i]);
+      tlink->value = tnode;
+      rnode->type_list.Push(tlink);
+      // insert type->relation node
+      LinkNode<RelationNode*>* rlink = make<LinkNode<RelationNode*> >();
+      rlink->value = rnode;
+      tnode->rel_list.Push(rlink);
+    }
+    // add the relation to the working queue.
+    this->AddToQueue(rnode);
+  } else {
+    LOG(FATAL) << "Do not know how to handle constraint type"
+               << constraint->type_key();
+  }
+}
+
+// Resolve a type in the solver context.
+Type TypeSolver::Resolve(const Type& type) {
+  auto it = tmap_.find(type);
+  if (it != tmap_.end()) {
+    return it->second->FindRoot()->resolved_type;
+  } else {
+    return type;
+  }
+}
+
+bool TypeSolver::Solve() {
+  // update until queue is empty
+  while (!update_queue_.empty()) {
+    RelationNode* rnode = update_queue_.front();
+    const auto& rel = rnode->rel;
+    update_queue_.pop();
+    CHECK(!rnode->resolved);
+    // update the relation with given evidence.
+    Array<Type> args;
+    for (auto* tlink = rnode->type_list.head; tlink != nullptr; tlink = tlink->next) {
+      args.push_back(tlink->value->FindRoot()->resolved_type);
+      CHECK_LE(args.size(), rel->args.size());
+    }
+    // call the function
+    bool resolved = rel->func(args, rel->num_inputs, rel->attrs, reporter_);
+    // mark inqueue as false after the function call
+    // so that rnode itself won't get enqueued again.
+    rnode->inqueue = false;
+
+    if (resolved) {
+      ++num_resolved_rels_;
+    }
+    rnode->resolved = resolved;
+  }
+  // This criterion is not necessarily right for all the possible cases
+  // TODO(tqchen): We should also count the number of in-complete types.
+  return num_resolved_rels_ == rel_nodes_.size();
+}
+
+
+// Expose type solver only for debugging purposes.
+TVM_REGISTER_API("relay._ir_pass._test_type_solver")
+.set_body([](runtime::TVMArgs args, runtime::TVMRetValue* ret) {
+    using runtime::PackedFunc;
+    using runtime::TypedPackedFunc;
+    auto solver = std::make_shared<TypeSolver>();
+
+    auto mod = [solver](std::string name) -> PackedFunc {
+      if (name == "Solve") {
+        return TypedPackedFunc<bool()>([solver]() {
+            return solver->Solve();
+          });
+      } else if (name == "Unify") {
+        return TypedPackedFunc<void(Type, Type)>([solver](Type lhs, Type rhs) {
+            solver->Unify(lhs, rhs);
+          });
+      } else if (name == "Resolve") {
+        return TypedPackedFunc<Type(Type)>([solver](Type t) {
+            return solver->Resolve(t);
+          });
+      } else if (name == "AddConstraint") {
+        return TypedPackedFunc<void(TypeConstraint)>([solver](TypeConstraint c) {
+            return solver->AddConstraint(c);
+          });
+      } else {
+        return PackedFunc();
+      }
+    };
+    *ret = runtime::TypedPackedFunc<runtime::PackedFunc(std::string)>(mod);
+  });
+
+}  // namespace relay
+}  // namespace tvm

src/relay/pass/type_solver.h

+/*!
+ *  Copyright (c) 2018 by Contributors
+ * \file type_solver.h
+ * \brief Solver logic for type inference.
+ */
+#ifndef TVM_RELAY_PASS_TYPE_SOLVER_H_
+#define TVM_RELAY_PASS_TYPE_SOLVER_H_
+
+#include <tvm/relay/type.h>
+#include <tvm/relay/pass.h>
+#include <vector>
+#include <queue>
+#include "../../common/arena.h"
+
+
+namespace tvm {
+namespace relay {
+
+/*!
+ * \brief Interface of type solver used in type inference.
+ *
+ * TypeSolver works on a list of constraints among incomplete types.
+ * The user will populate the constraints by AddConstraint and Assign.
+ * Then we can call Solve to trying to resolve the unknown.
+ *
+ * This can be viewed as "type program(computational graph)" of types, where
+ * the type constraint are operators of the graph and the incomplete
+ * types are intermediate value of the graph.
+ * If all the input types are concretely known, we should be able to
+ * just run a forward pass on the "type program" to get all the types.
+ *
+ * The list of constraints representation means we are storing it as a bipartite
+ * graph instead of a DAG. This is because some constraints might go both direction.
+ * TypeSolver could take advantage of bidirectional constraints to deduce input
+ * value given output ones. Never-the-less, we should keep in mind that
+ * there is a "forward direction" that the TypeSolver should take advantage of.
+ */
+class TypeSolver {
+ public:
+  TypeSolver();
+  ~TypeSolver();
+  /*!
+   * \brief Add a type constraint to the solver.
+   * \param constraint The constraint to be added.
+   */
+  void AddConstraint(const TypeConstraint& constraint);
+  /*!
+   * \brief Resolve type to the solution type in the solver.
+   * \param type The type to be resolved.
+   * \return The resolved type.
+   */
+  Type Resolve(const Type& type);
+  /*!
+   * \brief Start to solve the types using the current known information.
+   * \return Whether all the incomplete types has been fully resolved.
+   */
+  bool Solve();
+  /*!
+   * \brief Unify lhs and rhs.
+   * \param lhs The left operand.
+   * \param rhs The right operand
+   */
+  Type Unify(const Type& lhs, const Type& rhs);
+
+ private:
+  class Reporter;
+  struct TypeNode;
+  struct RelationNode;
+  // Internally the solver maintains a bipartite graph of Relation and Types.
+  // All the object in the structure is managed by a arena allocator
+  // which releases the memory upon distruction of the type solver.
+  /*!
+   * \brief Link list node
+   * \tparam T the content data type
+   */
+  template<typename T>
+  struct LinkNode {
+    /*! \brief The content value */
+    T value;
+    /*! \brief pointer to the next location */
+    LinkNode<T>* next{nullptr};
+  };
+  /*!
+   * \brief LinkedList structure
+   * \tparam T the content data type
+   */
+  template<typename T>
+  struct LinkedList {
+    /*! \brief Head pointer */
+    LinkNode<T>* head{nullptr};
+    /*! \brief Tail pointer */
+    LinkNode<T>* tail{nullptr};
+    /*!
+     * \brief Push a new node to the end of the linked list.
+     * \param node The node to be pushed.
+     */
+    void Push(LinkNode<T>* node) {
+      node->next = nullptr;
+      if (this->tail != nullptr) {
+        this->tail->next = node;
+        this->tail = node;
+      } else {
+        head = tail = node;
+      }
+    }
+  };
+  /*!
+   * \brief type node struct
+   *  TypeNode implements a union-find data structure(via parent)
+   *  that can unifies the same types to the name resolved_type.
+   *
+   *  It also contains collection of links to related Relations,
+   *  which is stored in rel_list.
+   */
+  struct TypeNode {
+    /*! \brief The final resolved type */
+    Type resolved_type;
+    /*! \brief type node in the union find algorithm */
+    TypeNode* parent{nullptr};
+    /*! \brief list of relations that is related to this type node */
+    LinkedList<RelationNode*> rel_list;
+    /*!
+     * \brief Find the root type node, perform path compression
+     * \return The root type node.
+     */
+    TypeNode* FindRoot() {
+      // fast path
+      if (this->parent == nullptr) return this;
+      // slow path with path compression.
+      TypeNode* root = this;
+      while (root->parent != nullptr) {
+        root = root->parent;
+      }
+      for (TypeNode* p = this; p != root;) {
+        TypeNode* parent = p->parent;
+        p->parent = root;
+        p = parent;
+      }
+      return root;
+    }
+  };
+  /*! \brief relation node */
+  struct RelationNode {
+    /*! \brief Whether the relation is in the queue to be solved */
+    bool inqueue{false};
+    /*! \brief Whether the relation is resolved */
+    bool resolved{false};
+    /*! \brief The corresponding type relation */
+    TypeRelation rel;
+    /*! \brief list types to this relation */
+    LinkedList<TypeNode*> type_list;
+  };
+  /*! \brief List of all allocated type nodes */
+  std::vector<TypeNode*> type_nodes_;
+  /*! \brief List of all allocated relation nodes */
+  std::vector<RelationNode*> rel_nodes_;
+  /*! \brief Number of resolved relations */
+  size_t num_resolved_rels_{0};
+  /*! \brief map from type node to types. */
+  std::unordered_map<Type, TypeNode*, NodeHash, NodeEqual> tmap_;
+  /*! \breif Internal queue to update the relation */
+  std::queue<RelationNode*> update_queue_;
+  /*! \brief allocator of all the internal node obhect*/
+  common::Arena arena_;
+  /*! \brief Reporter that reports back to self */
+  TypeReporter reporter_;
+  /*!
+   * \brief Create function to create a new node ptr via arena
+   * \tparam The type parameter
+   * \return The node pointer.
+   */
+  template<typename T>
+  T* make() {
+    T* ptr = arena_.Alloc<T>();
+    // call constructor
+    new (ptr) T();
+    return ptr;
+  }
+  /*!
+   * \brief GetTypeNode that is corresponds to t.
+   * if it do not exist, create a new one.
+   * \return The type node.
+   */
+  TypeNode* GetTypeNode(const Type& t) {
+    auto it = tmap_.find(t);
+    if (it != tmap_.end()) {
+      return it->second->FindRoot();
+    } else {
+      TypeNode* n = make<TypeNode>();
+      type_nodes_.push_back(n);
+      n->resolved_type = t;
+      tmap_[t] = n;
+      return n;
+    }
+  }
+  /*!
+   * \brief Add relation node rel to the update queue
+   * \param rel The relation node
+   */
+  void AddToQueue(RelationNode* rel) {
+    if (rel->inqueue) return;
+    CHECK(!rel->resolved);
+    rel->inqueue = true;
+    update_queue_.push(rel);
+  }
+  /*!
+   * \brief Merge rhs type node to lhs
+   * \param src The source operand
+   * \param dst The dst operand.
+   */
+  void MergeFromTo(TypeNode* src, TypeNode* dst) {
+    if (src == dst) return;
+    src->parent = dst;
+    // move the link to the to dst
+    for (auto* rlink = src->rel_list.head; rlink != nullptr;) {
+      // store next pointer first before rlink get moved
+      auto* next = rlink->next;
+      // if the relation is not yet resolved
+      // send the relation to the new
+      if (!rlink->value->resolved) {
+        this->AddToQueue(rlink->value);
+        dst->rel_list.Push(rlink);
+      }
+      rlink = next;
+    }
+  }
+};
+
+}  // namespace relay
+}  // namespace tvm
+#endif  // TVM_RELAY_PASS_TYPE_SOLVER_H_

src/relay/pass/type_visitor.h

@@ -108,11 +108,14 @@ struct TypeMutator : TypeFunctor<Type(const Type& n)> {
     for (const Type& t : type_rel->args) {
       new_args.push_back(this->VisitType(t));
     }
-    return TypeRelationNode::make(type_rel->name, type_rel->func_, new_args);
+    return TypeRelationNode::make(type_rel->func,
+                                  new_args,
+                                  type_rel->num_inputs,
+                                  type_rel->attrs);
   }
 
   Type VisitType_(const IncompleteTypeNode* op) override {
-    return GetRef<IncompleteType>(op);
+    return GetRef<Type>(op);
   }
 };
 

src/relay/pass/unifier.cc

-/*!
- *  Copyright (c) 2018 by Contributors
- * \file tvm/src/relay/pass/unifier.cc
- * \brief The type unifier which solves a system of equations between
- * incomplete types.
- */
-
-#include "./unifier.h"
-#include <tvm/relay/error.h>
-#include <tvm/relay/expr.h>
-#include <tvm/relay/logging.h>
-#include <tvm/relay/pass.h>
-#include <tvm/relay/type.h>
-#include "./type_subst.h"
-#include "./type_visitor.h"
-
-namespace tvm {
-namespace relay {
-
-using tvm::IRPrinter;
-using namespace tvm::runtime;
-
-UnionFind UnionFindNode::make(tvm::Map<IncompleteType, Type> uf_map) {
-  auto n = make_node<UnionFindNode>();
-  n->uf_map = uf_map;
-  return UnionFind(n);
-}
-
-void UnionFindNode::Insert(const IncompleteType& v) { this->uf_map.Set(v, v); }
-
-void UnionFindNode::debug() {
-  for (const auto& entry : this->uf_map) {
-    RELAY_LOG(INFO) << entry.first << " = " << entry.second << std::endl;
-  }
-}
-
-void UnionFindNode::AssertAlphaEqual(const Type& l, const Type& r) {
-  if (!AlphaEqual(l, r)) {
-    std::stringstream ss;
-    ss << "Incompatible parent types in UF:" << l << " and " << r;
-    throw UnionFindError(ss.str());
-  }
-}
-
-void UnionFindNode::Unify(const IncompleteType& v1, const Type& t) {
-  RELAY_LOG(INFO) << "UnionFindNode::Unify v1=" << v1 << ", t=" << t
-                  << std::endl;
-  auto parent1 = this->Find(v1);
-
-  // if t is a type var, then unify parents
-  const IncompleteTypeNode *tvn2 = t.as<IncompleteTypeNode>();
-  if (tvn2) {
-    auto v2 = GetRef<IncompleteType>(tvn2);
-    auto parent2 = this->Find(v2);
-
-    // if parents are exactly equal, then we're done
-    if (parent1 == parent2) {
-      return;
-    }
-
-    // if first parent is a type var, then can just set its union find map to
-    // second parent
-    if (const IncompleteTypeNode *pvn1 = parent1.as<IncompleteTypeNode>()) {
-      auto pv1 = GetRef<IncompleteType>(pvn1);
-      this->uf_map.Set(pv1, parent2);
-      return;
-    }
-
-    // if second parent is a type var but first isn't, can set second type var
-    if (const IncompleteTypeNode *pvn2 = parent2.as<IncompleteTypeNode>()) {
-      auto pv2 = GetRef<IncompleteType>(pvn2);
-      this->uf_map.Set(pv2, parent1);
-      return;
-    }
-
-    // if both parents are not type vars themselves, check alpha-equality
-    AssertAlphaEqual(parent1, parent2);
-    return;
-  }
-
-  // if t is not a type var, then unify with v1's parent if parent is a type
-  // var; else, check alpha-equality for compatibility
-  if (const IncompleteTypeNode *pvn1 = parent1.as<IncompleteTypeNode>()) {
-    auto pv1 = GetRef<IncompleteType>(pvn1);
-    this->uf_map.Set(pv1, t);
-    return;
-  }
-
-  AssertAlphaEqual(parent1, t);
-}
-
-Type UnionFindNode::Find(const IncompleteType& v) {
-  // The node has no mapping, so its representative is just itself.
-  if (this->uf_map.find(v) == this->uf_map.end()) {
-    return v;
-  }
-
-  Type parent = this->uf_map.at(v);
-
-  if (v == parent) {
-    return v;
-  }
-
-  // if parent is not a type var, then it must be the representative type
-  const IncompleteTypeNode *rep = parent.as<IncompleteTypeNode>();
-  if (!rep) {
-    return parent;
-  }
-
-  // otherwise, recurse and perform path compression
-  IncompleteType pv = GetRef<IncompleteType>(rep);
-  Type higher_up = this->Find(pv);
-  this->uf_map.Set(v, higher_up);
-  return higher_up;
-}
-
-TVM_REGISTER_API("relay._make.UnionFind")
-    .set_body([](TVMArgs args, TVMRetValue *ret) {
-      if (args.size() == 0) {
-        *ret = UnionFindNode::make({});
-      } else {
-        *ret = UnionFindNode::make(args[0]);
-      }
-    });
-
-TVM_STATIC_IR_FUNCTOR_REGISTER(IRPrinter, vtable)
-    .set_dispatch<UnionFindNode>([](const UnionFindNode *node,
-                                    tvm::IRPrinter *p) {
-      p->stream << "UnionFindNode(" << node->uf_map << ")";
-    });
-
-TypeUnifier TypeUnifierNode::make(UnionFind union_find) {
-  auto n = make_node<TypeUnifierNode>();
-  n->union_find = union_find;
-  return TypeUnifier(n);
-}
-
-void TypeUnifierNode::Insert(const IncompleteType& v) {
-  this->union_find->Insert(v);
-}
-
-Type TypeUnifierNode::Unify(const Type& t1, const Type& t2) {
-  RELAY_LOG(INFO) << "TypeUnifierNode::unify: t1=" << t1 << " t2=" << t2
-                  << std::endl;
-
-  Type unified = this->VisitType(t1, t2);
-  // TODO(@jroesch): Restore this code when we finish kind checker.
-  // if (!check_kind(unified)) {
-  // throw UnificationError("Invalid kinds in unified type");
-  // }
-  return unified;
-}
-
-struct IncompleteTypeSubst : TypeMutator {
-  const TypeUnifierNode *unifier;
-
-  IncompleteTypeSubst(const TypeUnifierNode *unifier) : unifier(unifier) {}
-
-  // type var: look it up in the type map and recurse
-  Type VisitType_(const IncompleteTypeNode* op) override {
-    auto tv = GetRef<IncompleteType>(op);
-    auto parent = unifier->union_find->Find(tv);
-    if (parent == tv) {
-      return tv;
-    }
-    return this->VisitType(parent);
-  }
-};
-
-Type TypeUnifierNode::Subst(const Type& t) {
-  IncompleteTypeSubst tvsubst(this);
-  // normalize first so substitutions in quantifiers will be correct
-  Type ret = tvsubst.VisitType(t);
-  // TODO(@jroesch): Restore this code when we finish kind checker.
-  // if (!check_kind(ret)) {
-  // std::stringstream ss;
-  // ss << "Invalid Kinds in substituted type!";
-  // ss << t << std::endl;
-  // ss << ret << std::endl;
-  // throw SubstitutionError(ss.str());
-  // }
-  return ret;
-}
-
-Type TypeUnifierNode::VisitType(const Type& t1, const Type t2) {
-  // When the right hand size is a type variable immediately unify.
-  if (const IncompleteTypeNode *tvn2 = t2.as<IncompleteTypeNode>()) {
-    return this->UnifyWithIncompleteType(t1, GetRef<IncompleteType>(tvn2));
-  } else {
-    return TypeFunctor<Type(const Type &t1, const Type t2)>::VisitType(t1, t2);
-  }
-}
-
-Type TypeUnifierNode::UnifyWithIncompleteType(const Type& t1,
-                                              const IncompleteType tv2) {
-  RELAY_LOG(INFO) << "unifyWithIncompleteType: t1=" << t1 << " t2=" << tv2
-                  << std::endl;
-  // Fix unify to return new representative
-  this->union_find->Unify(tv2, t1);
-  auto rep = this->union_find->Find(tv2);
-  RELAY_LOG(INFO) << "unifyWithIncompleteType: rep =" << rep << std::endl;
-  return rep;
-}
-
-Type TypeUnifierNode::VisitType_(const IncompleteTypeNode* t1, const Type rt2) {
-  IncompleteType tv1 = GetRef<IncompleteType>(t1);
-  RELAY_LOG(INFO) << "VisitType_: IncompleteTypeNode t1=" << t1 << " = " << rt2
-                  << std::endl;
-  this->union_find->Unify(tv1, rt2);
-  auto rep = this->union_find->Find(tv1);
-  RELAY_LOG(INFO) << "VisitType_: IncompleteTypeNode rep=" << rep << std::endl;
-  return rep;
-}
-
-Type TypeUnifierNode::VisitType_(const TypeParamNode* t1, const Type rt2) {
-  TypeParam ti1 = GetRef<TypeParam>(t1);
-
-  if (const TypeParamNode *tin2 = rt2.as<TypeParamNode>()) {
-    TypeParam ti2 = GetRef<TypeParam>(tin2);
-
-    if (ti1 != ti2) {
-      throw UnificationError("Attempting to unify non-matching TypeParams");
-    }
-
-    return ti1;
-  }
-
-  throw UnificationError("Unable to unify TypeParamNode");
-}
-
-Type TypeUnifierNode::VisitType_(const FuncTypeNode* t1, const Type rt2) {
-  FuncType ft1 = GetRef<FuncType>(t1);
-
-  if (const FuncTypeNode *tan2 = rt2.as<FuncTypeNode>()) {
-    FuncType ft2 = GetRef<FuncType>(tan2);
-
-    if (ft1->type_params.size() != ft2->type_params.size()) {
-      throw UnificationError(
-          "unable to unify functions with differing number of type parameters");
-    }
-
-    tvm::Map<TypeParam, Type> subst_map;
-
-    for (size_t i = 0; i < ft1->arg_types.size(); i++) {
-      subst_map.Set(ft1->type_params[i], ft2->type_params[i]);
-    }
-
-    ft1 = Downcast<FuncType>(TypeSubst(ft1, subst_map));
-
-    if (ft1->arg_types.size() != ft2->arg_types.size()) {
-      throw UnificationError("unable to unify functions of different arities");
-    }
-
-    tvm::Array<Type> unified_args;
-    for (size_t i = 0; i < ft1->arg_types.size(); i++) {
-      unified_args.push_back(
-          this->VisitType(ft1->arg_types[i], ft2->arg_types[i]));
-    }
-
-    Type unified_ret_type = this->VisitType(ft1->ret_type, ft2->ret_type);
-
-    return FuncTypeNode::make(unified_args, unified_ret_type, {}, {});
-  }
-
-  throw UnificationError("unable to unify function types");
-}
-
-Type TypeUnifierNode::VisitType_(const TensorTypeNode* t1, const Type rt2) {
-  TensorType tt1 = GetRef<TensorType>(t1);
-
-  if (const TensorTypeNode *ttn2 = rt2.as<TensorTypeNode>()) {
-    TensorType tt2 = GetRef<TensorType>(ttn2);
-
-    if (!AlphaEqual(tt1, tt2)) {
-      throw UnificationError("dtypes do not match");
-    }
-
-    RELAY_LOG(INFO) << "Unify Tensor Shape s1=" << tt1->shape
-                    << " s2= " << tt2->shape << std::endl;
-
-    if (tt1->shape.size() != tt2->shape.size()) {
-      throw UnificationError("shapes are not of the same length");
-    }
-
-    for (size_t i = 0U; i < tt1->shape.size(); i++) {
-      if (!tt1->shape[i].same_as(tt2->shape[i])) {
-        throw UnificationError("shapes do not match at index");
-      }
-    }
-
-    return rt2;
-  }
-
-  throw UnificationError("Cannot unify TensorTypeNode");
-}
-
-Type TypeUnifierNode::VisitType_(const TupleTypeNode* t1, const Type rt2) {
-  TupleType pt1 = GetRef<TupleType>(t1);
-
-  if (const TupleTypeNode *ptn2 = rt2.as<TupleTypeNode>()) {
-    TupleType pt2 = GetRef<TupleType>(ptn2);
-
-    std::vector<Type> unified_fields;
-    if (pt1->fields.size() != pt2->fields.size()) {
-      throw UnificationError("Product types are of different dimensions");
-    }
-
-    for (size_t i = 0U; i < pt1->fields.size(); i++) {
-      Type unified = this->VisitType(pt1->fields[i], pt2->fields[i]);
-      unified_fields.push_back(unified);
-    }
-
-    return TupleTypeNode::make(unified_fields);
-  }
-
-  throw UnificationError("Cannot unify TupleTypeNode");
-}
-
-Type TypeUnifierNode::VisitType_(const TypeRelationNode* tr1, const Type t2) {
-  throw InternalError("Cannot unify different type relations");
-}
-
-}  // namespace relay
-}  // namespace tvm

src/relay/pass/unifier.h

-/*!
- *  Copyright (c) 2018 by Contributors
- * \file include/tvm/relay/pass/unifier.h
- * \brief The type unifier which solves a system of equations between
- * incomplete types.
- */
-#ifndef TVM_RELAY_PASS_UNIFIER_H_
-#define TVM_RELAY_PASS_UNIFIER_H_
-
-#include <tvm/relay/expr.h>
-#include <string>
-#include "./type_functor.h"
-
-namespace tvm {
-namespace relay {
-
-struct UnionFindError : dmlc::Error {
-  explicit UnionFindError(const std::string& msg) : Error(msg) {}
-};
-
-struct UnificationError : dmlc::Error {
-  explicit UnificationError(const std::string& msg) : Error(msg) {}
-};
-
-struct SubstitutionError : dmlc::Error {
-  explicit SubstitutionError(const std::string& msg) : Error(msg) {}
-};
-
-/*! \brief A union-find data structure for the type-checker */
-class UnionFind;
-
-class UnionFindNode : public Node {
- public:
-  /*! \brief The inernal map from incomplete types to their representatives. */
-  tvm::Map<IncompleteType, Type> uf_map;
-
-  UnionFindNode() {}
-
-  void VisitAttrs(tvm::AttrVisitor* v) final { v->Visit("uf_map", &uf_map); }
-
-  TVM_DLL static UnionFind make(tvm::Map<IncompleteType, Type> uf_map);
-
-  /*! \brief Insert it into the union find.
-  * \param it The type to add to the union find.
-  */
-  void Insert(const IncompleteType& it);
-
-  /*! \brief Union operation, combine two equivalence classes.
-  * \param it The incomplete type to unify.
-  * \param ty The other type.
-  */
-  void Unify(const IncompleteType& it, const Type& t);
-
-  /*! \brief Find operation, returns the representative of the argument.
-  * \param it The element to lookup.
-  */
-  Type Find(const IncompleteType& it);
-
-  void debug();
-
-  void AssertAlphaEqual(const Type& l, const Type& r);
-
-  static constexpr const char* _type_key = "relay.UnionFind";
-  TVM_DECLARE_NODE_TYPE_INFO(UnionFindNode, Node);
-};
-
-class UnionFind : public NodeRef {
- public:
-  UnionFind() {}
-  explicit UnionFind(NodePtr<tvm::Node> p) : NodeRef(p) {}
-
-  // The union find structure is mutable so we do not use the standard macros
-  // and expose the pointer via `->`.
-  UnionFindNode* operator->() const {
-    return static_cast<UnionFindNode*>(node_.get());
-  }
-
-  using ContainerType = UnionFindNode;
-};
-
-class TypeUnifier;
-class TypeUnifierNode : public Node,
-                        private TypeFunctor<Type(const Type&, const Type)> {
- public:
-  UnionFind union_find;
-
-  TypeUnifierNode() {}
-
-  void VisitAttrs(tvm::AttrVisitor* v) final { v->Visit("union_find", &union_find); }
-
-  TVM_DLL static TypeUnifier make(UnionFind uf);
-
-  /*! \brief Introduces a new type var into the unifier */
-  void Insert(const IncompleteType& v);
-
-  /*! \brief Unifies two types if possible, throws a unification error if it
-   * cannot  */
-  Type Unify(const Type& t1, const Type& t2);
-
-  /*! \brief Attempts to substitute all type vars in t with concrete types,
-   * throws substitution error if it cannot concretize*/
-  Type Subst(const Type& t);
-
-  // /*! \brief Checks the kinds in the given type */
-  // Type CheckKinds(const Type& t);
-
-  static constexpr const char* _type_key = "relay.TypeUnifier";
-  TVM_DECLARE_NODE_TYPE_INFO(TypeUnifierNode, Node);
-
- private:
-  /*! \brief Unify incomplete type with another type. */
-  Type UnifyWithIncompleteType(const Type& t1, const IncompleteType tvn2);
-  /*! \brief Implements unification between two types with incomplete portions.
-   */
-  Type VisitType(const Type& t1, const Type t2) override;
-
-  // Visitor Cases
-  Type VisitType_(const IncompleteTypeNode* t1, const Type t2) override;
-  Type VisitType_(const TensorTypeNode* t1, const Type t2) override;
-  Type VisitType_(const TypeParamNode* t1, const Type t2) override;
-  Type VisitType_(const FuncTypeNode* t1, const Type t2) override;
-  Type VisitType_(const TupleTypeNode* t1, const Type t2) override;
-  Type VisitType_(const TypeRelationNode* s1, const Type t2) override;
-};
-
-class TypeUnifier : public NodeRef {
- public:
-  TypeUnifier() {}
-  explicit TypeUnifier(NodePtr<tvm::Node> p) : NodeRef(p) {}
-
-  // no const so that unifier can be mutable as a member of typechecker
-  inline TypeUnifierNode* operator->() const {
-    return static_cast<TypeUnifierNode*>(node_.get());
-  }
-
-  using ContainerType = TypeUnifierNode;
-};
-
-}  // namespace relay
-}  // namespace tvm
-#endif  // TVM_RELAY_PASS_UNIFIER_H_

tests/python/relay/test_ir_nodes.py

@@ -27,8 +27,8 @@ def test_tensor_type():
 
 
 def test_type_param():
-    tp = relay.TypeParam('name', relay.Kind.Shape)
-    assert tp.kind == relay.Kind.Shape
+    tp = relay.TypeParam('name', relay.Kind.Type)
+    assert tp.kind == relay.Kind.Type
     # assert tp.span  # TODO allow us to set span
     str(tp)
 

tests/python/relay/test_tyck_eval_integration.py → tests/python/relay/test_type_infer.py

@@ -76,10 +76,10 @@ def test_add_broadcast_op():
     assert_has_type(func.to_func(), expected_ty)
 
 def test_dual_op():
-    """Program: 
-       fn (x : Tensor[f32, (10, 10)]) { 
-         let t1 = log(x); 
-         let t2 = add(t1, x); 
+    """Program:
+       fn (x : Tensor[f32, (10, 10)]) {
+         let t1 = log(x);
+         let t2 = add(t1, x);
          return t1;
        }
     """
@@ -93,8 +93,8 @@ def test_dual_op():
 
 
 def test_decl():
-    """Program: 
-       def f(x : Tensor[f32, (10, 10)]) { 
+    """Program:
+       def f(x : Tensor[f32, (10, 10)]) {
            let lx = log(x);
            return lx;
        }
@@ -125,7 +125,7 @@ def test_recursion():
     n = b.param('n', ty='int32')
     data = b.param('data', ty='float32')
     with b.decl(f, n, data):
-        with b.if_scope(equal(n, convert(0.0))):
+        with b.if_scope(equal(n, convert(0))):
             b.ret(f(subtract(n, convert(1)), log(data)))
         with b.else_scope():
             b.ret(data)
@@ -152,11 +152,12 @@ def test_concat():
     assert_decl_has_type(ib.env, try_concat2, fn_ty)
 
 if __name__ == "__main__":
-    # test_monomorphic_let()
-    # test_single_op()
-    # test_add_op()
-    # test_add_broadcast_op()
-    # test_dual_op()
-    # test_decl()
-    # test_recursion()
+    test_recursion()
+
+    test_monomorphic_let()
+    test_single_op()
+    test_add_op()
+    test_add_broadcast_op()
+    test_dual_op()
+    test_decl()
     test_concat()

tests/python/relay/test_type_solver.py

+import tvm
+
+from tvm import relay
+from tvm.relay.ir_builder import scalar_type, convert, tensor_type
+
+
+def make_rel(name, args, num_inputs=None, attrs=None):
+    func = tvm.get_env_func("tvm.relay.type_relation." + name)
+    if num_inputs is None:
+        num_inputs = len(args) - 1
+    return relay.ty.TypeRelation(func, args, num_inputs, attrs)
+
+def make_solver():
+    solver = relay._ir_pass._test_type_solver()
+    solver.Solve = solver("Solve")
+    solver.Unify = solver("Unify")
+    solver.Resolve = solver("Resolve")
+    solver.AddConstraint = solver("AddConstraint")
+
+    def gen_type(name, args, out=None):
+        out = out if out else relay.ty.IncompleteType()
+        solver.AddConstraint(make_rel(name, args + [out]))
+        return out
+
+    solver.gen_type = gen_type
+    return solver
+
+
+def test_bcast():
+    solver = make_solver()
+    t0 = relay.ty.TensorType((10, 20), "float32")
+    t1 = relay.ty.TensorType((10, 1), "float32")
+    tc = relay.ty.TensorType((10, 1, 1), "float32")
+    t2 = solver.gen_type("Broadcast", [t0, t1])
+    t3 = solver.gen_type("Identity", [t2])
+    t4 = solver.gen_type("Broadcast", [t3, tc])
+    assert solver.Solve()
+    assert solver.Resolve(t2) == relay.ty.TensorType((10, 20), "float32")
+    assert solver.Resolve(t4) == relay.ty.TensorType((10, 10, 20), "float32")
+
+
+def test_backward_solving():
+    solver = make_solver()
+    t0 = relay.ty.TensorType((10, 20), "float32")
+    tc = relay.ty.TensorType((10, 1, 1), "float32")
+    t1 = relay.ty.IncompleteType()
+    t3 = solver.gen_type("Broadcast", [t0, t1])
+    t2 = solver.gen_type("Identity", [t1], out=tc)
+    assert solver.Solve()
+    assert solver.Resolve(t3) == relay.ty.TensorType((10, 10, 20), "float32")
+
+
+
+if __name__ == "__main__":
+    test_bcast()
+    test_backward_solving()

tutorials/nnvm/.gitignore

+*.pb
+*.mlmodel
+*.ttf
+*.txt
+*synset*txt
+*.cfg
+ssd_model
+*.names
+*.jpg
+*.pbtxt
+*.weights