[Refactor][Relay] Refactor Relay Python to use new FFI (#5077)

* refactor relay python

* revert relay/ir/*.py to relay

* Address comments

* remove direct access to analysis and transform namespace

docs/api/python/relay/base.rst

@@ -19,10 +19,6 @@ tvm.relay.base
 --------------
 .. automodule:: tvm.relay.base
 
-.. autofunction:: tvm.relay.base.register_relay_node
-
-.. autofunction:: tvm.relay.base.register_relay_attr_node
-
 .. autoclass:: tvm.relay.base.RelayNode
     :members:
 

python/tvm/relay/__init__.py

@@ -19,35 +19,37 @@
 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 prelude
+from . import loops
+from . import scope_builder
+from . import parser
+
 from . import transform
+from . import analysis
+from .analysis import 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 annotation
+from .op import vision
+from .op import contrib
 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 . import frontend
 from . import backend
 from . import quantize
@@ -55,15 +57,12 @@ from . import quantize
 # Dialects
 from . import qnn
 
-from .scope_builder import ScopeBuilder
-# Load Memory pass
-from . import memory_alloc
-
 # Required to traverse large programs
 setrecursionlimit(10000)
 
 # Span
 Span = base.Span
+SourceName = base.SourceName
 
 # Type
 Type = ty.Type
@@ -98,6 +97,7 @@ RefRead = expr.RefRead
 RefWrite = expr.RefWrite
 
 # ADT
+Pattern = adt.Pattern
 PatternWildcard = adt.PatternWildcard
 PatternVar = adt.PatternVar
 PatternConstructor = adt.PatternConstructor
@@ -111,9 +111,6 @@ Match = adt.Match
 var = expr.var
 const = expr.const
 bind = expr.bind
-module_pass = transform.module_pass
-function_pass = transform.function_pass
-alpha_equal = analysis.alpha_equal
 
 # TypeFunctor
 TypeFunctor = type_functor.TypeFunctor
@@ -125,6 +122,15 @@ ExprFunctor = expr_functor.ExprFunctor
 ExprVisitor = expr_functor.ExprVisitor
 ExprMutator = expr_functor.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
 
@@ -139,9 +145,3 @@ Pass = transform.Pass
 ModulePass = transform.ModulePass
 FunctionPass = transform.FunctionPass
 Sequential = transform.Sequential
-
-# Feature
-Feature = feature.Feature
-
-# CallGraph
-CallGraph = call_graph.CallGraph

python/tvm/relay/_analysis.py → python/tvm/relay/_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 IR."""
 import tvm._ffi
 
-tvm._ffi._init_api("relay._analysis", __name__)
+tvm._ffi._init_api("relay.ir", __name__)

python/tvm/relay/adt.py

@@ -17,9 +17,11 @@
 # 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 . import _make
+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."""
 
@@ -44,10 +46,10 @@ class PatternWildcard(Pattern):
         wildcard: PatternWildcard
             a wildcard pattern.
         """
-        self.__init_handle_by_constructor__(_make.PatternWildcard)
+        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."""
 
@@ -63,10 +65,10 @@ class PatternVar(Pattern):
         pv: PatternVar
             A variable pattern.
         """
-        self.__init_handle_by_constructor__(_make.PatternVar, 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."""
 
@@ -88,10 +90,10 @@ class PatternConstructor(Pattern):
         """
         if patterns is None:
             patterns = []
-        self.__init_handle_by_constructor__(_make.PatternConstructor, constructor, patterns)
+        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."""
 
@@ -111,10 +113,10 @@ class PatternTuple(Pattern):
         """
         if patterns is None:
             patterns = []
-        self.__init_handle_by_constructor__(_make.PatternTuple, patterns)
+        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."""
 
@@ -133,10 +135,10 @@ class Clause(Object):
         clause: Clause
             The Clause.
         """
-        self.__init_handle_by_constructor__(_make.Clause, 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."""
 
@@ -160,4 +162,4 @@ class Match(ExprWithOp):
         match: tvm.relay.Expr
             The match expression.
         """
-        self.__init_handle_by_constructor__(_make.Match, data, clauses, complete)
+        self.__init_handle_by_constructor__(_ffi_api.Match, data, clauses, complete)

python/tvm/relay/annotation.py → python/tvm/relay/analysis/__init__.py

@@ -14,8 +14,15 @@
 # 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 *
 
-# pylint: disable=wildcard-import, unused-import, unused-wildcard-import
-"""Annotation related operators."""
-# Re-export in a specific file name so that autodoc can pick it up
-from .op.annotation import *
+# Call graph
+from . import call_graph
+
+# Feature
+from . import feature
+
+CallGraph = call_graph.CallGraph

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

@@ -14,8 +14,7 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-# pylint: disable=no-else-return, unidiomatic-typecheck, undefined-variable
-"""The interface of expr function exposed from C++."""
+"""FFI APIs for Relay program analysis."""
 import tvm._ffi
 
-tvm._ffi._init_api("relay._base", __name__)
+tvm._ffi._init_api("relay.analysis", __name__)

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

@@ -22,9 +22,9 @@ configuring the passes and scripting them in Python.
 """
 from tvm.ir import RelayExpr, IRModule
 
-from . import _analysis
-from .ty import Type
+from . import _ffi_api
 from .feature import Feature
+from ..ty import Type
 
 
 def post_order_visit(expr, fvisit):
@@ -40,7 +40,7 @@ def post_order_visit(expr, fvisit):
     fvisit : function
         The visitor function to be applied.
     """
-    return _analysis.post_order_visit(expr, fvisit)
+    return _ffi_api.post_order_visit(expr, fvisit)
 
 
 def well_formed(expr):
@@ -56,7 +56,7 @@ def well_formed(expr):
     well_form : bool
         Whether the input expression is well formed
     """
-    return _analysis.well_formed(expr)
+    return _ffi_api.well_formed(expr)
 
 
 def check_kind(t, mod=None):
@@ -85,9 +85,9 @@ def check_kind(t, mod=None):
         assert check_kind(relay.TupleType([relay.TypeParam('tp1', relay.Kind.Type)])) == Type
     """
     if mod is not None:
-        return _analysis.check_kind(t, mod)
+        return _ffi_api.check_kind(t, mod)
     else:
-        return _analysis.check_kind(t)
+        return _ffi_api.check_kind(t)
 
 
 def check_constant(expr):
@@ -103,7 +103,7 @@ def check_constant(expr):
     result : bool
         Whether the expression is constant.
     """
-    return _analysis.check_constant(expr)
+    return _ffi_api.check_constant(expr)
 
 
 def free_vars(expr):
@@ -125,7 +125,7 @@ def free_vars(expr):
     neural networks: usually this means weights of previous
     are ordered first.
     """
-    return _analysis.free_vars(expr)
+    return _ffi_api.free_vars(expr)
 
 
 def bound_vars(expr):
@@ -141,7 +141,7 @@ def bound_vars(expr):
     free : List[tvm.relay.Var]
         The list of bound variables in post-DFS order.
     """
-    return _analysis.bound_vars(expr)
+    return _ffi_api.bound_vars(expr)
 
 
 def all_vars(expr):
@@ -157,7 +157,7 @@ def all_vars(expr):
     free : List[tvm.relay.Var]
         The list of all variables in post-DFS order.
     """
-    return _analysis.all_vars(expr)
+    return _ffi_api.all_vars(expr)
 
 
 def free_type_vars(expr, mod=None):
@@ -177,7 +177,7 @@ def free_type_vars(expr, mod=None):
         The list of free type variables in post-DFS order
     """
     use_mod = mod if mod is not None else IRModule()
-    return _analysis.free_type_vars(expr, use_mod)
+    return _ffi_api.free_type_vars(expr, use_mod)
 
 
 def bound_type_vars(expr, mod=None):
@@ -197,7 +197,7 @@ def bound_type_vars(expr, mod=None):
         The list of bound type variables in post-DFS order
     """
     use_mod = mod if mod is not None else IRModule()
-    return _analysis.bound_type_vars(expr, use_mod)
+    return _ffi_api.bound_type_vars(expr, use_mod)
 
 
 def all_type_vars(expr, mod=None):
@@ -217,7 +217,7 @@ def all_type_vars(expr, mod=None):
         The list of all type variables in post-DFS order
     """
     use_mod = mod if mod is not None else IRModule()
-    return _analysis.all_type_vars(expr, use_mod)
+    return _ffi_api.all_type_vars(expr, use_mod)
 
 
 def alpha_equal(lhs, rhs):
@@ -236,7 +236,7 @@ def alpha_equal(lhs, rhs):
     result : bool
         True iff lhs is alpha equal to rhs.
     """
-    return bool(_analysis._alpha_equal(lhs, rhs))
+    return bool(_ffi_api._alpha_equal(lhs, rhs))
 
 
 def assert_alpha_equal(lhs, rhs):
@@ -250,7 +250,7 @@ def assert_alpha_equal(lhs, rhs):
     rhs : tvm.relay.Expr
         One of the input Expression.
     """
-    _analysis._assert_alpha_equal(lhs, rhs)
+    _ffi_api._assert_alpha_equal(lhs, rhs)
 
 
 def graph_equal(lhs, rhs):
@@ -272,7 +272,7 @@ def graph_equal(lhs, rhs):
     result : bool
       True iff lhs is data-flow equivalent to rhs.
     """
-    return bool(_analysis._graph_equal(lhs, rhs))
+    return bool(_ffi_api._graph_equal(lhs, rhs))
 
 
 def assert_graph_equal(lhs, rhs):
@@ -289,7 +289,7 @@ def assert_graph_equal(lhs, rhs):
     rhs : tvm.relay.Expr
       One of the input Expression.
     """
-    _analysis._assert_graph_equal(lhs, rhs)
+    _ffi_api._assert_graph_equal(lhs, rhs)
 
 
 def collect_device_info(expr):
@@ -303,10 +303,10 @@ def collect_device_info(expr):
 
     Returns
     -------
-    ret : Dict[tvm.relay.expr, int]
+    ret : Dict[tvm.relay.ir.expr, int]
         A dictionary mapping tvm.relay.Expr to device type.
     """
-    return _analysis.CollectDeviceInfo(expr)
+    return _ffi_api.CollectDeviceInfo(expr)
 
 
 def collect_device_annotation_ops(expr):
@@ -319,11 +319,11 @@ def collect_device_annotation_ops(expr):
 
     Returns
     -------
-    ret : Dict[tvm.relay.expr, int]
+    ret : Dict[tvm.relay.Expr, int]
         A dictionary mapping tvm.relay.Expr to device type where the keys are
         annotation expressions.
     """
-    return _analysis.CollectDeviceAnnotationOps(expr)
+    return _ffi_api.CollectDeviceAnnotationOps(expr)
 
 
 def get_total_mac_number(expr):
@@ -340,7 +340,7 @@ def get_total_mac_number(expr):
     result : int64
       The number of MACs (multiply-accumulate) of a model
     """
-    return _analysis.GetTotalMacNumber(expr)
+    return _ffi_api.GetTotalMacNumber(expr)
 
 
 def unmatched_cases(match, mod=None):
@@ -360,7 +360,7 @@ def unmatched_cases(match, mod=None):
     missing_patterns : [tvm.relay.Pattern]
         Patterns that the match expression does not catch.
     """
-    return _analysis.unmatched_cases(match, mod)
+    return _ffi_api.unmatched_cases(match, mod)
 
 
 def detect_feature(a, b=None):
@@ -383,7 +383,7 @@ def detect_feature(a, b=None):
     """
     if isinstance(a, IRModule):
         a, b = b, a
-    return {Feature(int(x)) for x in _analysis.detect_feature(a, b)}
+    return {Feature(int(x)) for x in _ffi_api.detect_feature(a, b)}
 
 
 def structural_hash(value):
@@ -400,9 +400,9 @@ def structural_hash(value):
       The hash value
     """
     if isinstance(value, RelayExpr):
-        return int(_analysis._expr_hash(value))
+        return int(_ffi_api._expr_hash(value))
     elif isinstance(value, Type):
-        return int(_analysis._type_hash(value))
+        return int(_ffi_api._type_hash(value))
     else:
         msg = ("found value of type {0} expected" +
                "relay.Expr or relay.Type").format(type(value))
@@ -421,10 +421,10 @@ def extract_fused_functions(mod):
 
     Returns
     -------
-    ret : Dict[int, tvm.relay.expr.Function]
+    ret : Dict[int, tvm.relay.ir.expr.Function]
         A module containing only fused primitive functions
     """
-    ret_mod = _analysis.ExtractFusedFunctions()(mod)
+    ret_mod = _ffi_api.ExtractFusedFunctions()(mod)
     ret = {}
     for hash_, func in ret_mod.functions.items():
         ret[hash_] = func

python/tvm/relay/call_graph.py → python/tvm/relay/analysis/call_graph.py

@@ -18,9 +18,9 @@
 """Call graph used in Relay."""
 
 from tvm.ir import IRModule
-from .base import Object
-from .expr import GlobalVar
-from . import _analysis
+from tvm.runtime import Object
+from ..expr import GlobalVar
+from . import _ffi_api
 
 
 class CallGraph(Object):
@@ -39,7 +39,7 @@ class CallGraph(Object):
         call_graph: CallGraph
             A constructed call graph.
         """
-        self.__init_handle_by_constructor__(_analysis.CallGraph, module)
+        self.__init_handle_by_constructor__(_ffi_api.CallGraph, module)
 
     @property
     def module(self):
@@ -54,7 +54,7 @@ class CallGraph(Object):
         ret : tvm.ir.IRModule
             The contained IRModule
         """
-        return _analysis.GetModule(self)
+        return _ffi_api.GetModule(self)
 
     def ref_count(self, var):
         """Return the number of references to the global var
@@ -69,7 +69,7 @@ class CallGraph(Object):
             The number reference to the global var
         """
         var = self._get_global_var(var)
-        return _analysis.GetRefCountGlobalVar(self, var)
+        return _ffi_api.GetRefCountGlobalVar(self, var)
 
     def global_call_count(self, var):
         """Return the number of global function calls from a given global var.
@@ -84,7 +84,7 @@ class CallGraph(Object):
             The number of global function calls from the given var.
         """
         var = self._get_global_var(var)
-        return _analysis.GetGlobalVarCallCount(self, var)
+        return _ffi_api.GetGlobalVarCallCount(self, var)
 
     def is_recursive(self, var):
         """Return if the function corresponding to a var is a recursive
@@ -100,7 +100,7 @@ class CallGraph(Object):
             If the function corresponding to var is recurisve.
         """
         var = self._get_global_var(var)
-        return _analysis.IsRecursive(self, var)
+        return _ffi_api.IsRecursive(self, var)
 
     def _get_global_var(self, var):
         """Return the global var using a given name or GlobalVar.
@@ -137,8 +137,8 @@ class CallGraph(Object):
             The call graph represented in string.
         """
         var = self._get_global_var(var)
-        return _analysis.PrintCallGraphGlobalVar(self, var)
+        return _ffi_api.PrintCallGraphGlobalVar(self, var)
 
     def __str__(self):
         """Print the call graph in the topological order."""
-        return _analysis.PrintCallGraph(self)
+        return _ffi_api.PrintCallGraph(self)

python/tvm/relay/backend/compile_engine.py

@@ -22,7 +22,7 @@ import logging
 import numpy as np
 import tvm
 from tvm import te
-from ..base import register_relay_node, Object
+from tvm.runtime import Object
 from ... import target as _target
 from ... import autotvm
 from .. import expr as _expr
@@ -33,7 +33,7 @@ from . import _backend
 logger = logging.getLogger('compile_engine')
 
 
-@register_relay_node
+@tvm._ffi.register_object("relay.LoweredOutput")
 class LoweredOutput(Object):
     """Lowered output"""
     def __init__(self, outputs, implement):
@@ -41,7 +41,7 @@ class LoweredOutput(Object):
             _backend._make_LoweredOutput, outputs, implement)
 
 
-@register_relay_node
+@tvm._ffi.register_object("relay.CCacheKey")
 class CCacheKey(Object):
     """Key in the CompileEngine.
 
@@ -58,7 +58,7 @@ class CCacheKey(Object):
             _backend._make_CCacheKey, source_func, target)
 
 
-@register_relay_node
+@tvm._ffi.register_object("relay.CCacheValue")
 class CCacheValue(Object):
     """Value in the CompileEngine, including usage statistics.
     """
@@ -261,7 +261,7 @@ def lower_call(call, inputs, target):
     return LoweredOutput(outputs, best_impl)
 
 
-@register_relay_node
+@tvm._ffi.register_object("relay.CompileEngine")
 class CompileEngine(Object):
     """CompileEngine to get lowered code.
     """

python/tvm/relay/backend/interpreter.py

@@ -20,25 +20,25 @@ from __future__ import absolute_import
 
 import numpy as np
 
-from tvm.runtime import container
+import tvm._ffi
+from tvm.runtime import container, Object
 from tvm.ir import IRModule
 
 from . import _backend
 from .. import _make, analysis, transform
 from ... import nd
-from ..base import Object, register_relay_node
 from ..expr import Tuple, RefCreate, Call, Constant, GlobalVar, Function, const
 from ..scope_builder import ScopeBuilder
 
 
-@register_relay_node
+@tvm._ffi.register_object("relay.ConstructorValue")
 class ConstructorValue(Object):
     def __init__(self, tag, fields, constructor):
         self.__init_handle_by_constructor__(
             _make.ConstructorValue, tag, fields, constructor)
 
 
-@register_relay_node
+@tvm._ffi.register_object("relay.RefValue")
 class RefValue(Object):
     def __init__(self, value):
         self.__init_handle_by_constructor__(

python/tvm/relay/base.py

@@ -21,47 +21,17 @@ import tvm._ffi
 
 from tvm.runtime import Object
 from tvm.ir import SourceName, Span, Node as RelayNode
-from . import _make
-from . import _expr
-from . import _base
 
 
 __STD_PATH__ = os.path.join(os.path.dirname(os.path.realpath(__file__)), "std")
 
+
 @tvm._ffi.register_func("tvm.relay.std_path")
 def _std_path():
     return __STD_PATH__
 
 
-def register_relay_node(type_key=None):
-    """Register a Relay node type.
-
-    Parameters
-    ----------
-    type_key : str or cls
-        The type key of the node.
-    """
-    if not isinstance(type_key, str):
-        return tvm._ffi.register_object(
-            "relay." + type_key.__name__)(type_key)
-    return tvm._ffi.register_object(type_key)
-
-
-def register_relay_attr_node(type_key=None):
-    """Register a Relay attribute node.
-
-    Parameters
-    ----------
-    type_key : str or cls
-        The type key of the node.
-    """
-    if not isinstance(type_key, str):
-        return tvm._ffi.register_object(
-            "relay.attrs." + type_key.__name__)(type_key)
-    return tvm._ffi.register_object(type_key)
-
-
-@register_relay_node
+@tvm._ffi.register_object("relay.Id")
 class Id(Object):
     """Unique identifier(name) used in Var.
        Guaranteed to be stable across all passes.

python/tvm/relay/contrib.py

-# 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, unused-import, unused-wildcard-import
-"""Contrib operators."""
-# Re-export in a specific file name so that autodoc can pick it up
-from .op.contrib import *

python/tvm/relay/expr.py

@@ -20,13 +20,13 @@ from __future__ import absolute_import
 from numbers import Number as _Number
 
 import numpy as _np
+import tvm._ffi
 from tvm._ffi import base as _base
 from tvm.runtime import NDArray, convert, ndarray as _nd
 from tvm.ir import RelayExpr, GlobalVar, BaseFunc
 
-from .base import RelayNode, register_relay_node
-from . import _make
-from . import _expr
+from .base import RelayNode
+from . import _ffi_api
 from . import ty as _ty
 
 # alias relay expr as Expr.
@@ -54,7 +54,7 @@ class ExprWithOp(RelayExpr):
         result : tvm.relay.Expr
             The result expression.
         """
-        return _make.cast(self, dtype)
+        return _ffi_api.cast(self, dtype)
 
     def __neg__(self):
         return _op_make.negative(self)
@@ -160,7 +160,7 @@ class ExprWithOp(RelayExpr):
         """
         return Call(self, args)
 
-@register_relay_node
+@tvm._ffi.register_object("relay.Constant")
 class Constant(ExprWithOp):
     """A constant expression in Relay.
 
@@ -170,10 +170,10 @@ class Constant(ExprWithOp):
         The data content of the constant expression.
     """
     def __init__(self, data):
-        self.__init_handle_by_constructor__(_make.Constant, data)
+        self.__init_handle_by_constructor__(_ffi_api.Constant, data)
 
 
-@register_relay_node
+@tvm._ffi.register_object("relay.Tuple")
 class Tuple(ExprWithOp):
     """Tuple expression that groups several fields together.
 
@@ -183,7 +183,7 @@ class Tuple(ExprWithOp):
         The fields in the tuple.
     """
     def __init__(self, fields):
-        self.__init_handle_by_constructor__(_make.Tuple, fields)
+        self.__init_handle_by_constructor__(_ffi_api.Tuple, fields)
 
     def __getitem__(self, index):
         if index >= len(self):
@@ -197,7 +197,7 @@ class Tuple(ExprWithOp):
         raise TypeError("astype cannot be used on tuple")
 
 
-@register_relay_node
+@tvm._ffi.register_object("relay.Var")
 class Var(ExprWithOp):
     """A local variable in Relay.
 
@@ -216,7 +216,7 @@ class Var(ExprWithOp):
     """
     def __init__(self, name_hint, type_annotation=None):
         self.__init_handle_by_constructor__(
-            _make.Var, name_hint, type_annotation)
+            _ffi_api.Var, name_hint, type_annotation)
 
     @property
     def name_hint(self):
@@ -225,7 +225,7 @@ class Var(ExprWithOp):
         return name
 
 
-@register_relay_node
+@tvm._ffi.register_object("relay.Function")
 class Function(BaseFunc):
     """A function declaration expression.
 
@@ -254,7 +254,7 @@ class Function(BaseFunc):
             type_params = convert([])
 
         self.__init_handle_by_constructor__(
-            _make.Function, params, body, ret_type, type_params, attrs)
+            _ffi_api.Function, params, body, ret_type, type_params, attrs)
 
     def __call__(self, *args):
         """Invoke the global function.
@@ -282,12 +282,12 @@ class Function(BaseFunc):
         func : Function
             A new copy of the function
         """
-        return _expr.FunctionWithAttr(
+        return _ffi_api.FunctionWithAttr(
             self, attr_key, convert(attr_value))
 
 
 
-@register_relay_node
+@tvm._ffi.register_object("relay.Call")
 class Call(ExprWithOp):
     """Function call node in Relay.
 
@@ -313,10 +313,10 @@ class Call(ExprWithOp):
         if not type_args:
             type_args = []
         self.__init_handle_by_constructor__(
-            _make.Call, op, args, attrs, type_args)
+            _ffi_api.Call, op, args, attrs, type_args)
 
 
-@register_relay_node
+@tvm._ffi.register_object("relay.Let")
 class Let(ExprWithOp):
     """Let variable binding expression.
 
@@ -333,10 +333,10 @@ class Let(ExprWithOp):
     """
     def __init__(self, variable, value, body):
         self.__init_handle_by_constructor__(
-            _make.Let, variable, value, body)
+            _ffi_api.Let, variable, value, body)
 
 
-@register_relay_node
+@tvm._ffi.register_object("relay.If")
 class If(ExprWithOp):
     """A conditional expression in Relay.
 
@@ -353,10 +353,10 @@ class If(ExprWithOp):
     """
     def __init__(self, cond, true_branch, false_branch):
         self.__init_handle_by_constructor__(
-            _make.If, cond, true_branch, false_branch)
+            _ffi_api.If, cond, true_branch, false_branch)
 
 
-@register_relay_node
+@tvm._ffi.register_object("relay.TupleGetItem")
 class TupleGetItem(ExprWithOp):
     """Get index-th item from a tuple.
 
@@ -370,10 +370,10 @@ class TupleGetItem(ExprWithOp):
     """
     def __init__(self, tuple_value, index):
         self.__init_handle_by_constructor__(
-            _make.TupleGetItem, tuple_value, index)
+            _ffi_api.TupleGetItem, tuple_value, index)
 
 
-@register_relay_node
+@tvm._ffi.register_object("relay.RefCreate")
 class RefCreate(ExprWithOp):
     """Create a new reference from initial value.
     Parameters
@@ -382,10 +382,10 @@ class RefCreate(ExprWithOp):
        The initial value.
     """
     def __init__(self, value):
-        self.__init_handle_by_constructor__(_make.RefCreate, value)
+        self.__init_handle_by_constructor__(_ffi_api.RefCreate, value)
 
 
-@register_relay_node
+@tvm._ffi.register_object("relay.RefRead")
 class RefRead(ExprWithOp):
     """Get the value inside the reference.
     Parameters
@@ -394,10 +394,10 @@ class RefRead(ExprWithOp):
          The reference.
     """
     def __init__(self, ref):
-        self.__init_handle_by_constructor__(_make.RefRead, ref)
+        self.__init_handle_by_constructor__(_ffi_api.RefRead, ref)
 
 
-@register_relay_node
+@tvm._ffi.register_object("relay.RefWrite")
 class RefWrite(ExprWithOp):
     """
     Update the value inside the reference.
@@ -410,7 +410,7 @@ class RefWrite(ExprWithOp):
         The new value.
     """
     def __init__(self, ref, value):
-        self.__init_handle_by_constructor__(_make.RefWrite, ref, value)
+        self.__init_handle_by_constructor__(_ffi_api.RefWrite, ref, value)
 
 
 class TempExpr(ExprWithOp):
@@ -427,7 +427,7 @@ class TempExpr(ExprWithOp):
         -------
         The corresponding normal expression.
         """
-        return _expr.TempExprRealize(self)
+        return _ffi_api.TempExprRealize(self)
 
 
 class TupleWrapper(object):
@@ -587,4 +587,4 @@ def bind(expr, binds):
     result : tvm.relay.Expr
         The expression or function after binding.
     """
-    return _expr.Bind(expr, binds)
+    return _ffi_api.Bind(expr, binds)

python/tvm/relay/frontend/tflite.py

@@ -26,8 +26,8 @@ from .. import analysis
 from .. import expr as _expr
 from .. import op as _op
 from .. import qnn as _qnn
-from ..util import get_scalar_from_constant
 from ... import nd as _nd
+from .util import get_scalar_from_constant
 from .common import ExprTable
 from .common import infer_shape as _infer_shape
 

python/tvm/relay/util.py → python/tvm/relay/frontend/util.py

@@ -18,7 +18,7 @@
 """ Utility functions that are used across many directories. """
 from __future__ import absolute_import
 import numpy as np
-from . import expr as _expr
+from .. import expr as _expr
 
 def get_scalar_from_constant(expr):
     """ Returns scalar value from Relay constant scalar. """

python/tvm/relay/image.py

-# 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, unused-import, unused-wildcard-import
-"""Image network related operators."""
-# Re-export in a specific file name so that autodoc can pick it up
-from .op.image import *

python/tvm/relay/nn.py

-# 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, unused-import, unused-wildcard-import
-"""Neural network related operators."""
-# Re-export in a specific file name so that autodoc can pick it up
-from .op.nn import *

python/tvm/relay/op/__init__.py

@@ -41,7 +41,6 @@ from . import _tensor_grad
 from . import _transform
 from . import _reduce
 from . import _algorithm
-from ..base import register_relay_node
 
 
 def _register_op_make():

python/tvm/relay/op/op.py

@@ -19,13 +19,12 @@
 import tvm._ffi
 from tvm.driver import lower, build
 
-from ..base import register_relay_node
 from ..expr import RelayExpr
 from ...target import get_native_generic_func, GenericFunc
 from ...runtime import Object
 from . import _make
 
-@register_relay_node
+@tvm._ffi.register_object("relay.Op")
 class Op(RelayExpr):
     """A Relay operator definition."""
 

python/tvm/relay/op/transform.py

@@ -38,7 +38,7 @@ def cast(data, dtype):
     result : relay.Expr
         The casted result.
     """
-    from .. import _make as _relay_make
+    from .. import _ffi_api as _relay_make
     return _relay_make.cast(data, dtype)
 
 
@@ -55,7 +55,7 @@ def cast_like(data, dtype_like):
     result : relay.Expr
         The casted result.
     """
-    from .. import _make as _relay_make
+    from .. import _ffi_api as _relay_make
     return _relay_make.cast_like(data, dtype_like)
 
 

python/tvm/relay/qnn/op/legalizations.py

@@ -21,7 +21,7 @@ from __future__ import absolute_import
 import tvm
 from tvm import relay
 from .. import op as reg
-from ...util import get_scalar_from_constant
+from ...frontend.util import get_scalar_from_constant
 
 #################################################
 # Register the functions for different operators.

python/tvm/relay/quantize/_annotate.py

@@ -24,7 +24,6 @@ from .. import expr as _expr
 from .. import analysis as _analysis
 from .. import op as _op
 from ..op import op as _reg
-from ..base import register_relay_node
 from . import _quantize
 from .quantize import QAnnotateKind, current_qconfig, quantize_context
 from .quantize import _forward_op
@@ -58,7 +57,7 @@ _reg.register_pattern("relay.op.annotation.simulated_quantize",
 _reg.register_injective_schedule("annotation.cast_hint")
 
 
-@register_relay_node
+@tvm._ffi.register_object("relay.QAnnotateExpr")
 class QAnnotateExpr(_expr.TempExpr):
     """A special kind of Expr for Annotating.
 

python/tvm/relay/quantize/_partition.py

@@ -19,7 +19,6 @@
 import tvm
 from .. import expr as _expr
 from .. import analysis as _analysis
-from ..base import register_relay_node
 from ..op import op as _reg
 from . import _quantize
 from .quantize import _forward_op
@@ -30,7 +29,7 @@ def register_partition_function(op_name, frewrite=None, level=10):
     return _register(frewrite) if frewrite is not None else _register
 
 
-@register_relay_node
+@tvm._ffi.register_object("relay.QPartitionExpr")
 class QPartitionExpr(_expr.TempExpr):
     def __init__(self, expr):
         self.__init_handle_by_constructor__(

python/tvm/relay/quantize/quantize.py

@@ -17,12 +17,12 @@
 #pylint: disable=unused-argument, not-context-manager
 """Automatic quantization toolkit."""
 import tvm.ir
+from tvm.runtime import Object
 
 from . import _quantize
 from ._calibrate import calibrate
 from .. import expr as _expr
 from .. import transform as _transform
-from ..base import Object, register_relay_node
 
 
 class QAnnotateKind(object):
@@ -52,7 +52,7 @@ def _forward_op(ref_call, args):
         ref_call.op, args, ref_call.attrs, ref_call.type_args)
 
 
-@register_relay_node("relay.quantize.QConfig")
+@tvm._ffi.register_object("relay.quantize.QConfig")
 class QConfig(Object):
     """Configure the quantization behavior by setting config variables.
 

python/tvm/relay/_expr.py → python/tvm/relay/transform/__init__.py

@@ -14,8 +14,9 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-# pylint: disable=no-else-return, unidiomatic-typecheck, undefined-variable
-"""The interface of expr function exposed from C++."""
-import tvm._ffi
+# pylint: disable=wildcard-import, redefined-builtin, invalid-name
+"""The Relay IR namespace containing transformations."""
+# transformation passes
+from .transform import *
 
-tvm._ffi._init_api("relay._expr", __name__)
+from . import memory_alloc

python/tvm/relay/_transform.py → python/tvm/relay/transform/_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 Relay type inference and checking."""
+"""FFI APIs for Relay transformation passes."""
 import tvm._ffi
 
 tvm._ffi._init_api("relay._transform", __name__)

python/tvm/relay/memory_alloc.py → python/tvm/relay/transform/memory_alloc.py

@@ -19,12 +19,13 @@
 A pass for manifesting explicit memory allocations.
 """
 import numpy as np
-from .expr_functor import ExprMutator
-from .scope_builder import ScopeBuilder
+from ..expr_functor import ExprMutator
+from ..scope_builder import ScopeBuilder
 from . import transform
-from . import op, ty, expr
-from .. import DataType, register_func
-from .backend import compile_engine
+from .. import op
+from ... import DataType, register_func
+from .. import ty, expr
+from ..backend import compile_engine
 
 
 def is_primitive(call):

python/tvm/relay/ty.py

@@ -20,10 +20,10 @@ from tvm.ir import Type, TypeKind, TypeVar, GlobalTypeVar
 from tvm.ir import TypeConstraint, FuncType, TupleType, IncompleteType
 from tvm.ir import TypeCall, TypeRelation, TensorType, RelayRefType as RefType
 
-from .base import RelayNode, register_relay_node
-from . import _make
+from .base import RelayNode
+from . import _ffi_api
 
-Any = _make.Any
+Any = _ffi_api.Any
 
 def type_has_any(tensor_type):
     """Check whether type has any as a shape.
@@ -36,7 +36,7 @@ def type_has_any(tensor_type):
     has_any : bool
         The check result.
     """
-    return _make.IsDynamic(tensor_type)
+    return _ffi_api.IsDynamic(tensor_type)
 
 
 def ShapeVar(name):

python/tvm/relay/vision.py

-# 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, unused-import, unused-wildcard-import
-"""Vision network related operators."""
-# Re-export in a specific file name so that autodoc can pick it up
-from .op.vision import *

src/relay/analysis/alpha_equal.cc

@@ -581,7 +581,7 @@ bool AlphaEqual(const Expr& lhs, const Expr& rhs) {
   return AlphaEqualHandler(false, false).ExprEqual(lhs, rhs);
 }
 
-TVM_REGISTER_GLOBAL("relay._analysis._alpha_equal")
+TVM_REGISTER_GLOBAL("relay.analysis._alpha_equal")
 .set_body_typed([](ObjectRef a, ObjectRef b) {
   return AlphaEqualHandler(false, false).Equal(a, b);
 });
@@ -591,18 +591,18 @@ TVM_REGISTER_GLOBAL("ir.type_alpha_equal")
   return AlphaEqual(a, b);
 });
 
-TVM_REGISTER_GLOBAL("relay._analysis._assert_alpha_equal")
+TVM_REGISTER_GLOBAL("relay.analysis._assert_alpha_equal")
 .set_body_typed([](ObjectRef a, ObjectRef b) {
   bool alpha_equal = AlphaEqualHandler(false, true).Equal(a, b);
   CHECK(alpha_equal) << AsText(a, true) << " and " << AsText(b, true) << " are not alpha equal";
 });
 
-TVM_REGISTER_GLOBAL("relay._analysis._graph_equal")
+TVM_REGISTER_GLOBAL("relay.analysis._graph_equal")
 .set_body_typed([](ObjectRef a, ObjectRef b) {
   return AlphaEqualHandler(true, false).Equal(a, b);
 });
 
-TVM_REGISTER_GLOBAL("relay._analysis._assert_graph_equal")
+TVM_REGISTER_GLOBAL("relay.analysis._assert_graph_equal")
 .set_body_typed([](ObjectRef a, ObjectRef b) {
   bool graph_equal = AlphaEqualHandler(true, true).Equal(a, b);
   CHECK(graph_equal) << AsText(a, true) << " and " << AsText(b, true) << " are not graph equal";

src/relay/analysis/call_graph.cc

@@ -299,24 +299,24 @@ TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
   p->stream << "CallGraph: \n" << GetRef<CallGraph>(node);
 });
 
-TVM_REGISTER_GLOBAL("relay._analysis.CallGraph")
+TVM_REGISTER_GLOBAL("relay.analysis.CallGraph")
 .set_body_typed([](IRModule module) {
   return CallGraph(module);
 });
 
-TVM_REGISTER_GLOBAL("relay._analysis.PrintCallGraph")
+TVM_REGISTER_GLOBAL("relay.analysis.PrintCallGraph")
 .set_body_typed([](CallGraph call_graph) {
   std::stringstream ss;
   ss << call_graph;
   return ss.str();
 });
 
-TVM_REGISTER_GLOBAL("relay._analysis.GetModule")
+TVM_REGISTER_GLOBAL("relay.analysis.GetModule")
 .set_body_typed([](CallGraph call_graph) {
   return call_graph->module;
 });
 
-TVM_REGISTER_GLOBAL("relay._analysis.PrintCallGraphGlobalVar")
+TVM_REGISTER_GLOBAL("relay.analysis.PrintCallGraphGlobalVar")
 .set_body_typed([](CallGraph call_graph, GlobalVar var) {
   const auto* entry_node = call_graph[var];
   std::stringstream ss;
@@ -324,19 +324,19 @@ TVM_REGISTER_GLOBAL("relay._analysis.PrintCallGraphGlobalVar")
   return ss.str();
 });
 
-TVM_REGISTER_GLOBAL("relay._analysis.GetRefCountGlobalVar")
+TVM_REGISTER_GLOBAL("relay.analysis.GetRefCountGlobalVar")
 .set_body_typed([](CallGraph call_graph, GlobalVar var) {
   const auto* entry_node = call_graph[var];
   return static_cast<int>(entry_node->GetRefCount());
 });
 
-TVM_REGISTER_GLOBAL("relay._analysis.GetGlobalVarCallCount")
+TVM_REGISTER_GLOBAL("relay.analysis.GetGlobalVarCallCount")
 .set_body_typed([](CallGraph call_graph, GlobalVar var) {
   const auto* entry_node = call_graph[var];
   return static_cast<int>(entry_node->size());
 });
 
-TVM_REGISTER_GLOBAL("relay._analysis.IsRecursive")
+TVM_REGISTER_GLOBAL("relay.analysis.IsRecursive")
 .set_body_typed([](CallGraph call_graph, GlobalVar var) {
   const auto* entry_node = call_graph[var];
   return entry_node->IsRecursive();

src/relay/analysis/extract_fused_functions.cc

@@ -74,7 +74,7 @@ Pass ExtractFusedFunctions() {
                     "ExtractFusedFunctions");
 }
 
-TVM_REGISTER_GLOBAL("relay._analysis.ExtractFusedFunctions").set_body_typed(ExtractFusedFunctions);
+TVM_REGISTER_GLOBAL("relay.analysis.ExtractFusedFunctions").set_body_typed(ExtractFusedFunctions);
 
 }  // namespace transform
 

src/relay/analysis/feature.cc

@@ -104,7 +104,7 @@ Array<Integer> PyDetectFeature(const Expr& expr, const IRModule& mod) {
   return static_cast<Array<Integer>>(fs);
 }
 
-TVM_REGISTER_GLOBAL("relay._analysis.detect_feature")
+TVM_REGISTER_GLOBAL("relay.analysis.detect_feature")
 .set_body_typed(PyDetectFeature);
 
 }  // namespace relay

src/relay/analysis/kind_check.cc

@@ -186,7 +186,7 @@ Kind KindCheck(const Type& t, const IRModule& mod) {
   return kc.Check(t);
 }
 
-TVM_REGISTER_GLOBAL("relay._analysis.check_kind")
+TVM_REGISTER_GLOBAL("relay.analysis.check_kind")
 .set_body([](TVMArgs args, TVMRetValue* ret) {
     if (args.size() == 1) {
       *ret = KindCheck(args[0], IRModule({}, {}));

src/relay/analysis/mac_count.cc

@@ -206,7 +206,7 @@ int64_t GetTotalMacNumber(const Expr& expr) {
   return MacCounter::GetTotalMacNumber(expr);
 }
 
-TVM_REGISTER_GLOBAL("relay._analysis.GetTotalMacNumber")
+TVM_REGISTER_GLOBAL("relay.analysis.GetTotalMacNumber")
 .set_body_typed(GetTotalMacNumber);
 
 }  // namespace mac_count

src/relay/analysis/match_exhaustion.cc

@@ -310,7 +310,7 @@ Array<Pattern> UnmatchedCases(const Match& match, const IRModule& mod) {
 }
 
 // expose for testing only
-TVM_REGISTER_GLOBAL("relay._analysis.unmatched_cases")
+TVM_REGISTER_GLOBAL("relay.analysis.unmatched_cases")
 .set_body_typed(
   [](const Match& match, const IRModule& mod_ref) {
     IRModule call_mod = mod_ref;

src/relay/analysis/type_solver.cc

@@ -659,7 +659,7 @@ bool TypeSolver::Solve() {
 }
 
 // Expose type solver only for debugging purposes.
-TVM_REGISTER_GLOBAL("relay._analysis._test_type_solver")
+TVM_REGISTER_GLOBAL("relay.analysis._test_type_solver")
 .set_body([](runtime::TVMArgs args, runtime::TVMRetValue* ret) {
     using runtime::PackedFunc;
     using runtime::TypedPackedFunc;

src/relay/analysis/util.cc

@@ -274,10 +274,10 @@ tvm::Array<Var> AllVars(const Expr& expr) {
   return VarVisitor().All(expr);
 }
 
-TVM_REGISTER_GLOBAL("relay._analysis.free_vars")
+TVM_REGISTER_GLOBAL("relay.analysis.free_vars")
 .set_body_typed(FreeVars);
 
-TVM_REGISTER_GLOBAL("relay._analysis.bound_vars")
+TVM_REGISTER_GLOBAL("relay.analysis.bound_vars")
   .set_body([](TVMArgs args, TVMRetValue* ret) {
       ObjectRef x = args[0];
       if (x.as<ExprNode>()) {
@@ -287,10 +287,10 @@ TVM_REGISTER_GLOBAL("relay._analysis.bound_vars")
       }
     });
 
-TVM_REGISTER_GLOBAL("relay._analysis.all_vars")
+TVM_REGISTER_GLOBAL("relay.analysis.all_vars")
 .set_body_typed(AllVars);
 
-TVM_REGISTER_GLOBAL("relay._analysis.free_type_vars")
+TVM_REGISTER_GLOBAL("relay.analysis.free_type_vars")
 .set_body([](TVMArgs args, TVMRetValue* ret) {
     ObjectRef x = args[0];
     IRModule mod = args[1];
@@ -301,7 +301,7 @@ TVM_REGISTER_GLOBAL("relay._analysis.free_type_vars")
     }
   });
 
-TVM_REGISTER_GLOBAL("relay._analysis.bound_type_vars")
+TVM_REGISTER_GLOBAL("relay.analysis.bound_type_vars")
   .set_body([](TVMArgs args, TVMRetValue* ret) {
       ObjectRef x = args[0];
       IRModule mod = args[1];
@@ -312,7 +312,7 @@ TVM_REGISTER_GLOBAL("relay._analysis.bound_type_vars")
       }
     });
 
-TVM_REGISTER_GLOBAL("relay._analysis.all_type_vars")
+TVM_REGISTER_GLOBAL("relay.analysis.all_type_vars")
   .set_body([](TVMArgs args, TVMRetValue* ret) {
       ObjectRef x = args[0];
       IRModule mod = args[1];

src/relay/analysis/well_formed.cc

@@ -125,7 +125,7 @@ bool WellFormed(const Expr& e) {
   return WellFormedChecker().CheckWellFormed(e);
 }
 
-TVM_REGISTER_GLOBAL("relay._analysis.well_formed")
+TVM_REGISTER_GLOBAL("relay.analysis.well_formed")
 .set_body_typed(WellFormed);
 
 }  // namespace relay

src/relay/backend/compile_engine.cc

@@ -86,7 +86,7 @@ bool IsDynamic(const Type& ty) {
 }
 
 // TODO(@jroesch): MOVE ME
-TVM_REGISTER_GLOBAL("relay._make.IsDynamic")
+TVM_REGISTER_GLOBAL("relay.ir.IsDynamic")
 .set_body_typed(IsDynamic);
 
 Array<IndexExpr> GetShape(const Array<IndexExpr>& shape) {

src/relay/ir/adt.cc

@@ -34,7 +34,7 @@ PatternWildcard PatternWildcardNode::make() {
 
 TVM_REGISTER_NODE_TYPE(PatternWildcardNode);
 
-TVM_REGISTER_GLOBAL("relay._make.PatternWildcard")
+TVM_REGISTER_GLOBAL("relay.ir.PatternWildcard")
 .set_body_typed(PatternWildcardNode::make);
 
 TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
@@ -50,7 +50,7 @@ PatternVar PatternVarNode::make(tvm::relay::Var var) {
 
 TVM_REGISTER_NODE_TYPE(PatternVarNode);
 
-TVM_REGISTER_GLOBAL("relay._make.PatternVar")
+TVM_REGISTER_GLOBAL("relay.ir.PatternVar")
 .set_body_typed(PatternVarNode::make);
 
 TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
@@ -69,7 +69,7 @@ PatternConstructor PatternConstructorNode::make(Constructor constructor,
 
 TVM_REGISTER_NODE_TYPE(PatternConstructorNode);
 
-TVM_REGISTER_GLOBAL("relay._make.PatternConstructor")
+TVM_REGISTER_GLOBAL("relay.ir.PatternConstructor")
 .set_body_typed(PatternConstructorNode::make);
 
 TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
@@ -87,7 +87,7 @@ PatternTuple PatternTupleNode::make(tvm::Array<Pattern> patterns) {
 
 TVM_REGISTER_NODE_TYPE(PatternTupleNode);
 
-TVM_REGISTER_GLOBAL("relay._make.PatternTuple")
+TVM_REGISTER_GLOBAL("relay.ir.PatternTuple")
 .set_body_typed(PatternTupleNode::make);
 
 TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
@@ -105,7 +105,7 @@ Clause ClauseNode::make(Pattern lhs, Expr rhs) {
 
 TVM_REGISTER_NODE_TYPE(ClauseNode);
 
-TVM_REGISTER_GLOBAL("relay._make.Clause")
+TVM_REGISTER_GLOBAL("relay.ir.Clause")
 .set_body_typed(ClauseNode::make);
 
 TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
@@ -125,7 +125,7 @@ Match MatchNode::make(Expr data, tvm::Array<Clause> clauses, bool complete) {
 
 TVM_REGISTER_NODE_TYPE(MatchNode);
 
-TVM_REGISTER_GLOBAL("relay._make.Match")
+TVM_REGISTER_GLOBAL("relay.ir.Match")
 .set_body_typed(MatchNode::make);
 
 TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)

src/relay/ir/expr.cc

@@ -38,7 +38,7 @@ Constant ConstantNode::make(runtime::NDArray data) {
 
 TVM_REGISTER_NODE_TYPE(ConstantNode);
 
-TVM_REGISTER_GLOBAL("relay._make.Constant")
+TVM_REGISTER_GLOBAL("relay.ir.Constant")
 .set_body_typed(ConstantNode::make);
 
 TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
@@ -71,7 +71,7 @@ Tuple TupleNode::make(tvm::Array<relay::Expr> fields) {
 
 TVM_REGISTER_NODE_TYPE(TupleNode);
 
-TVM_REGISTER_GLOBAL("relay._make.Tuple")
+TVM_REGISTER_GLOBAL("relay.ir.Tuple")
 .set_body_typed(TupleNode::make);
 
 TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
@@ -96,7 +96,7 @@ Var VarNode::make(std::string name_hint, Type type_annotation) {
 
 TVM_REGISTER_NODE_TYPE(VarNode);
 
-TVM_REGISTER_GLOBAL("relay._make.Var")
+TVM_REGISTER_GLOBAL("relay.ir.Var")
 .set_body_typed(static_cast<Var (*)(std::string, Type)>(VarNode::make));
 
 TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
@@ -123,7 +123,7 @@ Call CallNode::make(Expr op, Array<Expr> args, Attrs attrs,
 
 TVM_REGISTER_NODE_TYPE(CallNode);
 
-TVM_REGISTER_GLOBAL("relay._make.Call")
+TVM_REGISTER_GLOBAL("relay.ir.Call")
 .set_body_typed(CallNode::make);
 
 TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
@@ -143,7 +143,7 @@ Let LetNode::make(Var var, Expr value, Expr body) {
 
 TVM_REGISTER_NODE_TYPE(LetNode);
 
-TVM_REGISTER_GLOBAL("relay._make.Let")
+TVM_REGISTER_GLOBAL("relay.ir.Let")
 .set_body_typed(LetNode::make);
 
 TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
@@ -163,7 +163,7 @@ If IfNode::make(Expr cond, Expr true_branch, Expr false_branch) {
 
 TVM_REGISTER_NODE_TYPE(IfNode);
 
-TVM_REGISTER_GLOBAL("relay._make.If")
+TVM_REGISTER_GLOBAL("relay.ir.If")
 .set_body_typed(IfNode::make);
 
 TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
@@ -182,7 +182,7 @@ TupleGetItem TupleGetItemNode::make(Expr tuple, int index) {
 
 TVM_REGISTER_NODE_TYPE(TupleGetItemNode);
 
-TVM_REGISTER_GLOBAL("relay._make.TupleGetItem")
+TVM_REGISTER_GLOBAL("relay.ir.TupleGetItem")
 .set_body_typed(TupleGetItemNode::make);
 
 TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
@@ -199,7 +199,7 @@ RefCreate RefCreateNode::make(Expr value) {
 
 TVM_REGISTER_NODE_TYPE(RefCreateNode);
 
-TVM_REGISTER_GLOBAL("relay._make.RefCreate")
+TVM_REGISTER_GLOBAL("relay.ir.RefCreate")
 .set_body_typed(RefCreateNode::make);
 
 TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
@@ -216,7 +216,7 @@ RefRead RefReadNode::make(Expr ref) {
 
 TVM_REGISTER_NODE_TYPE(RefReadNode);
 
-TVM_REGISTER_GLOBAL("relay._make.RefRead")
+TVM_REGISTER_GLOBAL("relay.ir.RefRead")
 .set_body_typed(RefReadNode::make);
 
 TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
@@ -234,7 +234,7 @@ RefWrite RefWriteNode::make(Expr ref, Expr value) {
 
 TVM_REGISTER_NODE_TYPE(RefWriteNode);
 
-TVM_REGISTER_GLOBAL("relay._make.RefWrite")
+TVM_REGISTER_GLOBAL("relay.ir.RefWrite")
 .set_body_typed(RefWriteNode::make);
 
 TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
@@ -243,12 +243,12 @@ TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
   p->stream << "RefWriteNode(" << node->ref << ", " << node->value << ")";
 });
 
-TVM_REGISTER_GLOBAL("relay._expr.TempExprRealize")
+TVM_REGISTER_GLOBAL("relay.ir.TempExprRealize")
 .set_body_typed([](TempExpr temp) {
   return temp->Realize();
 });
 
-TVM_REGISTER_GLOBAL("relay._make.Any")
+TVM_REGISTER_GLOBAL("relay.ir.Any")
 .set_body_typed([]() { return Any::make(); });
 
 }  // namespace relay

src/relay/ir/expr_functor.cc

@@ -347,7 +347,7 @@ void PostOrderVisit(const Expr& e, std::function<void(const Expr&)> fvisit) {
   ExprApplyVisit(fvisit).VisitExpr(e);
 }
 
-TVM_REGISTER_GLOBAL("relay._analysis.post_order_visit")
+TVM_REGISTER_GLOBAL("relay.analysis.post_order_visit")
 .set_body_typed([](Expr expr, PackedFunc f) {
     PostOrderVisit(expr, [f](const Expr& n) {
         f(n);
@@ -443,7 +443,7 @@ Expr Bind(const Expr& expr, const tvm::Map<Var, Expr>& args_map) {
   }
 }
 
-TVM_REGISTER_GLOBAL("relay._expr.Bind")
+TVM_REGISTER_GLOBAL("relay.ir.Bind")
 .set_body([](TVMArgs args, TVMRetValue* ret) {
     ObjectRef input = args[0];
     if (input->IsInstance<ExprNode>()) {

src/relay/ir/function.cc

@@ -62,7 +62,7 @@ bool FunctionNode::UseDefaultCompiler() const {
 
 TVM_REGISTER_NODE_TYPE(FunctionNode);
 
-TVM_REGISTER_GLOBAL("relay._make.Function")
+TVM_REGISTER_GLOBAL("relay.ir.Function")
 .set_body_typed([](tvm::Array<Var> params,
                    Expr body,
                    Type ret_type,
@@ -80,7 +80,7 @@ TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
             << node->attrs << ")";
 });
 
-TVM_REGISTER_GLOBAL("relay._expr.FunctionWithAttr")
+TVM_REGISTER_GLOBAL("relay.ir.FunctionWithAttr")
 .set_body_typed(
     [](Function func, std::string name, ObjectRef ref) {
       return WithAttr(std::move(func), name, ref);

src/relay/ir/hash.cc

@@ -423,12 +423,12 @@ size_t StructuralHash::operator()(const Expr& expr) const {
   return RelayHashHandler().ExprHash(expr);
 }
 
-TVM_REGISTER_GLOBAL("relay._analysis._expr_hash")
+TVM_REGISTER_GLOBAL("relay.analysis._expr_hash")
 .set_body_typed([](ObjectRef ref) {
   return static_cast<int64_t>(RelayHashHandler().Hash(ref));
 });
 
-TVM_REGISTER_GLOBAL("relay._analysis._type_hash")
+TVM_REGISTER_GLOBAL("relay.analysis._type_hash")
 .set_body_typed([](Type type) {
   return static_cast<int64_t>(RelayHashHandler().TypeHash(type));
 });

src/relay/op/tensor/transform.cc

@@ -82,7 +82,7 @@ Expr MakeCast(Expr data,
   return CallNode::make(op, {data}, Attrs(attrs), {});
 }
 
-TVM_REGISTER_GLOBAL("relay._make.cast")
+TVM_REGISTER_GLOBAL("relay.ir.cast")
 .set_body_typed(MakeCast);
 
 RELAY_REGISTER_OP("cast")
@@ -138,7 +138,7 @@ Expr MakeCastLike(Expr data,
 }
 
 
-TVM_REGISTER_GLOBAL("relay._make.cast_like")
+TVM_REGISTER_GLOBAL("relay.ir.cast_like")
 .set_body_typed(MakeCastLike);
 
 RELAY_REGISTER_OP("cast_like")

src/relay/transforms/device_annotation.cc

@@ -560,10 +560,10 @@ Map<Expr, Integer> CollectDeviceAnnotationOps(const Expr& expr) {
   return AnnotatationVisitor::GetAnnotations(expr);
 }
 
-TVM_REGISTER_GLOBAL("relay._analysis.CollectDeviceInfo")
+TVM_REGISTER_GLOBAL("relay.analysis.CollectDeviceInfo")
 .set_body_typed(CollectDeviceInfo);
 
-TVM_REGISTER_GLOBAL("relay._analysis.CollectDeviceAnnotationOps")
+TVM_REGISTER_GLOBAL("relay.analysis.CollectDeviceAnnotationOps")
 .set_body_typed(CollectDeviceAnnotationOps);
 
 namespace transform {

src/relay/transforms/fold_constant.cc

@@ -73,7 +73,7 @@ bool ConstantCheck(const Expr& e) {
   return ConstantChecker().Check(e);
 }
 
-TVM_REGISTER_GLOBAL("relay._analysis.check_constant")
+TVM_REGISTER_GLOBAL("relay.analysis.check_constant")
 .set_body_typed(ConstantCheck);
 
 // TODO(tvm-team) consider combine dead-code with constant folder.

tests/python/relay/test_feature.py → tests/python/relay/test_analysis_feature.py

@@ -18,9 +18,8 @@
 import tvm
 from tvm import te
 from tvm import relay
-from tvm.relay.analysis import detect_feature
+from tvm.relay.analysis import detect_feature, Feature
 from tvm.relay.transform import gradient
-from tvm.relay.feature import Feature
 from tvm.relay.prelude import Prelude
 from tvm.relay.testing import run_infer_type
 

tests/python/relay/test_backend_graph_runtime.py

@@ -17,10 +17,8 @@
 import numpy as np
 
 import tvm
-from tvm import te
 from tvm import relay
 from tvm.contrib import graph_runtime
-from tvm.relay.scope_builder import ScopeBuilder
 from tvm.relay.op import add
 from tvm.relay.testing.config import ctx_list
 

tests/python/relay/test_call_graph.py

@@ -25,7 +25,7 @@ def test_callgraph_construct():
     x = relay.var("x", shape=(2, 3))
     y = relay.var("y", shape=(2, 3))
     mod["g1"] = relay.Function([x, y], x + y)
-    call_graph = relay.CallGraph(mod)
+    call_graph = relay.analysis.CallGraph(mod)
     assert "g1" in str(call_graph)
     assert relay.alpha_equal(mod, call_graph.module)
 
@@ -38,7 +38,7 @@ def test_print_element():
     x1 = relay.var("x1", shape=(2, 3))
     y1 = relay.var("y1", shape=(2, 3))
     mod["g1"] = relay.Function([x1, y1], x1 - y1)
-    call_graph = relay.CallGraph(mod)
+    call_graph = relay.analysis.CallGraph(mod)
 
     assert "#refs = 0" in str(call_graph.print_var("g0"))
     assert "#refs = 0" in str(call_graph.print_var("g1"))
@@ -54,13 +54,13 @@ def test_global_call_count():
     y1 = relay.var("y1", shape=(2, 3))
     g1 = relay.GlobalVar("g1")
     mod[g1] = relay.Function([x1, y1], g0(x1, y1))
-    call_graph = relay.CallGraph(mod)
+    call_graph = relay.analysis.CallGraph(mod)
 
     p0 = relay.var("p0", shape=(2, 3))
     p1 = relay.var("p1", shape=(2, 3))
     func = relay.Function([p0, p1], g0(p0, p1) * g1(p0, p1))
     mod["main"] = func
-    call_graph = relay.CallGraph(mod)
+    call_graph = relay.analysis.CallGraph(mod)
 
     assert call_graph.global_call_count(g0) == 0
     assert call_graph.global_call_count(g1) == 1
@@ -77,13 +77,13 @@ def test_ref_count():
     y1 = relay.var("y1", shape=(2, 3))
     g1 = relay.GlobalVar("g1")
     mod[g1] = relay.Function([x1, y1], x1 - y1)
-    call_graph = relay.CallGraph(mod)
+    call_graph = relay.analysis.CallGraph(mod)
 
     p0 = relay.var("p0", shape=(2, 3))
     p1 = relay.var("p1", shape=(2, 3))
     func = relay.Function([p0, p1], g0(p0, p1) * g1(p0, p1))
     mod["main"] = func
-    call_graph = relay.CallGraph(mod)
+    call_graph = relay.analysis.CallGraph(mod)
 
     assert call_graph.ref_count(g0) == 1
     assert call_graph.ref_count(g1) == 1
@@ -100,13 +100,13 @@ def test_nested_ref():
     y1 = relay.var("y1", shape=(2, 3))
     g1 = relay.GlobalVar("g1")
     mod[g1] = relay.Function([x1, y1], g0(x1, y1))
-    call_graph = relay.CallGraph(mod)
+    call_graph = relay.analysis.CallGraph(mod)
 
     p0 = relay.var("p0", shape=(2, 3))
     p1 = relay.var("p1", shape=(2, 3))
     func = relay.Function([p0, p1], g0(p0, p1) * g1(p0, p1))
     mod["main"] = func
-    call_graph = relay.CallGraph(mod)
+    call_graph = relay.analysis.CallGraph(mod)
 
     assert call_graph.ref_count(g0) == 2
     assert call_graph.ref_count(g1) == 1
@@ -138,7 +138,7 @@ def test_recursive_func():
     mod[sum_up] = func
     iarg = relay.var('i', shape=[], dtype='int32')
     mod["main"] = relay.Function([iarg], sum_up(iarg))
-    call_graph = relay.CallGraph(mod)
+    call_graph = relay.analysis.CallGraph(mod)
 
     assert call_graph.is_recursive(sum_up)
     assert call_graph.ref_count(sum_up) == 2

tests/python/relay/test_pass_annotation.py

@@ -19,7 +19,6 @@ import json
 import numpy as np
 
 import tvm
-from tvm import te
 from tvm import relay
 from tvm.contrib import graph_runtime
 from tvm.relay.expr_functor import ExprMutator

tests/python/relay/test_memory_alloc.py → tests/python/relay/test_pass_memory_alloc.py

@@ -18,7 +18,7 @@ import tvm
 from tvm import te
 import numpy as np
 from tvm import relay
-from tvm.relay import memory_alloc
+from tvm.relay.transform import memory_alloc
 
 def check_vm_alloc(func, check_fn):
     mod = tvm.IRModule()

tests/python/relay/test_pass_partition_graph.py

@@ -25,7 +25,7 @@ from tvm import relay
 from tvm import runtime
 from tvm.relay import transform
 from tvm.contrib import util
-from tvm.relay.annotation import compiler_begin, compiler_end
+from tvm.relay.op.annotation import compiler_begin, compiler_end
 from tvm.relay.expr_functor import ExprMutator
 
 # Leverage the pass manager to write a simple white list based annotator

tests/python/relay/test_pass_to_a_normal_form.py

@@ -22,7 +22,7 @@ from tvm.relay.analysis import alpha_equal, detect_feature
 from tvm.relay import op, create_executor, transform
 from tvm.relay.prelude import Prelude
 from tvm.relay.testing import add_nat_definitions, count
-from tvm.relay.feature import Feature
+from tvm.relay.analysis import Feature
 
 
 def run_opt_pass(expr, passes):

tests/python/relay/test_pass_to_cps.py

@@ -16,15 +16,14 @@
 # under the License.
 import numpy as np
 import tvm
-from tvm import te
 from tvm import relay
 from tvm.relay.analysis import alpha_equal, detect_feature
 from tvm.relay.transform import to_cps, un_cps
-from tvm.relay.feature import Feature
+from tvm.relay.analysis import Feature
 from tvm.relay.prelude import Prelude
 from tvm.relay.testing import add_nat_definitions, make_nat_expr, rand, run_infer_type, run_opt_pass
 from tvm.relay import create_executor
-from tvm.relay import Function, transform
+from tvm.relay import transform
 
 
 def test_id():

tests/python/relay/test_pass_to_graph_normal_form.py

@@ -16,9 +16,9 @@
 # under the License.
 import numpy as np
 import tvm
-from tvm import te
 from tvm import relay
-from tvm.relay import op, create_executor, transform, Feature
+from tvm.relay import op, create_executor, transform
+from tvm.relay.analysis import Feature
 from tvm.relay.analysis import detect_feature
 
 

tests/python/relay/test_type_solver.py

@@ -15,7 +15,6 @@
 # specific language governing permissions and limitations
 # under the License.
 import tvm
-from tvm import te
 from tvm import relay
 import pytest
 
@@ -27,7 +26,7 @@ def make_rel(name, args, num_inputs=None, attrs=None):
     return relay.ty.TypeRelation(func, args, num_inputs, attrs)
 
 def make_solver():
-    solver = relay._analysis._test_type_solver()
+    solver = relay.analysis._ffi_api._test_type_solver()
     solver.Solve = solver("Solve")
     solver.Unify = solver("Unify")
     solver.Resolve = solver("Resolve")

tests/python/relay/test_vm.py

@@ -18,7 +18,6 @@ import numpy as np
 import pytest
 
 import tvm
-from tvm import te
 from tvm import runtime
 from tvm import relay
 from tvm.relay.scope_builder import ScopeBuilder