[RELAY] Remove re-exports of tvm.transform (#5337)

docs/api/python/ir.rst

@@ -21,3 +21,11 @@ tvm.ir
    :members:
    :imported-members:
    :autosummary:
+
+
+tvm.transform
+-------------
+.. automodule:: tvm.transform
+   :members:
+   :imported-members:
+   :autosummary:

docs/dev/convert_layout.rst

@@ -227,7 +227,7 @@ ConvertLayout pass is extremely easy to use. The pass is not a part of default r
 
     # Convert the layout to NCHW
     # RemoveUnunsedFunctions is used to clean up the graph.
-    seq = relay.transform.Sequential([relay.transform.RemoveUnusedFunctions(),
+    seq = tvm.transform.Sequential([relay.transform.RemoveUnusedFunctions(),
                                       relay.transform.ConvertLayout('NCHW')])
     with relay.transform.PassContext(opt_level=3):
         mod = seq(mod)

docs/dev/relay_pass_infra.rst

@@ -582,7 +582,7 @@ using ``Sequential`` associated with other types of passes.
     func = relay.Function([x], z2)
 
     # Customize the optimization pipeline.
-    seq = _transform.Sequential([
+    seq = tvm.transform.Sequential([
         relay.transform.InferType(),
         relay.transform.FoldConstant(),
         relay.transform.EliminateCommonSubexpr(),
@@ -609,7 +609,7 @@ sequential pass example could be like the following to enable IR dumping for
 
 .. code:: python
 
-    seq = _transform.Sequential([
+    seq = tvm.transform.Sequential([
         relay.transform.InferType(),
         relay.transform.FoldConstant(),
         relay.transform.PrintIR(),

include/tvm/ir/transform.h

@@ -361,9 +361,11 @@ TVM_DLL Pass CreateModulePass(
 
 /*!
  * \brief A special trace pass that prints the header and IR to LOG(INFO).
+ * \param header The header to be attached to the output.
+ * \param show_meta_data Whether should we show meta data.
  * \return The pass.
  */
-TVM_DLL Pass PrintIR(std::string header);
+TVM_DLL Pass PrintIR(std::string header = "", bool show_meta_data = false);
 
 }  // namespace transform
 }  // namespace tvm

python/tvm/ir/json_compact.py

@@ -106,7 +106,7 @@ def create_updater_06_to_07():
         "relay.PassInfo": _rename("transform.PassInfo"),
         "relay.PassContext": _rename("transform.PassContext"),
         "relay.ModulePass": _rename("transform.ModulePass"),
-        "relay.Sequantial": _rename("transform.Sequantial"),
+        "relay.Sequential": _rename("transform.Sequential"),
         # TIR
         "Variable": _update_tir_var("tir.Var"),
         "SizeVar": _update_tir_var("tir.SizeVar"),

python/tvm/ir/transform.py

@@ -329,7 +329,7 @@ def module_pass(pass_func=None, opt_level=None, name=None, required=None):
     return create_module_pass
 
 
-def PrintIR(header):
+def PrintIR(header="", show_meta_data=False):
     """A special trace pass that prints the header and IR.
 
     Parameters
@@ -337,8 +337,11 @@ def PrintIR(header):
     header : str
         The header to be displayed along with the dump.
 
+    show_meta_data : bool
+        A boolean flag to indicate if meta data should be printed.
+
     Returns
     --------
     The pass
     """
-    return _ffi_transform_api.PrintIR(header)
+    return _ffi_transform_api.PrintIR(header, show_meta_data)

python/tvm/relay/__init__.py

@@ -128,20 +128,9 @@ Prelude = prelude.Prelude
 # Scope builder
 ScopeBuilder = scope_builder.ScopeBuilder
 
-module_pass = transform.module_pass
-function_pass = transform.function_pass
-
 # Parser
 fromtext = parser.fromtext
 
 # Param Serialization
 save_param_dict = param_dict.save_param_dict
 load_param_dict = param_dict.load_param_dict
-
-# Pass manager
-PassInfo = transform.PassInfo
-PassContext = transform.PassContext
-Pass = transform.Pass
-ModulePass = transform.ModulePass
-FunctionPass = transform.FunctionPass
-Sequential = transform.Sequential

python/tvm/relay/backend/interpreter.py

@@ -210,10 +210,10 @@ class Interpreter(Executor):
         opt_mod : tvm.IRModule
             The optimized module.
         """
-        seq = transform.Sequential([transform.SimplifyInference(),
-                                    transform.FuseOps(0),
-                                    transform.ToANormalForm(),
-                                    transform.InferType()])
+        seq = tvm.transform.Sequential([transform.SimplifyInference(),
+                                        transform.FuseOps(0),
+                                        transform.ToANormalForm(),
+                                        transform.InferType()])
         return seq(self.mod)
 
     def _make_executor(self, expr=None):

python/tvm/relay/qnn/transform.py

@@ -60,7 +60,7 @@ def CanonicalizeOps():
 
     Returns
     -------
-    ret : tvm.relay.Pass
+    ret : tvm.transform.Pass
         The registered pass that canonicalizes QNN ops to Relay ops.
     """
 
@@ -108,7 +108,7 @@ def Legalize():
 
     Returns
     -------
-    ret : tvm.relay.Pass
+    ret : tvm.transform.Pass
         The registered pass that legalizes QNN ops.
     """
 

python/tvm/relay/quantize/quantize.py

@@ -17,6 +17,7 @@
 #pylint: disable=unused-argument, not-context-manager
 """Automatic quantization toolkit."""
 import tvm.ir
+import tvm
 from tvm.runtime import Object
 
 from . import _quantize
@@ -240,7 +241,7 @@ def partition():
 
     Returns
     -------
-    ret: tvm.relay.Pass
+    ret: tvm.transform.Pass
         The registered pass for VTA rewrite.
     """
     return _quantize.QuantizePartition()
@@ -253,7 +254,7 @@ def annotate():
 
     Returns
     -------
-    ret: tvm.relay.Pass
+    ret: tvm.transform.Pass
         The registered pass for quantization annotation.
     """
     return _quantize.QuantizeAnnotate()
@@ -267,7 +268,7 @@ def realize():
 
     Returns
     -------
-    ret: tvm.relay.Pass
+    ret: tvm.transform.Pass
         The registered pass for quantization realization.
     """
     return _quantize.QuantizeRealize()
@@ -298,11 +299,12 @@ def prerequisite_optimize(mod, params=None):
     """ Prerequisite optimization passes for quantization. Perform
     "SimplifyInference", "FoldScaleAxis", "FoldConstant", and
     "CanonicalizeOps" optimization before quantization. """
-    optimize = _transform.Sequential([_transform.SimplifyInference(),
-                                      _transform.FoldConstant(),
-                                      _transform.FoldScaleAxis(),
-                                      _transform.CanonicalizeOps(),
-                                      _transform.FoldConstant()])
+    optimize = tvm.transform.Sequential(
+        [_transform.SimplifyInference(),
+         _transform.FoldConstant(),
+         _transform.FoldScaleAxis(),
+         _transform.CanonicalizeOps(),
+         _transform.FoldConstant()])
 
     if params:
         mod['main'] = _bind_params(mod['main'], params)
@@ -336,19 +338,20 @@ def quantize(mod, params=None, dataset=None):
     """
     mod = prerequisite_optimize(mod, params)
 
-    calibrate_pass = _transform.module_pass(calibrate(dataset), opt_level=1,
-                                            name="QuantizeCalibrate")
+    calibrate_pass = tvm.transform.module_pass(
+        calibrate(dataset), opt_level=1,
+        name="QuantizeCalibrate")
     quant_passes = [partition(),
                     annotate(),
                     calibrate_pass]
     if not current_qconfig().do_simulation:
         quant_passes.append(realize())
     quant_passes.append(_transform.FoldConstant())
-    quantize_seq = _transform.Sequential(quant_passes)
-    with _transform.PassContext(opt_level=3,
-                                required_pass=["QuantizeAnnotate",
-                                               "QuantizeCalibrate",
-                                               "QuantizeRealize"]):
+    quantize_seq = tvm.transform.Sequential(quant_passes)
+    with tvm.transform.PassContext(opt_level=3,
+                                   required_pass=["QuantizeAnnotate",
+                                                  "QuantizeCalibrate",
+                                                  "QuantizeRealize"]):
         with quantize_context():
             mod = quantize_seq(mod)
 

python/tvm/relay/testing/__init__.py

@@ -47,7 +47,7 @@ from .py_converter import to_python, run_as_python
 from ..transform import gradient
 
 def run_opt_pass(expr, opt_pass):
-    assert isinstance(opt_pass, transform.Pass)
+    assert isinstance(opt_pass, tvm.transform.Pass)
     mod = tvm.IRModule.from_expr(expr)
     mod = opt_pass(mod)
     entry = mod["main"]

python/tvm/relay/testing/py_converter.py

@@ -95,8 +95,8 @@ class PythonConverter(ExprFunctor):
 
         # necessary pass: SimplifyInference (otherwise we can't generate code for some operators)
         # and fusion (to get primitive functions)
-        opts = relay.transform.Sequential([relay.transform.SimplifyInference(),
-                                           relay.transform.FuseOps(fuse_opt_level=0)])
+        opts = tvm.transform.Sequential([relay.transform.SimplifyInference(),
+                                         relay.transform.FuseOps(fuse_opt_level=0)])
         mod = opts(mod)
         optimized = mod['main']
         return optimized if isinstance(unwrapped, Function) else optimized.body

python/tvm/relay/transform/transform.py

@@ -22,10 +22,9 @@ import types
 import inspect
 import functools
 
-import tvm
+import tvm.ir
 from tvm import te
 from tvm.runtime import ndarray as _nd
-from tvm.ir.transform import PassInfo, PassContext, Pass, ModulePass, Sequential, module_pass
 
 from tvm import relay
 from . import _ffi_api
@@ -78,12 +77,13 @@ def build_config(opt_level=2,
     pass_context: PassContext
         The pass context for optimizations.
     """
-    return PassContext(opt_level, fallback_device, required_pass,
-                       disabled_pass, trace)
+    return tvm.ir.transform.PassContext(
+        opt_level, fallback_device, required_pass,
+        disabled_pass, trace)
 
 
 @tvm._ffi.register_object("relay.FunctionPass")
-class FunctionPass(Pass):
+class FunctionPass(tvm.ir.transform.Pass):
     """A pass that works on each tvm.relay.Function in a module. A function
     pass class should be created through `function_pass`.
     """
@@ -94,7 +94,7 @@ def InferType():
 
     Returns
     -------
-    ret : tvm.relay.Pass
+    ret : tvm.transform.Pass
         The registered type inference pass.
     """
     return _ffi_api.InferType()
@@ -106,7 +106,7 @@ def FoldScaleAxis():
 
     Returns
     -------
-    ret : tvm.relay.Pass
+    ret : tvm.transform.Pass
         The registered pass to fold expressions.
 
     Note
@@ -123,7 +123,7 @@ def BackwardFoldScaleAxis():
 
     Returns
     -------
-    ret : tvm.relay.Pass
+    ret : tvm.transform.Pass
         The registered pass to backward fold expressions.
 
     Note
@@ -144,7 +144,7 @@ def RemoveUnusedFunctions(entry_functions=None):
 
     Returns
     -------
-    ret : tvm.relay.Pass
+    ret : tvm.transform.Pass
         The registered pass to remove unused functions.
     """
     if entry_functions is None:
@@ -156,7 +156,7 @@ def ForwardFoldScaleAxis():
 
     Returns
     -------
-    ret : tvm.relay.Pass
+    ret : tvm.transform.Pass
         The registered pass to forward fold expressions.
 
     Note
@@ -174,7 +174,7 @@ def SimplifyInference():
 
     Returns
     -------
-    ret: tvm.relay.Pass
+    ret: tvm.transform.Pass
         The registered pass to perform operator simplification.
     """
     return _ffi_api.SimplifyInference()
@@ -185,7 +185,7 @@ def FastMath():
 
     Returns
     -------
-    ret: tvm.relay.Pass
+    ret: tvm.transform.Pass
         The registered pass to perform fast math operations.
     """
     return _ffi_api.FastMath()
@@ -198,7 +198,7 @@ def CanonicalizeOps():
 
     Returns
     -------
-    ret: tvm.relay.Pass
+    ret: tvm.transform.Pass
         The registered pass performing the canonicalization.
     """
     return _ffi_api.CanonicalizeOps()
@@ -214,7 +214,7 @@ def DeadCodeElimination(inline_once=False):
 
     Returns
     -------
-    ret: tvm.relay.Pass
+    ret: tvm.transform.Pass
         The registered pass that eliminates the dead code in a Relay program.
     """
     return _ffi_api.DeadCodeElimination(inline_once)
@@ -227,7 +227,7 @@ def LazyGradientInit():
 
     Returns
     -------
-    ret: tvm.relay.Pass
+    ret: tvm.transform.Pass
         A pass which delays and/or reduces memory allocation,
         by lazily allocating 0 or one filled tensors.
     """
@@ -238,7 +238,7 @@ def FoldConstant():
 
     Returns
     -------
-    ret : tvm.relay.Pass
+    ret : tvm.transform.Pass
         The registered pass for constant folding.
     """
     return _ffi_api.FoldConstant()
@@ -255,7 +255,7 @@ def FuseOps(fuse_opt_level=-1):
 
     Returns
     -------
-    ret : tvm.relay.Pass
+    ret : tvm.transform.Pass
         The registered pass for operator fusion.
     """
     return _ffi_api.FuseOps(fuse_opt_level)
@@ -272,7 +272,7 @@ def CombineParallelConv2D(min_num_branches=3):
 
     Returns
     -------
-    ret: tvm.relay.Pass
+    ret: tvm.transform.Pass
         The registered pass that combines parallel conv2d operators.
     """
     return _ffi_api.CombineParallelConv2D(min_num_branches)
@@ -304,7 +304,7 @@ def CombineParallelDense(min_num_branches=3):
 
     Returns
     -------
-    ret: tvm.relay.Pass
+    ret: tvm.transform.Pass
         The registered pass that combines parallel dense operators.
     """
     return _ffi_api.CombineParallelDense(min_num_branches)
@@ -318,7 +318,7 @@ def AlterOpLayout():
 
     Returns
     -------
-    ret : tvm.relay.Pass
+    ret : tvm.transform.Pass
         The registered pass that alters the layout of operators.
     """
     return _ffi_api.AlterOpLayout()
@@ -366,7 +366,7 @@ def Legalize(legalize_map_attr_name="FTVMLegalize"):
 
     Returns
     -------
-    ret : tvm.relay.Pass
+    ret : tvm.transform.Pass
         The registered pass that rewrites an expr.
     """
     return _ffi_api.Legalize(legalize_map_attr_name)
@@ -387,7 +387,7 @@ def MergeComposite(pattern_table):
 
     Returns
     -------
-    ret : tvm.relay.Pass
+    ret : tvm.transform.Pass
         The registered pass that merges operators into a single composite
         relay function.
     """
@@ -413,7 +413,7 @@ def MergeCompilerRegions():
 
     Returns
     -------
-    ret : tvm.relay.Pass
+    ret : tvm.transform.Pass
         The registered pass that merges compiler regions.
     """
     return _ffi_api.MergeCompilerRegions()
@@ -433,7 +433,7 @@ def RewriteAnnotatedOps(fallback_device):
 
     Returns
     -------
-    ret: tvm.relay.Pass
+    ret: tvm.transform.Pass
         The registered pass that rewrites an expression with annotated
         `on_device` operators.
     """
@@ -448,7 +448,7 @@ def ToANormalForm():
 
     Returns
     -------
-    ret: Union[tvm.relay.Pass, tvm.relay.Expr]
+    ret: Union[tvm.transform.Pass, tvm.relay.Expr]
         The registered pass that transforms an expression into A Normal Form.
     """
     return _ffi_api.ToANormalForm()
@@ -462,7 +462,7 @@ def ToCPS(expr, mod=None):
 
     Returns
     -------
-    result: tvm.relay.Pass
+    result: tvm.transform.Pass
         The registered pass that transforms an expression into CPS.
     """
     return _ffi_api.to_cps(expr, mod)
@@ -481,7 +481,7 @@ def EtaExpand(expand_constructor=False, expand_global_var=False):
 
     Returns
     -------
-    ret: tvm.relay.Pass
+    ret: tvm.transform.Pass
         The registered pass that eta expands an expression.
     """
     return _ffi_api.EtaExpand(expand_constructor, expand_global_var)
@@ -492,7 +492,7 @@ def ToGraphNormalForm():
 
     Returns
     -------
-    ret : tvm.relay.Pass
+    ret : tvm.transform.Pass
         The registered pass that transforms an expression into Graph Normal Form.
     """
     return _ffi_api.ToGraphNormalForm()
@@ -509,7 +509,7 @@ def EliminateCommonSubexpr(fskip=None):
 
     Returns
     -------
-    ret : tvm.relay.Pass
+    ret : tvm.transform.Pass
         The registered pass that eliminates common subexpressions.
     """
     return _ffi_api.EliminateCommonSubexpr(fskip)
@@ -527,7 +527,7 @@ def PartialEvaluate():
 
     Returns
     -------
-    ret: tvm.relay.Pass
+    ret: tvm.transform.Pass
         The registered pass that performs partial evaluation on an expression.
     """
     return _ffi_api.PartialEvaluate()
@@ -539,7 +539,7 @@ def CanonicalizeCast():
 
     Returns
     -------
-    ret : tvm.relay.Pass
+    ret : tvm.transform.Pass
         The registered pass that canonicalizes cast expression.
     """
     return _ffi_api.CanonicalizeCast()
@@ -551,36 +551,19 @@ def LambdaLift():
 
     Returns
     -------
-    ret : tvm.relay.Pass
+    ret : tvm.transform.Pass
         The registered pass that lifts the lambda function.
     """
     return _ffi_api.LambdaLift()
 
 
-def PrintIR(show_meta_data=True):
-    """
-    Print the IR for a module to help debugging.
-
-    Parameters
-    ----------
-    show_meta_data : bool
-        A boolean flag to indicate if meta data should be printed.
-
-    Returns
-    -------
-    ret : tvm.relay.Pass
-        The registered pass that prints the module IR.
-    """
-    return _ffi_api.PrintIR(show_meta_data)
-
-
 def PartitionGraph():
     """Partition a Relay program into regions that can be executed on different
     backends.
 
     Returns
     -------
-    ret: tvm.relay.Pass
+    ret: tvm.transform.Pass
         The registered pass that partitions the Relay program.
     """
     return _ffi_api.PartitionGraph()
@@ -598,7 +581,7 @@ def AnnotateTarget(targets):
 
     Returns
     -------
-    ret : tvm.relay.Pass
+    ret : tvm.transform.Pass
         The annotated pass that wrapps ops with subgraph_start and
         subgraph_end.
     """
@@ -614,7 +597,7 @@ def Inline():
 
     Returns
     -------
-    ret: tvm.relay.Pass
+    ret: tvm.transform.Pass
         The registered pass that performs inlining for a Relay IR module.
     """
     return _ffi_api.Inline()
@@ -809,7 +792,7 @@ def function_pass(pass_func=None, opt_level=None, name=None, required=None):
     def create_function_pass(pass_arg):
         """Internal function that creates a function pass"""
         fname = name if name else pass_arg.__name__
-        info = PassInfo(opt_level, fname, required)
+        info = tvm.transform.PassInfo(opt_level, fname, required)
         if inspect.isclass(pass_arg):
             return _wrap_class_function_pass(pass_arg, info)
         if not isinstance(pass_arg, (types.FunctionType, types.LambdaType)):

src/ir/transform.cc

@@ -474,10 +474,10 @@ TVM_REGISTER_GLOBAL("transform.ExitPassContext")
 .set_body_typed(PassContext::Internal::ExitScope);
 
 
-Pass PrintIR(std::string header) {
-  auto pass_func =[header](IRModule mod, const PassContext& ctx) {
+Pass PrintIR(std::string header, bool show_meta_data) {
+  auto pass_func =[header, show_meta_data](IRModule mod, const PassContext& ctx) {
     LOG(INFO) << "PrintIR(" << header << "):\n"
-              << mod;
+              << AsText(mod, show_meta_data);
     return mod;
   };
   return CreateModulePass(pass_func, 0, "PrintIR", {});

src/relay/transforms/print_ir.cc

-/*
- * 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.
- */
-
-/*!
- *
- * \file src/relay/transforms/print_ir.cc
- *
- * \brief Print the module IR to help debugging.
- */
-#include <tvm/relay/expr.h>
-#include <tvm/relay/transform.h>
-
-namespace tvm {
-namespace relay {
-
-namespace transform {
-
-Pass PrintIR(bool show_meta_data) {
-  runtime::TypedPackedFunc<IRModule(IRModule, PassContext)> pass_func =
-    [=](IRModule m, PassContext pc) {
-      LOG(INFO) << "Dumping the module IR: " << std::endl << AsText(m, show_meta_data);
-      return m;
-  };
-  return CreateModulePass(pass_func, 0, "PrintIR", {});
-}
-
-TVM_REGISTER_GLOBAL("relay._transform.PrintIR")
-.set_body_typed(PrintIR);
-
-}  // namespace transform
-
-}  // namespace relay
-}  // namespace tvm

tests/python/relay/test_op_level10.py

@@ -53,10 +53,10 @@ def test_checkpoint_alpha_equal():
     df = transform.gradient(run_infer_type(f))
 
     # run PE and DCE
-    with transform.PassContext(opt_level=3):
+    with tvm.transform.PassContext(opt_level=3):
         passes = [transform.PartialEvaluate(),
                   transform.DeadCodeElimination(inline_once=True)]
-        mod = transform.Sequential(passes)(tvm.IRModule.from_expr(df))
+        mod = tvm.transform.Sequential(passes)(tvm.IRModule.from_expr(df))
         df = mod["main"]
 
     df_parsed = relay.parser.fromtext(
@@ -109,10 +109,10 @@ def test_checkpoint_alpha_equal_tuple():
     df = transform.gradient(run_infer_type(f))
 
     # run PE and DCE
-    with transform.PassContext(opt_level=3):
+    with tvm.transform.PassContext(opt_level=3):
         passes = [transform.PartialEvaluate(),
                   transform.DeadCodeElimination(inline_once=True)]
-        mod = transform.Sequential(passes)(tvm.IRModule.from_expr(df))
+        mod = tvm.transform.Sequential(passes)(tvm.IRModule.from_expr(df))
         df = mod["main"]
 
     df_parsed = relay.parser.fromtext(

tests/python/relay/test_pass_alter_op_layout.py

@@ -26,8 +26,8 @@ from tvm.relay.testing.temp_op_attr import TempOpAttr
 def run_opt_pass(expr, passes):
     passes = passes if isinstance(passes, list) else [passes]
     mod = tvm.IRModule.from_expr(expr)
-    seq = transform.Sequential(passes)
-    with transform.PassContext(opt_level=3):
+    seq = tvm.transform.Sequential(passes)
+    with tvm.transform.PassContext(opt_level=3):
         mod = seq(mod)
     entry = mod["main"]
     return entry if isinstance(expr, relay.Function) else entry.body

tests/python/relay/test_pass_annotation.py

@@ -28,8 +28,8 @@ from tvm.relay import transform
 def run_opt_pass(expr, passes):
     passes = passes if isinstance(passes, list) else [passes]
     mod = tvm.IRModule.from_expr(expr)
-    seq = transform.Sequential(passes)
-    with transform.PassContext(opt_level=3):
+    seq = tvm.transform.Sequential(passes)
+    with tvm.transform.PassContext(opt_level=3):
         mod = seq(mod)
     return mod["main"]
 

tests/python/relay/test_pass_canonicalize_cast.py

@@ -54,9 +54,9 @@ def test_canonicalize_cast():
         bias2 = relay.var("bias2", shape=(16, 1, 1), dtype="int32")
         y = before(data, conv_weight, bias1, bias2)
         mod = tvm.IRModule.from_expr(y)
-        seq = _transform.Sequential([_transform.InferType(), _transform.CanonicalizeCast(),
+        seq = tvm.transform.Sequential([_transform.InferType(), _transform.CanonicalizeCast(),
                                      _transform.InferType()])
-        with _transform.PassContext(opt_level=3):
+        with tvm.transform.PassContext(opt_level=3):
             mod = seq(mod)
         y = mod["main"]
         y_expected = expected(data, conv_weight, bias1, bias2)

tests/python/relay/test_pass_combine_parallel_conv2d.py

@@ -26,7 +26,7 @@ def run_combine_parallel(expr, min_num_branches=3):
     return mod["main"]
 
 def run_opt_pass(expr, opt_pass):
-    assert isinstance(opt_pass, transform.Pass)
+    assert isinstance(opt_pass, tvm.transform.Pass)
     mod = tvm.IRModule.from_expr(expr)
     mod = opt_pass(mod)
     return mod["main"]

tests/python/relay/test_pass_combine_parallel_dense.py

@@ -26,7 +26,7 @@ def run_combine_parallel(expr, min_num_branches=3):
     return mod["main"]
 
 def run_opt_pass(expr, opt_pass):
-    assert isinstance(opt_pass, transform.Pass)
+    assert isinstance(opt_pass, tvm.transform.Pass)
     mod = tvm.IRModule.from_expr(expr)
     mod = opt_pass(mod)
     return mod["main"]

tests/python/relay/test_pass_convert_op_layout.py

@@ -26,8 +26,8 @@ from tvm.relay import transform, analysis
 def run_opt_pass(expr, passes):
     passes = passes if isinstance(passes, list) else [passes]
     mod = tvm.IRModule.from_expr(expr)
-    seq = transform.Sequential(passes)
-    with transform.PassContext(opt_level=3):
+    seq = tvm.transform.Sequential(passes)
+    with tvm.transform.PassContext(opt_level=3):
         mod = seq(mod)
     entry = mod["main"]
     return entry if isinstance(expr, relay.Function) else entry.body

tests/python/relay/test_pass_dead_code_elimination.py

@@ -47,7 +47,7 @@ e = env()
 
 
 def run_opt_pass(expr, opt_pass):
-    assert isinstance(opt_pass, transform.Pass)
+    assert isinstance(opt_pass, tvm.transform.Pass)
     mod = tvm.IRModule.from_expr(expr)
     mod = opt_pass(mod)
     entry = mod["main"]

tests/python/relay/test_pass_eliminate_common_subexpr.py

@@ -24,7 +24,7 @@ from tvm.relay import transform, analysis
 
 
 def run_opt_pass(expr, opt_pass):
-    assert isinstance(opt_pass, transform.Pass)
+    assert isinstance(opt_pass, tvm.transform.Pass)
     mod = tvm.IRModule.from_expr(expr)
     mod = opt_pass(mod)
     entry = mod["main"]

tests/python/relay/test_pass_eta_expand.py

@@ -33,8 +33,8 @@ def test_eta_expand_global_var():
             @aux
         }
     """)
-    seq = _transform.Sequential([_transform.EtaExpand(expand_global_var=True)])
-    with _transform.PassContext(opt_level=3):
+    seq = tvm.transform.Sequential([_transform.EtaExpand(expand_global_var=True)])
+    with tvm.transform.PassContext(opt_level=3):
         mod = seq(mod)
     expected = relay.fromtext(r"""
         v0.0.4
@@ -62,8 +62,8 @@ def test_eta_expand_constructor():
             Cons
         }
     """)
-    seq = _transform.Sequential([_transform.EtaExpand(expand_constructor=True)])
-    with _transform.PassContext(opt_level=3):
+    seq = tvm.transform.Sequential([_transform.EtaExpand(expand_constructor=True)])
+    with tvm.transform.PassContext(opt_level=3):
         mod = seq(mod)
     expected = relay.fromtext(r"""
         v0.0.4

tests/python/relay/test_pass_fold_constant.py

@@ -24,7 +24,7 @@ from tvm.relay.testing import run_infer_type, create_workload
 
 
 def run_opt_pass(expr, opt_pass):
-    assert isinstance(opt_pass, transform.Pass)
+    assert isinstance(opt_pass, tvm.transform.Pass)
 
     mod = tvm.IRModule.from_expr(expr)
     mod = opt_pass(mod)
@@ -174,7 +174,7 @@ def test_fold_batch_norm():
         add = relay.add(conv, bias)
         return relay.Function(relay.analysis.free_vars(add), add)
 
-    remove_bn_pass = transform.Sequential([
+    remove_bn_pass = tvm.transform.Sequential([
         relay.transform.InferType(),
         relay.transform.SimplifyInference(),
         relay.transform.FoldConstant(),

tests/python/relay/test_pass_fold_scale_axis.py

@@ -26,7 +26,7 @@ def _get_positive_scale(size):
 
 
 def run_opt_pass(expr, opt_pass):
-    assert isinstance(opt_pass, transform.Pass)
+    assert isinstance(opt_pass, tvm.transform.Pass)
     mod = tvm.IRModule.from_expr(expr)
     mod = opt_pass(mod)
     entry = mod["main"]

tests/python/relay/test_pass_lazy_gradient_init.py

@@ -80,7 +80,7 @@ def test_add_tuple():
 
   mod["main"] = y
   mod = transform.LazyGradientInit()(mod)
-  mod = transform.PrintIR(show_meta_data=True)(mod)
+  mod = tvm.transform.PrintIR(show_meta_data=True)(mod)
   y = mod["main"]
 
   assert mod["main"].checked_type == relay.FuncType([t], tensor_type)
@@ -116,7 +116,7 @@ def test_mult():
 def test_ret_tuple():
   """Test tuple return type. Check types and semantic equivalence."""
   mod = tvm.IRModule()
-  
+
   shape = (10, 10)
   dtype = 'float32'
   t = relay.TensorType(shape, dtype)
@@ -141,7 +141,7 @@ def test_ret_tuple():
 def test_add_broadcast():
   """Test adding matrices of different size. Check types and semantic equivalence."""
   mod = tvm.IRModule()
-  
+
   shape1 = (3, 4, 1)
   shape2 = (1, 5)
   dtype = 'float32'
@@ -173,7 +173,7 @@ def test_reverse_ad_identity():
   """Simple test with reverse mode ad."""
   # of f(x) = x
   mod = tvm.IRModule()
-  
+
   shape = (10, 10)
   dtype = 'float32'
   t = relay.TensorType(shape, dtype)
@@ -201,7 +201,7 @@ def test_reverse_ad_identity():
 def test_multivar_reverse_ad():
   """Simple test with multivariate reverse mode ad."""
   mod = tvm.IRModule()
-  
+
   shape = (10, 10)
   dtype = 'float32'
   t = relay.TensorType(shape, dtype)
@@ -232,7 +232,7 @@ def test_multivar_reverse_ad():
 def test_after_partial_eval():
   """Test transformation following reverse mode ad and PartialEval"""
   mod = tvm.IRModule()
-  
+
   shape = (10, 10)
   dtype = 'float32'
   t = relay.TensorType(shape, dtype)
@@ -248,7 +248,7 @@ def test_after_partial_eval():
   mod["main"] = back_func
   back_func = mod["main"]
 
-  seq = transform.Sequential([
+  seq = tvm.transform.Sequential([
     transform.PartialEvaluate(),
     transform.LazyGradientInit(),
     transform.DeadCodeElimination()
@@ -270,7 +270,7 @@ def test_after_partial_eval():
 def test_before_partial_eval():
   """Test transformation before PartialEval"""
   mod = tvm.IRModule()
-  
+
   shape = (10, 10)
   dtype = 'float32'
   t = relay.TensorType(shape, dtype)
@@ -284,7 +284,7 @@ def test_before_partial_eval():
   back_func = run_infer_type(back_func)
 
   mod["main"] = back_func
-  seq = transform.Sequential([
+  seq = tvm.transform.Sequential([
     transform.LazyGradientInit(),
     transform.PartialEvaluate(),
     transform.DeadCodeElimination()
@@ -306,7 +306,7 @@ def test_before_partial_eval():
 def test_zeros():
   """Simple test using "zeros" op"""
   mod = tvm.IRModule()
-  
+
   shape = (10, 10)
   dtype = 'float32'
   t = relay.TensorType(shape, dtype)
@@ -328,7 +328,7 @@ def test_zeros():
 def test_ones():
   """Simple test using "ones" op"""
   mod = tvm.IRModule()
-  
+
   shape = (10, 10)
   dtype = 'float32'
   t = relay.TensorType(shape, dtype)
@@ -350,7 +350,7 @@ def test_ones():
 def test_zeros_like():
   """Simple test using "zeros_like" op"""
   mod = tvm.IRModule()
-  
+
   shape = (10, 10)
   dtype = 'float32'
   t = relay.TensorType(shape, dtype)
@@ -372,7 +372,7 @@ def test_zeros_like():
 def test_ones_like():
   """Simple test using "ones_like" op"""
   mod = tvm.IRModule()
-  
+
   shape = (10, 10)
   dtype = 'float32'
   t = relay.TensorType(shape, dtype)

tests/python/relay/test_pass_legalize.py

@@ -28,8 +28,8 @@ from tvm.relay.testing.temp_op_attr import TempOpAttr
 def run_opt_pass(expr, passes):
     passes = passes if isinstance(passes, list) else [passes]
     mod = tvm.IRModule.from_expr(expr)
-    seq = transform.Sequential(passes)
-    with transform.PassContext(opt_level=3):
+    seq = tvm.transform.Sequential(passes)
+    with tvm.transform.PassContext(opt_level=3):
         mod = seq(mod)
     entry = mod["main"]
     return entry if isinstance(expr, relay.Function) else entry.body

tests/python/relay/test_pass_mac_count.py

@@ -23,7 +23,7 @@ from tvm.relay import analysis, transform
 
 
 def run_opt_pass(expr, opt_pass):
-    assert isinstance(opt_pass, transform.Pass)
+    assert isinstance(opt_pass, tvm.transform.Pass)
     mod = tvm.IRModule.from_expr(expr)
     mod = opt_pass(mod)
     entry = mod["main"]

tests/python/relay/test_pass_manager.py

@@ -129,13 +129,13 @@ def test_module_pass():
     opt_tester = OptTester(mod)
     pass_ctx = None
 
-    @_transform.module_pass(opt_level=opt_level, name=pass_name)
+    @tvm.transform.module_pass(opt_level=opt_level, name=pass_name)
     def transform(expr, ctx):
         return opt_tester.transform(expr, ctx)
 
     def test_pass_registration():
         mod_pass = transform
-        assert isinstance(mod_pass, _transform.ModulePass)
+        assert isinstance(mod_pass, tvm.transform.ModulePass)
         pass_info = mod_pass.info
         assert pass_info.name == pass_name
         assert pass_info.opt_level == opt_level
@@ -143,8 +143,8 @@ def test_module_pass():
     def test_pass_registration_no_decorator():
         def direct_transform(expr, ctx):
             return opt_tester.transform(expr, ctx)
-        mod_pass = _transform.module_pass(direct_transform, opt_level=3)
-        assert isinstance(mod_pass, _transform.ModulePass)
+        mod_pass = tvm.transform.module_pass(direct_transform, opt_level=3)
+        assert isinstance(mod_pass, tvm.transform.ModulePass)
         pass_info = mod_pass.info
         assert pass_info.name == "direct_transform"
         assert pass_info.opt_level == 3
@@ -285,7 +285,7 @@ def test_function_pass():
 
 
 def test_module_class_pass():
-    @relay.transform.module_pass(opt_level=1)
+    @tvm.transform.module_pass(opt_level=1)
     class TestPipeline:
         """Simple test function to replace one argument to another."""
         def __init__(self, new_mod, replace):
@@ -309,7 +309,7 @@ def test_module_class_pass():
 
 
 def test_pass_info():
-    info = relay.transform.PassInfo(opt_level=1, name="xyz")
+    info = tvm.transform.PassInfo(opt_level=1, name="xyz")
     assert info.opt_level == 1
     assert info.name == "xyz"
 
@@ -350,7 +350,7 @@ def test_sequential_pass():
     opt_tester = OptTester(mod)
     pass_ctx = None
 
-    @_transform.module_pass(opt_level=1)
+    @tvm.transform.module_pass(opt_level=1)
     def mod_transform(expr, ctx):
         return opt_tester.transform(expr, ctx)
 
@@ -367,21 +367,21 @@ def test_sequential_pass():
         passes = [module_pass, function_pass]
         opt_level = 2
         pass_name = "sequential"
-        sequential = _transform.Sequential(passes=passes, opt_level=opt_level)
+        sequential = tvm.transform.Sequential(passes=passes, opt_level=opt_level)
         pass_info = sequential.info
         assert pass_info.name == pass_name
         assert pass_info.opt_level == opt_level
 
     def test_no_pass():
         passes = []
-        sequential = _transform.Sequential(opt_level=1, passes=passes)
+        sequential = tvm.transform.Sequential(opt_level=1, passes=passes)
         ret_mod = sequential(mod)
         mod_func = ret_mod[v_sub]
         check_func(sub, mod_func)
 
     def test_only_module_pass():
         passes = [module_pass]
-        sequential = _transform.Sequential(opt_level=1, passes=passes)
+        sequential = tvm.transform.Sequential(opt_level=1, passes=passes)
         with relay.build_config(required_pass=["mod_transform"]):
             ret_mod = sequential(mod)
         # Check the subtract function.
@@ -396,7 +396,7 @@ def test_sequential_pass():
     def test_only_function_pass():
         # Check the subtract function.
         passes = [function_pass]
-        sequential = _transform.Sequential(opt_level=1, passes=passes)
+        sequential = tvm.transform.Sequential(opt_level=1, passes=passes)
         with relay.build_config(required_pass=["func_transform"]):
             ret_mod = sequential(mod)
         _, new_sub = extract_var_func(ret_mod, v_sub.name_hint)
@@ -411,7 +411,7 @@ def test_sequential_pass():
         # function pass.
         mod = tvm.IRModule({v_sub: sub, v_log: log})
         passes = [module_pass, function_pass]
-        sequential = _transform.Sequential(opt_level=1, passes=passes)
+        sequential = tvm.transform.Sequential(opt_level=1, passes=passes)
         required = ["mod_transform", "func_transform"]
         with relay.build_config(required_pass=required):
             ret_mod = sequential(mod)
@@ -482,7 +482,7 @@ def test_sequential_with_scoping():
         z1 = relay.add(z, z)
         return relay.Function([x], z1)
 
-    seq = _transform.Sequential([
+    seq = tvm.transform.Sequential([
         relay.transform.InferType(),
         relay.transform.FoldConstant(),
         relay.transform.EliminateCommonSubexpr(),
@@ -507,10 +507,10 @@ def test_print_ir(capfd):
     y = relay.multiply(y, relay.const(2, "float32"))
     func = relay.Function([x], y)
 
-    seq = _transform.Sequential([
+    seq = tvm.transform.Sequential([
         relay.transform.InferType(),
         relay.transform.FoldConstant(),
-        relay.transform.PrintIR(),
+        tvm.transform.PrintIR(),
         relay.transform.DeadCodeElimination()
     ])
 
@@ -520,7 +520,7 @@ def test_print_ir(capfd):
 
     out = capfd.readouterr().err
 
-    assert "Dumping the module IR" in out
+    assert "PrintIR" in out
     assert "multiply" in out
 
 __TRACE_COUNTER__ = 0
@@ -539,7 +539,7 @@ def test_print_debug_callback():
     y = relay.multiply(y, relay.const(2, "float32"))
     func = relay.Function([x], y)
 
-    seq = _transform.Sequential([
+    seq = tvm.transform.Sequential([
         relay.transform.InferType(),
         relay.transform.FoldConstant(),
         relay.transform.DeadCodeElimination()

tests/python/relay/test_pass_partial_eval.py

@@ -38,8 +38,8 @@ def check_eval(expr, expected_result, mod=None, rtol=1e-07):
 def run_opt_pass(expr, passes):
     passes = passes if isinstance(passes, list) else [passes]
     mod = tvm.IRModule.from_expr(expr)
-    seq = transform.Sequential(passes)
-    with transform.PassContext(opt_level=3):
+    seq = tvm.transform.Sequential(passes)
+    with tvm.transform.PassContext(opt_level=3):
        mod = seq(mod)
     entry = mod["main"]
     return entry if isinstance(expr, relay.Function) else entry.body
@@ -58,7 +58,7 @@ def dcpe(expr, mod=None, grad=False):
     if mod:
         assert isinstance(expr, Function)
         mod["main"] = expr
-        seq = transform.Sequential(passes)
+        seq = tvm.transform.Sequential(passes)
         mod = seq(mod)
         return mod["main"]
     return run_opt_pass(expr, passes)

tests/python/relay/test_pass_partition_graph.py

@@ -496,7 +496,7 @@ def test_function_lifting():
         op_list = ["nn.batch_norm", "nn.conv2d"]
         mod = WhiteListAnnotator(op_list, "test_compiler")(mod)
 
-        opt_pass = transform.Sequential([
+        opt_pass = tvm.transform.Sequential([
             transform.InferType(),
             transform.PartitionGraph(),
             transform.SimplifyInference(),
@@ -578,7 +578,7 @@ def test_function_lifting_inline():
         op_list = ["nn.batch_norm", "nn.conv2d"]
         mod = WhiteListAnnotator(op_list, "test_compiler")(mod)
 
-        opt_pass = transform.Sequential([
+        opt_pass = tvm.transform.Sequential([
             transform.InferType(),
             transform.PartitionGraph(),
             transform.SimplifyInference(),
@@ -878,13 +878,13 @@ def test_dnnl_fuse():
         # This is required for constant folding
         mod["main"] = bind_params_by_name(mod["main"], params)
 
-        remove_bn_pass = transform.Sequential([
+        remove_bn_pass = tvm.transform.Sequential([
             transform.InferType(),
             transform.SimplifyInference(),
             transform.FoldConstant(),
             transform.FoldScaleAxis(),
         ])
-        composite_partition = transform.Sequential([
+        composite_partition = tvm.transform.Sequential([
             remove_bn_pass,
             transform.MergeComposite(pattern_table),
             transform.AnnotateTarget("dnnl"),

tests/python/relay/test_pass_qnn_legalize.py

@@ -37,8 +37,8 @@ def alpha_equal(x, y):
 def run_opt_pass(expr, passes):
     passes = passes if isinstance(passes, list) else [passes]
     mod = tvm.IRModule.from_expr(expr)
-    seq = transform.Sequential(passes)
-    with transform.PassContext(opt_level=3):
+    seq = tvm.transform.Sequential(passes)
+    with tvm.transform.PassContext(opt_level=3):
         mod = seq(mod)
     entry = mod["main"]
     return entry if isinstance(expr, relay.Function) else entry.body

tests/python/relay/test_pass_to_a_normal_form.py

@@ -28,8 +28,8 @@ from tvm.relay.analysis import Feature
 def run_opt_pass(expr, passes):
     passes = passes if isinstance(passes, list) else [passes]
     mod = tvm.IRModule.from_expr(expr)
-    seq = transform.Sequential(passes)
-    with transform.PassContext(opt_level=3):
+    seq = tvm.transform.Sequential(passes)
+    with tvm.transform.PassContext(opt_level=3):
        mod = seq(mod)
     entry = mod["main"]
     return entry if isinstance(expr, relay.Function) else entry.body

tests/python/relay/test_pass_to_cps.py

@@ -71,7 +71,8 @@ def test_cps_pe():
         x = run_infer_type(x)
         y = un_cps(x)
         y = run_infer_type(y)
-        x = run_opt_pass(x, transform.Sequential([transform.PartialEvaluate(), transform.DeadCodeElimination(inline_once=True)]))
+        x = run_opt_pass(x, tvm.transform.Sequential(
+            [transform.PartialEvaluate(), transform.DeadCodeElimination(inline_once=True)]))
         assert Feature.fRefCreate not in detect_feature(x)
     unit = relay.Function([], relay.const(0., dtype='float32'))
     f_ref = relay.Var("f_ref")

tutorials/dev/relay_pass_infra.py

@@ -29,7 +29,7 @@ introduced an infrastructure to manage the optimization passes.
 The optimizations of a Relay program could be applied at various granularity,
 namely function-level and module-level using :py:class:`tvm.relay.transform.FunctionPass`
 and py:class:`tvm.relay.transform.ModulePass`
-respectively. Or users can rely on py:class:`tvm.relay.transform.Sequential` to apply a sequence of passes
+respectively. Or users can rely on py:class:`tvm.transform.Sequential` to apply a sequence of passes
 on a Relay program where the dependencies between passes can be resolved by the
 pass infra. For more details about each type of these passes, please refer to
 the :ref:`relay-pass-infra`
@@ -130,22 +130,22 @@ print(mod)
 # fusion, as this pass generates let bindings for each expression to
 # canonicalize a Relay program.
 #
-# Relay, hence, provides :py:class:`tvm.relay.transform.Sequential` to alleviate developers from handling
+# Relay, hence, provides :py:class:`tvm.transform.Sequential` to alleviate developers from handling
 # these issues explicitly by specifying the required passes of each pass and
 # packing them as a whole to execute. For example, the same passes can now be
-# applied using the sequential style as the following. :py:class:`tvm.relay.transform.Sequential` is
+# applied using the sequential style as the following. :py:class:`tvm.transform.Sequential` is
 # similiar to `torch.nn.sequential <https://pytorch.org/docs/stable/nn.html#torch.nn.Sequential>`_
 # and `mxnet.gluon.block <https://mxnet.incubator.apache.org/api/python/docs/_modules/mxnet/gluon/block.html>`_.
 # For example, `torch.nn.sequential` is used to contain a sequence of PyTorch
 # `Modules` that will be added to build a network. It focuses on the network
-# layers. Instead, the :py:class:`tvm.relay.transform.Sequential` in our pass infra works on the optimizing
+# layers. Instead, the :py:class:`tvm.transform.Sequential` in our pass infra works on the optimizing
 # pass.
 
-# Now let's execute some passes through :py:class:`tvm.relay.transform.Sequential`
+# Now let's execute some passes through :py:class:`tvm.transform.Sequential`
 f = example()
 mod = tvm.IRModule.from_expr(f)
 # Glob the interested passes.
-seq = relay.transform.Sequential([relay.transform.FoldConstant(),
+seq = tvm.transform.Sequential([relay.transform.FoldConstant(),
                                   relay.transform.EliminateCommonSubexpr(),
                                   relay.transform.FuseOps(fuse_opt_level=2)])
 mod1 = seq(mod)
@@ -156,7 +156,7 @@ print(mod1)
 # identical addition operations. This is because `EliminateCommonSubexpr`
 # was not actually performed. The reason is because only the passes that have
 # optimization level less or equal to 2 will be executed by default under
-# :py:class:`tvm.relay.transform.Sequential`. The pass infra,
+# :py:class:`tvm.transform.Sequential`. The pass infra,
 # however, provides a configuration interface
 # for users to customize the optimization level that they want to execute.
 
@@ -186,7 +186,7 @@ with relay.build_config(opt_level=3):
     mod4 = seq(mod)
 print(mod4)
 
-seq1 = relay.transform.Sequential([relay.transform.AlterOpLayout()])
+seq1 = tvm.transform.Sequential([relay.transform.AlterOpLayout()])
 with relay.build_config(opt_level=3):
     with tvm.target.create("llvm"):
         mod5 = seq1(mod)
@@ -237,11 +237,11 @@ print(mod3)
 
 f = example()
 mod = tvm.IRModule.from_expr(f)
-seq = relay.transform.Sequential([relay.transform.FoldConstant(),
-                                  relay.transform.PrintIR(False),
-                                  relay.transform.EliminateCommonSubexpr(),
-                                  relay.transform.FuseOps(),
-                                  relay.transform.PrintIR(False)])
+seq = tvm.transform.Sequential([relay.transform.FoldConstant(),
+                                tvm.transform.PrintIR(),
+                                relay.transform.EliminateCommonSubexpr(),
+                                relay.transform.FuseOps(),
+                                tvm.transform.PrintIR()])
 with relay.build_config(opt_level=3):
     mod = seq(mod)
 

vta/python/vta/top/graphpack.py

@@ -24,7 +24,7 @@ from tvm.relay import ExprMutator
 
 def run_opt_pass(expr, opt_pass):
     """Exectue a relay pass."""
-    assert isinstance(opt_pass, transform.Pass)
+    assert isinstance(opt_pass, tvm.transform.Pass)
     mod = tvm.IRModule.from_expr(expr)
     mod = opt_pass(mod)
     entry = mod["main"]