revert relay/ir/*.py to relay

docs/api/python/relay/base.rst

@@ -15,16 +15,16 @@
     specific language governing permissions and limitations
     under the License.
 
-tvm.relay.ir.base
+tvm.relay.base
 --------------
-.. automodule:: tvm.relay.ir.base
+.. automodule:: tvm.relay.base
 
-.. autofunction:: tvm.relay.ir.base.register_relay_node
+.. autofunction:: tvm.relay.base.register_relay_node
 
-.. autofunction:: tvm.relay.ir.base.register_relay_attr_node
+.. autofunction:: tvm.relay.base.register_relay_attr_node
 
-.. autoclass:: tvm.relay.ir.base.RelayNode
+.. autoclass:: tvm.relay.base.RelayNode
     :members:
 
-.. autoclass:: tvm.relay.ir.base.Id
+.. autoclass:: tvm.relay.base.Id
     :members:

docs/api/python/relay/expr.rst

@@ -15,56 +15,56 @@
     specific language governing permissions and limitations
     under the License.
 
-tvm.relay.ir.expr
+tvm.relay.expr
 --------------
 
-.. automodule:: tvm.relay.ir.expr
+.. automodule:: tvm.relay.expr
 
-.. autofunction:: tvm.relay.ir.expr.var
+.. autofunction:: tvm.relay.expr.var
 
-.. autofunction:: tvm.relay.ir.expr.const
+.. autofunction:: tvm.relay.expr.const
 
-.. autofunction:: tvm.relay.ir.expr.bind
+.. autofunction:: tvm.relay.expr.bind
 
-.. autoclass:: tvm.relay.ir.expr.Expr
+.. autoclass:: tvm.relay.expr.Expr
     :members:
 
-.. autoclass:: tvm.relay.ir.expr.Constant
+.. autoclass:: tvm.relay.expr.Constant
     :members:
 
-.. autoclass:: tvm.relay.ir.expr.Tuple
+.. autoclass:: tvm.relay.expr.Tuple
     :members:
 
-.. autoclass:: tvm.relay.ir.expr.Function
+.. autoclass:: tvm.relay.expr.Function
     :members:
 
-.. autoclass:: tvm.relay.ir.expr.Call
+.. autoclass:: tvm.relay.expr.Call
     :members:
 
-.. autoclass:: tvm.relay.ir.expr.Let
+.. autoclass:: tvm.relay.expr.Let
     :members:
 
-.. autoclass:: tvm.relay.ir.expr.If
+.. autoclass:: tvm.relay.expr.If
     :members:
 
-.. autoclass:: tvm.relay.ir.expr.TupleGetItem
+.. autoclass:: tvm.relay.expr.TupleGetItem
     :members:
 
-.. autoclass:: tvm.relay.ir.expr.RefCreate
+.. autoclass:: tvm.relay.expr.RefCreate
     :members:
 
-.. autoclass:: tvm.relay.ir.expr.RefRead
+.. autoclass:: tvm.relay.expr.RefRead
     :members:
 
-.. autoclass:: tvm.relay.ir.expr.RefWrite
+.. autoclass:: tvm.relay.expr.RefWrite
     :members:
 
-.. autoclass:: tvm.relay.ir.expr.TupleGetItem
+.. autoclass:: tvm.relay.expr.TupleGetItem
     :members:
 
-.. autoclass:: tvm.relay.ir.expr.TempExpr
+.. autoclass:: tvm.relay.expr.TempExpr
     :members:
 
-.. autoclass:: tvm.relay.ir.expr.TupleWrapper
+.. autoclass:: tvm.relay.expr.TupleWrapper
     :members:
 

docs/api/python/relay/scope_builder.rst

@@ -15,10 +15,10 @@
     specific language governing permissions and limitations
     under the License.
 
-tvm.relay.ir.scope_builder
+tvm.relay.scope_builder
 -----------------------
 
-.. automodule:: tvm.relay.ir.scope_builder
+.. automodule:: tvm.relay.scope_builder
 
-.. autoclass:: tvm.relay.ir.scope_builder.ScopeBuilder
+.. autoclass:: tvm.relay.scope_builder.ScopeBuilder
     :members:

python/tvm/autotvm/graph_tuner/utils/traverse_graph.py

@@ -21,8 +21,8 @@
 import tvm
 from tvm import relay, autotvm
 from tvm.relay import transform
-from tvm.relay.ir import Call, Function, TupleGetItem, Var, Constant, Tuple
-from tvm.relay.ir import TupleType, TensorType
+from tvm.relay.expr import Call, Function, TupleGetItem, Var, Constant, Tuple
+from tvm.relay.ty import TupleType, TensorType
 from tvm.autotvm.task import TaskExtractEnv
 
 from .utils import has_multiple_inputs, is_boundary_node, is_skipped_node

python/tvm/relay/__init__.py

@@ -19,9 +19,16 @@
 import os
 from sys import setrecursionlimit
 
-from . import ir
-from .ir import adt, expr, ty, base, scope_builder
-from .ir import prelude, loops, parser
+from . import base
+from . import ty
+from . import expr
+from . import type_functor
+from . import expr_functor
+from . import adt
+from . import prelude
+from . import loops
+from . import scope_builder
+from . import parser
 
 from . import transform
 from . import analysis
@@ -60,65 +67,66 @@
 setrecursionlimit(10000)
 
 # Span
-Span = ir.Span
+Span = base.Span
+SourceName = base.SourceName
 
 # Type
-Type = ir.Type
-TupleType = ir.TupleType
-TensorType = ir.TensorType
-TypeKind = ir.TypeKind
-TypeVar = ir.TypeVar
-ShapeVar = ir.ShapeVar
-TypeConstraint = ir.TypeConstraint
-FuncType = ir.FuncType
-TypeRelation = ir.TypeRelation
-IncompleteType = ir.IncompleteType
-scalar_type = ir.scalar_type
-RefType = ir.RefType
-GlobalTypeVar = ir.GlobalTypeVar
-TypeCall = ir.TypeCall
-Any = ir.Any
+Type = ty.Type
+TupleType = ty.TupleType
+TensorType = ty.TensorType
+TypeKind = ty.TypeKind
+TypeVar = ty.TypeVar
+ShapeVar = ty.ShapeVar
+TypeConstraint = ty.TypeConstraint
+FuncType = ty.FuncType
+TypeRelation = ty.TypeRelation
+IncompleteType = ty.IncompleteType
+scalar_type = ty.scalar_type
+RefType = ty.RefType
+GlobalTypeVar = ty.GlobalTypeVar
+TypeCall = ty.TypeCall
+Any = ty.Any
 
 # Expr
-Expr = ir.Expr
-Constant = ir.Constant
-Tuple = ir.Tuple
-Var = ir.Var
-GlobalVar = ir.GlobalVar
-Function = ir.Function
-Call = ir.Call
-Let = ir.Let
-If = ir.If
-TupleGetItem = ir.TupleGetItem
-RefCreate = ir.RefCreate
-RefRead = ir.RefRead
-RefWrite = ir.RefWrite
+Expr = expr.RelayExpr
+Constant = expr.Constant
+Tuple = expr.Tuple
+Var = expr.Var
+GlobalVar = expr.GlobalVar
+Function = expr.Function
+Call = expr.Call
+Let = expr.Let
+If = expr.If
+TupleGetItem = expr.TupleGetItem
+RefCreate = expr.RefCreate
+RefRead = expr.RefRead
+RefWrite = expr.RefWrite
 
 # ADT
-Pattern = ir.Pattern
-PatternWildcard = ir.PatternWildcard
-PatternVar = ir.PatternVar
-PatternConstructor = ir.PatternConstructor
-PatternTuple = ir.PatternTuple
-Constructor = ir.Constructor
-TypeData = ir.TypeData
-Clause = ir.Clause
-Match = ir.Match
+Pattern = adt.Pattern
+PatternWildcard = adt.PatternWildcard
+PatternVar = adt.PatternVar
+PatternConstructor = adt.PatternConstructor
+PatternTuple = adt.PatternTuple
+Constructor = adt.Constructor
+TypeData = adt.TypeData
+Clause = adt.Clause
+Match = adt.Match
 
 # helper functions
-var = ir.var
-const = ir.const
-bind = ir.bind
+var = expr.var
+const = expr.const
+bind = expr.bind
 
 # TypeFunctor
-TypeFunctor = ir.TypeFunctor
-TypeVisitor = ir.TypeVisitor
-TypeMutator = ir.TypeMutator
+TypeFunctor = type_functor.TypeFunctor
+TypeVisitor = type_functor.TypeVisitor
+TypeMutator = type_functor.TypeMutator
 
 # ExprFunctor
-ExprFunctor = ir.ExprFunctor
-ExprVisitor = ir.ExprVisitor
-ExprMutator = ir.ExprMutator
+ExprFunctor = expr_functor.ExprFunctor
+ExprVisitor = expr_functor.ExprVisitor
+ExprMutator = expr_functor.ExprMutator
 
 # Prelude
 Prelude = prelude.Prelude

python/tvm/relay/ir/_parser.py → python/tvm/relay/_parser.py

@@ -40,11 +40,11 @@ def __new__(cls, *args, **kwds):
 import tvm.ir._ffi_api
 from tvm.ir import IRModule
 
-from . import Span, SourceName
+from .base import Span, SourceName
 from . import adt
 from . import expr
 from . import ty
-from .. import op
+from . import op
 
 PYTHON_VERSION = sys.version_info.major
 try:
@@ -56,9 +56,9 @@ def __new__(cls, *args, **kwds):
                     .format(version=PYTHON_VERSION))
 
 try:
-    from ..grammar.py3.RelayVisitor import RelayVisitor
-    from ..grammar.py3.RelayParser import RelayParser
-    from ..grammar.py3.RelayLexer import RelayLexer
+    from .grammar.py3.RelayVisitor import RelayVisitor
+    from .grammar.py3.RelayParser import RelayParser
+    from .grammar.py3.RelayLexer import RelayLexer
 except ImportError:
     raise Exception("Couldn't find ANTLR parser. Try building with USE_ANTLR=ON.")
 

python/tvm/relay/ir/adt.py → python/tvm/relay/adt.py

@@ -17,8 +17,10 @@
 # pylint: disable=no-else-return, unidiomatic-typecheck, invalid-name, unused-import
 """Algebraic data types in Relay."""
 from tvm.ir import Constructor, TypeData
+from tvm.runtime import Object
+import tvm._ffi
 
-from .base import RelayNode, register_relay_node, Object
+from .base import RelayNode
 from . import _ffi_api
 from .ty import Type
 from .expr import ExprWithOp, RelayExpr, Call
@@ -28,7 +30,7 @@ class Pattern(RelayNode):
     """Base type for pattern matching constructs."""
 
 
-@register_relay_node
+@tvm._ffi.register_object("relay.PatternWildcard")
 class PatternWildcard(Pattern):
     """Wildcard pattern in Relay: Matches any ADT and binds nothing."""
 
@@ -47,7 +49,7 @@ def __init__(self):
         self.__init_handle_by_constructor__(_ffi_api.PatternWildcard)
 
 
-@register_relay_node
+@tvm._ffi.register_object("relay.PatternVar")
 class PatternVar(Pattern):
     """Variable pattern in Relay: Matches anything and binds it to the variable."""
 
@@ -66,7 +68,7 @@ def __init__(self, var):
         self.__init_handle_by_constructor__(_ffi_api.PatternVar, var)
 
 
-@register_relay_node
+@tvm._ffi.register_object("relay.PatternConstructor")
 class PatternConstructor(Pattern):
     """Constructor pattern in Relay: Matches an ADT of the given constructor, binds recursively."""
 
@@ -91,7 +93,7 @@ def __init__(self, constructor, patterns=None):
         self.__init_handle_by_constructor__(_ffi_api.PatternConstructor, constructor, patterns)
 
 
-@register_relay_node
+@tvm._ffi.register_object("relay.PatternTuple")
 class PatternTuple(Pattern):
     """Constructor pattern in Relay: Matches a tuple, binds recursively."""
 
@@ -114,7 +116,7 @@ def __init__(self, patterns=None):
         self.__init_handle_by_constructor__(_ffi_api.PatternTuple, patterns)
 
 
-@register_relay_node
+@tvm._ffi.register_object("relay.Clause")
 class Clause(Object):
     """Clause for pattern matching in Relay."""
 
@@ -136,7 +138,7 @@ def __init__(self, lhs, rhs):
         self.__init_handle_by_constructor__(_ffi_api.Clause, lhs, rhs)
 
 
-@register_relay_node
+@tvm._ffi.register_object("relay.Match")
 class Match(ExprWithOp):
     """Pattern matching expression in Relay."""