refactor relay python

docs/api/python/relay/base.rst

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

docs/api/python/relay/expr.rst

@@ -15,56 +15,56 @@
     specific language governing permissions and limitations
     under the License.
 
-tvm.relay.expr
+tvm.relay.ir.expr
 --------------
 
-.. automodule:: tvm.relay.expr
+.. automodule:: tvm.relay.ir.expr
 
-.. autofunction:: tvm.relay.expr.var
+.. autofunction:: tvm.relay.ir.expr.var
 
-.. autofunction:: tvm.relay.expr.const
+.. autofunction:: tvm.relay.ir.expr.const
 
-.. autofunction:: tvm.relay.expr.bind
+.. autofunction:: tvm.relay.ir.expr.bind
 
-.. autoclass:: tvm.relay.expr.Expr
+.. autoclass:: tvm.relay.ir.expr.Expr
     :members:
 
-.. autoclass:: tvm.relay.expr.Constant
+.. autoclass:: tvm.relay.ir.expr.Constant
     :members:
 
-.. autoclass:: tvm.relay.expr.Tuple
+.. autoclass:: tvm.relay.ir.expr.Tuple
     :members:
 
-.. autoclass:: tvm.relay.expr.Function
+.. autoclass:: tvm.relay.ir.expr.Function
     :members:
 
-.. autoclass:: tvm.relay.expr.Call
+.. autoclass:: tvm.relay.ir.expr.Call
     :members:
 
-.. autoclass:: tvm.relay.expr.Let
+.. autoclass:: tvm.relay.ir.expr.Let
     :members:
 
-.. autoclass:: tvm.relay.expr.If
+.. autoclass:: tvm.relay.ir.expr.If
     :members:
 
-.. autoclass:: tvm.relay.expr.TupleGetItem
+.. autoclass:: tvm.relay.ir.expr.TupleGetItem
     :members:
 
-.. autoclass:: tvm.relay.expr.RefCreate
+.. autoclass:: tvm.relay.ir.expr.RefCreate
     :members:
 
-.. autoclass:: tvm.relay.expr.RefRead
+.. autoclass:: tvm.relay.ir.expr.RefRead
     :members:
 
-.. autoclass:: tvm.relay.expr.RefWrite
+.. autoclass:: tvm.relay.ir.expr.RefWrite
     :members:
 
-.. autoclass:: tvm.relay.expr.TupleGetItem
+.. autoclass:: tvm.relay.ir.expr.TupleGetItem
     :members:
 
-.. autoclass:: tvm.relay.expr.TempExpr
+.. autoclass:: tvm.relay.ir.expr.TempExpr
     :members:
 
-.. autoclass:: tvm.relay.expr.TupleWrapper
+.. autoclass:: tvm.relay.ir.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.scope_builder
+tvm.relay.ir.scope_builder
 -----------------------
 
-.. automodule:: tvm.relay.scope_builder
+.. automodule:: tvm.relay.ir.scope_builder
 
-.. autoclass:: tvm.relay.scope_builder.ScopeBuilder
+.. autoclass:: tvm.relay.ir.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.expr import Call, Function, TupleGetItem, Var, Constant, Tuple
-from tvm.relay.ty import TupleType, TensorType
+from tvm.relay.ir import Call, Function, TupleGetItem, Var, Constant, Tuple
+from tvm.relay.ir 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,111 +19,115 @@
 import os
 from sys import setrecursionlimit
 
-from . import call_graph
-from . import base
-from . import ty
-from . import expr
-from . import type_functor
-from . import expr_functor
-from . import adt
-from . import analysis
+from . import ir
+from .ir import adt, expr, ty, base, scope_builder
+from .ir import prelude, loops, parser
+
 from . import transform
+from . import analysis
+from .analysis import call_graph, feature, alpha_equal
 from .build_module import build, create_executor, optimize
 from .transform import build_config
-from . import prelude
-from . import parser
 from . import debug
 from . import param_dict
-from . import feature
 from .backend import vm
 
 # Root operators
 from .op import Op
+from .op import nn
+from .op import image
+from .op import vision
+from .op import annotation
 from .op.reduce import *
 from .op.tensor import *
 from .op.transform import *
 from .op.algorithm import *
-from . import nn
-from . import annotation
-from . import vision
-from . import contrib
-from . import image
+from .op.nn import *
+from .op.vision import *
+from .op.contrib import *
+from .op.image import *
 from . import frontend
 from . import backend
 from . import quantize
 
 # Dialects
 from . import qnn
 
-from .scope_builder import ScopeBuilder
 # Load Memory pass
-from . import memory_alloc
+from .transform import memory_alloc
 
 # Required to traverse large programs
 setrecursionlimit(10000)
 
 # Span
-Span = base.Span
+Span = ir.Span
 
 # Type
-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
+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
 
 # Expr
-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
+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
 
 # ADT
-PatternWildcard = adt.PatternWildcard
-PatternVar = adt.PatternVar
-PatternConstructor = adt.PatternConstructor
-PatternTuple = adt.PatternTuple
-Constructor = adt.Constructor
-TypeData = adt.TypeData
-Clause = adt.Clause
-Match = adt.Match
+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
 
 # helper functions
-var = expr.var
-const = expr.const
-bind = expr.bind
-module_pass = transform.module_pass
-function_pass = transform.function_pass
-alpha_equal = analysis.alpha_equal
+var = ir.var
+const = ir.const
+bind = ir.bind
 
 # TypeFunctor
-TypeFunctor = type_functor.TypeFunctor
-TypeVisitor = type_functor.TypeVisitor
-TypeMutator = type_functor.TypeMutator
+TypeFunctor = ir.TypeFunctor
+TypeVisitor = ir.TypeVisitor
+TypeMutator = ir.TypeMutator
 
 # ExprFunctor
-ExprFunctor = expr_functor.ExprFunctor
-ExprVisitor = expr_functor.ExprVisitor
-ExprMutator = expr_functor.ExprMutator
+ExprFunctor = ir.ExprFunctor
+ExprVisitor = ir.ExprVisitor
+ExprMutator = ir.ExprMutator
+
+# Prelude
+Prelude = prelude.Prelude
+
+# Scope builder
+ScopeBuilder = scope_builder.ScopeBuilder
+
+module_pass = transform.module_pass
+function_pass = transform.function_pass
 
 # Parser
 fromtext = parser.fromtext

python/tvm/relay/analysis/__init__.py

@@ -0,0 +1,26 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=wildcard-import, redefined-builtin, invalid-name
+"""The Relay IR namespace containing the analysis passes."""
+# Analysis passes
+from .analysis import *
+
+# Call graph
+from . import call_graph
+
+# Feature
+from . import feature

python/tvm/relay/_analysis.py → python/tvm/relay/analysis/_ffi_api.py

@@ -14,7 +14,7 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-"""FFI exposing the passes for Relay program analysis."""
+"""FFI APIs for Relay program analysis."""
 import tvm._ffi
 
-tvm._ffi._init_api("relay._analysis", __name__)
+tvm._ffi._init_api("relay.analysis", __name__)