.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "how_to/extend_tvm/use_pass_infra.py" .. LINE NUMBERS ARE GIVEN BELOW. .. only:: html .. note:: :class: sphx-glr-download-link-note Click :ref:`here ` to download the full example code .. rst-class:: sphx-glr-example-title .. _sphx_glr_how_to_extend_tvm_use_pass_infra.py: .. _tutorial-use-pass-infra: How to Use TVM Pass Infra ========================= **Author**: `Zhi Chen `_ As the number of optimization passes increases in Relay/tir, it becomes intractable to execute them and maintain their dependencies manually. Therefore, we have introduced an infrastructure to manage the optimization passes and make it applicable to different layers of the IR in the TVM stack. The optimizations of a Relay/tir program could be applied at various granularity, namely function-level and module-level using :py:class:`tvm.relay.transform.FunctionPass`/ :py:class:`tvm.tir.transform.PrimFuncPass` and :py:class:`tvm.transform.ModulePass` respectively. Or users can rely on :py:class:`tvm.transform.Sequential` to apply a sequence of passes on a Relay/tir 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:`pass-infra` This tutorial mainly demonstrates how developers can use the pass infra to perform a certain optimization and create an optimization pipeline for a Relay program. The same approach can be used for tir as well. .. GENERATED FROM PYTHON SOURCE LINES 42-49 .. code-block:: default import numpy as np import tvm from tvm import te import tvm.relay as relay .. GENERATED FROM PYTHON SOURCE LINES 55-60 Create An Example Relay Program ------------------------------- First of all, we create a simple Relay program for the tutorial. This program will be used by various optimizations of the examples in this tutorial. Similarly, users can write a tir primitive function and apply the tir passes. .. GENERATED FROM PYTHON SOURCE LINES 60-78 .. code-block:: default def example(): shape = (1, 64, 54, 54) c_data = np.empty(shape).astype("float32") c = relay.const(c_data) weight = relay.var("weight", shape=(64, 64, 3, 3)) x = relay.var("x", relay.TensorType((1, 64, 56, 56), "float32")) conv = relay.nn.conv2d(x, weight) y = relay.add(c, c) y = relay.multiply(y, relay.const(2, "float32")) y = relay.add(conv, y) z = relay.add(y, c) z1 = relay.add(y, c) z2 = relay.add(z, z1) return relay.Function([x, weight], z2) .. GENERATED FROM PYTHON SOURCE LINES 79-89 Optimize the Program -------------------- Now we would like to optimize the program. Relay features a host of optimizations. We will select some of them to apply on this example program. There are multiple ways to optimize a Relay program. Below we will provide examples for each of them. Manually Apply Optimization Passes ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ .. GENERATED FROM PYTHON SOURCE LINES 89-106 .. code-block:: default # Let's first create a relay Module which contains one or multiple Relay # functions for optimization. f = example() mod = tvm.IRModule.from_expr(f) # Now we can apply constant folding on the module. # fold_const here is a callback that doesn't take any parameters. fold_const = relay.transform.FoldConstant() # Then, we can invoke the pass on the given module. Note that the constant # folding pass works at the function-level. That being said, each function in # the module will be applied with the optimization. Users don't need to iterate # through individual functions manually to apply this pass. mod = fold_const(mod) # We can see from the updated program that the constants are folded. print(mod) .. rst-class:: sphx-glr-script-out .. code-block:: none def @main(%x: Tensor[(1, 64, 56, 56), float32] /* ty=Tensor[(1, 64, 56, 56), float32] */, %weight: Tensor[(64, 64, 3, 3), float32] /* ty=Tensor[(64, 64, 3, 3), float32] */) -> Tensor[(1, 64, 54, 54), float32] { %0 = nn.conv2d(%x, %weight, padding=[0, 0, 0, 0]) /* ty=Tensor[(1, 64, 54, 54), float32] */; %1 = add(%0, meta[relay.Constant][0] /* ty=Tensor[(1, 64, 54, 54), float32] */) /* ty=Tensor[(1, 64, 54, 54), float32] */; %2 = add(%1, meta[relay.Constant][1] /* ty=Tensor[(1, 64, 54, 54), float32] */) /* ty=Tensor[(1, 64, 54, 54), float32] */; %3 = add(%1, meta[relay.Constant][1] /* ty=Tensor[(1, 64, 54, 54), float32] */) /* ty=Tensor[(1, 64, 54, 54), float32] */; add(%2, %3) /* ty=Tensor[(1, 64, 54, 54), float32] */ } .. GENERATED FROM PYTHON SOURCE LINES 107-109 More optimizations can be applied in the similar manner. For instance, we can eliminate the common expressions that used by `z` and `z1`. .. GENERATED FROM PYTHON SOURCE LINES 109-112 .. code-block:: default mod = relay.transform.EliminateCommonSubexpr()(mod) print(mod) .. rst-class:: sphx-glr-script-out .. code-block:: none def @main(%x: Tensor[(1, 64, 56, 56), float32] /* ty=Tensor[(1, 64, 56, 56), float32] */, %weight: Tensor[(64, 64, 3, 3), float32] /* ty=Tensor[(64, 64, 3, 3), float32] */) -> Tensor[(1, 64, 54, 54), float32] { %0 = nn.conv2d(%x, %weight, padding=[0, 0, 0, 0]) /* ty=Tensor[(1, 64, 54, 54), float32] */; %1 = add(%0, meta[relay.Constant][0] /* ty=Tensor[(1, 64, 54, 54), float32] */) /* ty=Tensor[(1, 64, 54, 54), float32] */; %2 = add(%1, meta[relay.Constant][1] /* ty=Tensor[(1, 64, 54, 54), float32] */) /* ty=Tensor[(1, 64, 54, 54), float32] */; add(%2, %2) /* ty=Tensor[(1, 64, 54, 54), float32] */ } .. GENERATED FROM PYTHON SOURCE LINES 113-116 Some optimizations, such as fusion, are parametric as well. For example, opt level 0 will not allow operators to be fused together. Users can pass the `fuse_opt_level` to enable this. .. GENERATED FROM PYTHON SOURCE LINES 116-122 .. code-block:: default mod = relay.transform.FuseOps(fuse_opt_level=0)(mod) # We can observe that the optimized module contains functions that only have # a signle primitive op. print(mod) .. rst-class:: sphx-glr-script-out .. code-block:: none def @main(%x: Tensor[(1, 64, 56, 56), float32] /* ty=Tensor[(1, 64, 56, 56), float32] */, %weight: Tensor[(64, 64, 3, 3), float32] /* ty=Tensor[(64, 64, 3, 3), float32] */) -> Tensor[(1, 64, 54, 54), float32] { %0 = fn (%p03: Tensor[(1, 64, 56, 56), float32] /* ty=Tensor[(1, 64, 56, 56), float32] */, %p12: Tensor[(64, 64, 3, 3), float32] /* ty=Tensor[(64, 64, 3, 3), float32] */, Primitive=1) -> Tensor[(1, 64, 54, 54), float32] { nn.conv2d(%p03, %p12, padding=[0, 0, 0, 0]) /* ty=Tensor[(1, 64, 54, 54), float32] */ } /* ty=fn (Tensor[(1, 64, 56, 56), float32], Tensor[(64, 64, 3, 3), float32]) -> Tensor[(1, 64, 54, 54), float32] */; %1 = %0(%x, %weight) /* ty=Tensor[(1, 64, 54, 54), float32] */; %2 = fn (%p02: Tensor[(1, 64, 54, 54), float32] /* ty=Tensor[(1, 64, 54, 54), float32] */, %p11: Tensor[(1, 64, 54, 54), float32] /* ty=Tensor[(1, 64, 54, 54), float32] */, Primitive=1) -> Tensor[(1, 64, 54, 54), float32] { add(%p02, %p11) /* ty=Tensor[(1, 64, 54, 54), float32] */ } /* ty=fn (Tensor[(1, 64, 54, 54), float32], Tensor[(1, 64, 54, 54), float32]) -> Tensor[(1, 64, 54, 54), float32] */; %3 = %2(%1, meta[relay.Constant][0] /* ty=Tensor[(1, 64, 54, 54), float32] */) /* ty=Tensor[(1, 64, 54, 54), float32] */; %4 = fn (%p01: Tensor[(1, 64, 54, 54), float32] /* ty=Tensor[(1, 64, 54, 54), float32] */, %p1: Tensor[(1, 64, 54, 54), float32] /* ty=Tensor[(1, 64, 54, 54), float32] */, Primitive=1) -> Tensor[(1, 64, 54, 54), float32] { add(%p01, %p1) /* ty=Tensor[(1, 64, 54, 54), float32] */ } /* ty=fn (Tensor[(1, 64, 54, 54), float32], Tensor[(1, 64, 54, 54), float32]) -> Tensor[(1, 64, 54, 54), float32] */; %5 = %4(%3, meta[relay.Constant][1] /* ty=Tensor[(1, 64, 54, 54), float32] */) /* ty=Tensor[(1, 64, 54, 54), float32] */; %6 = fn (%p0: Tensor[(1, 64, 54, 54), float32] /* ty=Tensor[(1, 64, 54, 54), float32] */, Primitive=1) -> Tensor[(1, 64, 54, 54), float32] { add(%p0, %p0) /* ty=Tensor[(1, 64, 54, 54), float32] */ } /* ty=fn (Tensor[(1, 64, 54, 54), float32]) -> Tensor[(1, 64, 54, 54), float32] */; %6(%5) /* ty=Tensor[(1, 64, 54, 54), float32] */ } .. GENERATED FROM PYTHON SOURCE LINES 123-142 Use Sequential to Apply a Sequence of Passes ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Applying passes as above is actually tedious and it may require users to have better understanding about the dependencies between them. For example, fusion currently doesn't work well on let bindings. Therefore, we would not be able to fuse operators that were fusable if :py:func:`relay.transform.ToANormalForm` is applied before fusion, as this pass generates let bindings for each expression to canonicalize a Relay program. 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.transform.Sequential` is similar to `torch.nn.sequential `_ and `mxnet.gluon.block `_. 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.transform.Sequential` in our pass infra works on the optimizing pass. .. GENERATED FROM PYTHON SOURCE LINES 142-157 .. code-block:: default # 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 = tvm.transform.Sequential( [ relay.transform.FoldConstant(), relay.transform.EliminateCommonSubexpr(), relay.transform.FuseOps(fuse_opt_level=2), ] ) mod1 = seq(mod) print(mod1) .. rst-class:: sphx-glr-script-out .. code-block:: none def @main(%x: Tensor[(1, 64, 56, 56), float32] /* ty=Tensor[(1, 64, 56, 56), float32] */, %weight: Tensor[(64, 64, 3, 3), float32] /* ty=Tensor[(64, 64, 3, 3), float32] */) -> Tensor[(1, 64, 54, 54), float32] { %4 = fn (%p0: Tensor[(1, 64, 56, 56), float32] /* ty=Tensor[(1, 64, 56, 56), float32] */, %p1: Tensor[(64, 64, 3, 3), float32] /* ty=Tensor[(64, 64, 3, 3), float32] */, %p2: Tensor[(1, 64, 54, 54), float32] /* ty=Tensor[(1, 64, 54, 54), float32] */, %p3: Tensor[(1, 64, 54, 54), float32] /* ty=Tensor[(1, 64, 54, 54), float32] */, Primitive=1) -> Tensor[(1, 64, 54, 54), float32] { %0 = nn.conv2d(%p0, %p1, padding=[0, 0, 0, 0]) /* ty=Tensor[(1, 64, 54, 54), float32] */; %1 = add(%0, %p2) /* ty=Tensor[(1, 64, 54, 54), float32] */; %2 = add(%1, %p3) /* ty=Tensor[(1, 64, 54, 54), float32] */; %3 = add(%1, %p3) /* ty=Tensor[(1, 64, 54, 54), float32] */; add(%2, %3) /* ty=Tensor[(1, 64, 54, 54), float32] */ } /* ty=fn (Tensor[(1, 64, 56, 56), float32], Tensor[(64, 64, 3, 3), float32], Tensor[(1, 64, 54, 54), float32], Tensor[(1, 64, 54, 54), float32]) -> Tensor[(1, 64, 54, 54), float32] */; %4(%x, %weight, meta[relay.Constant][0] /* ty=Tensor[(1, 64, 54, 54), float32] */, meta[relay.Constant][1] /* ty=Tensor[(1, 64, 54, 54), float32] */) /* ty=Tensor[(1, 64, 54, 54), float32] */ } .. GENERATED FROM PYTHON SOURCE LINES 158-165 From the transformed Relay program, we can see that there are still two 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.transform.Sequential`. The pass infra, however, provides a configuration interface for users to customize the optimization level that they want to execute. .. GENERATED FROM PYTHON SOURCE LINES 165-170 .. code-block:: default with tvm.transform.PassContext(opt_level=3): mod2 = seq(mod) print(mod2) .. rst-class:: sphx-glr-script-out .. code-block:: none def @main(%x: Tensor[(1, 64, 56, 56), float32] /* ty=Tensor[(1, 64, 56, 56), float32] */, %weight: Tensor[(64, 64, 3, 3), float32] /* ty=Tensor[(64, 64, 3, 3), float32] */) -> Tensor[(1, 64, 54, 54), float32] { %3 = fn (%p0: Tensor[(1, 64, 56, 56), float32] /* ty=Tensor[(1, 64, 56, 56), float32] */, %p1: Tensor[(64, 64, 3, 3), float32] /* ty=Tensor[(64, 64, 3, 3), float32] */, %p2: Tensor[(1, 64, 54, 54), float32] /* ty=Tensor[(1, 64, 54, 54), float32] */, %p3: Tensor[(1, 64, 54, 54), float32] /* ty=Tensor[(1, 64, 54, 54), float32] */, Primitive=1) -> Tensor[(1, 64, 54, 54), float32] { %0 = nn.conv2d(%p0, %p1, padding=[0, 0, 0, 0]) /* ty=Tensor[(1, 64, 54, 54), float32] */; %1 = add(%0, %p2) /* ty=Tensor[(1, 64, 54, 54), float32] */; %2 = add(%1, %p3) /* ty=Tensor[(1, 64, 54, 54), float32] */; add(%2, %2) /* ty=Tensor[(1, 64, 54, 54), float32] */ } /* ty=fn (Tensor[(1, 64, 56, 56), float32], Tensor[(64, 64, 3, 3), float32], Tensor[(1, 64, 54, 54), float32], Tensor[(1, 64, 54, 54), float32]) -> Tensor[(1, 64, 54, 54), float32] */; %3(%x, %weight, meta[relay.Constant][0] /* ty=Tensor[(1, 64, 54, 54), float32] */, meta[relay.Constant][1] /* ty=Tensor[(1, 64, 54, 54), float32] */) /* ty=Tensor[(1, 64, 54, 54), float32] */ } .. GENERATED FROM PYTHON SOURCE LINES 171-178 Now we can see that only one of the two identical additions is kept. In addition, users can selectively disable some passes using the `disabled_pass` config, which is similar to the `-fno-xxx` option used the general purpose compilers, such as Clang and GCC. For example, we can disable EliminateCommonSubexpr as following. The printed module will again show two identical addition operations. .. GENERATED FROM PYTHON SOURCE LINES 178-183 .. code-block:: default with tvm.transform.PassContext(opt_level=3, disabled_pass=["EliminateCommonSubexpr"]): mod3 = seq(mod) print(mod3) .. rst-class:: sphx-glr-script-out .. code-block:: none def @main(%x: Tensor[(1, 64, 56, 56), float32] /* ty=Tensor[(1, 64, 56, 56), float32] */, %weight: Tensor[(64, 64, 3, 3), float32] /* ty=Tensor[(64, 64, 3, 3), float32] */) -> Tensor[(1, 64, 54, 54), float32] { %4 = fn (%p0: Tensor[(1, 64, 56, 56), float32] /* ty=Tensor[(1, 64, 56, 56), float32] */, %p1: Tensor[(64, 64, 3, 3), float32] /* ty=Tensor[(64, 64, 3, 3), float32] */, %p2: Tensor[(1, 64, 54, 54), float32] /* ty=Tensor[(1, 64, 54, 54), float32] */, %p3: Tensor[(1, 64, 54, 54), float32] /* ty=Tensor[(1, 64, 54, 54), float32] */, Primitive=1) -> Tensor[(1, 64, 54, 54), float32] { %0 = nn.conv2d(%p0, %p1, padding=[0, 0, 0, 0]) /* ty=Tensor[(1, 64, 54, 54), float32] */; %1 = add(%0, %p2) /* ty=Tensor[(1, 64, 54, 54), float32] */; %2 = add(%1, %p3) /* ty=Tensor[(1, 64, 54, 54), float32] */; %3 = add(%1, %p3) /* ty=Tensor[(1, 64, 54, 54), float32] */; add(%2, %3) /* ty=Tensor[(1, 64, 54, 54), float32] */ } /* ty=fn (Tensor[(1, 64, 56, 56), float32], Tensor[(64, 64, 3, 3), float32], Tensor[(1, 64, 54, 54), float32], Tensor[(1, 64, 54, 54), float32]) -> Tensor[(1, 64, 54, 54), float32] */; %4(%x, %weight, meta[relay.Constant][0] /* ty=Tensor[(1, 64, 54, 54), float32] */, meta[relay.Constant][1] /* ty=Tensor[(1, 64, 54, 54), float32] */) /* ty=Tensor[(1, 64, 54, 54), float32] */ } .. GENERATED FROM PYTHON SOURCE LINES 184-194 Implement a Pass Using Python Decorator ------------------------------------------ The next example illustrates how we can orchestrate a customized optimization pipeline through the pass infra using Python decorators. This functionality greatly eases the implementation of passes. For example, users can simply define a decorated class to do function-level optimizations as the following example shows. `transform_function` wraps a class to replace all constants with a multiple of `c`. Later on, each function in a given module will be visited and each constant in the function will be replaced when we invoke the customized pass. .. GENERATED FROM PYTHON SOURCE LINES 194-221 .. code-block:: default @relay.transform.function_pass(opt_level=1) class CustomPipeline: """Simple test function to replace one argument to another.""" def __init__(self, multiplier): self.multiplier = multiplier # This function can define a pass. def transform_function(self, func, mod, ctx): obj = self class ReplaceConstant(tvm.relay.ExprMutator): def visit_constant(self, c): return relay.multiply(obj.multiplier, c) return ReplaceConstant().visit(func) f = example() mod = tvm.IRModule.from_expr(f) custom_pass = CustomPipeline(multiplier=relay.const(3, "float32")) assert custom_pass.info.name == "CustomPipeline" mod3 = custom_pass(mod) print(mod3) .. rst-class:: sphx-glr-script-out .. code-block:: none def @main(%x: Tensor[(1, 64, 56, 56), float32] /* ty=Tensor[(1, 64, 56, 56), float32] */, %weight: Tensor[(64, 64, 3, 3), float32] /* ty=Tensor[(64, 64, 3, 3), float32] */) -> Tensor[(1, 64, 54, 54), float32] { %0 = multiply(3f /* ty=float32 */, meta[relay.Constant][0] /* ty=Tensor[(1, 64, 54, 54), float32] */) /* ty=Tensor[(1, 64, 54, 54), float32] */; %1 = add(%0, %0) /* ty=Tensor[(1, 64, 54, 54), float32] */; %2 = multiply(3f /* ty=float32 */, 2f /* ty=float32 */) /* ty=float32 */; %3 = nn.conv2d(%x, %weight, padding=[0, 0, 0, 0]) /* ty=Tensor[(1, 64, 54, 54), float32] */; %4 = multiply(%1, %2) /* ty=Tensor[(1, 64, 54, 54), float32] */; %5 = add(%3, %4) /* ty=Tensor[(1, 64, 54, 54), float32] */; %6 = add(%5, %0) /* ty=Tensor[(1, 64, 54, 54), float32] */; %7 = add(%5, %0) /* ty=Tensor[(1, 64, 54, 54), float32] */; add(%6, %7) /* ty=Tensor[(1, 64, 54, 54), float32] */ } .. GENERATED FROM PYTHON SOURCE LINES 222-228 Debug a Pass ------------ TVM provides users a plug-and-play style debugging pass that print the IR after a certain pass is done through a special pass (``PrintIR``) to dump the IR of the whole module. A slightly modified version of the sequential pass example could be like the following to enable IR dumping for ``FoldConstant`` optimization. .. GENERATED FROM PYTHON SOURCE LINES 228-240 .. code-block:: default f = example() mod = tvm.IRModule.from_expr(f) seq = tvm.transform.Sequential( [ relay.transform.FoldConstant(), tvm.transform.PrintIR(), relay.transform.EliminateCommonSubexpr(), relay.transform.FuseOps(), ] ) .. GENERATED FROM PYTHON SOURCE LINES 241-252 By inserting the ``PrintIR`` pass after ``FoldConstant``, the pass infra will dump out the module IR when ``FoldConstant`` is done. Users can plug in this pass after any pass they want to debug for viewing the optimization effect. There is a more flexible debugging mechanism. One can implement a ``PassInstrument`` class to execute arbitrary code not only before and/or after each pass but also at entering/exiting ``PassContext``. See :ref:`pass_instrument_cpp_backend` for more details. Here we use :py::func`tvm.instrument.pass_instrument` decorator to implement a PassInsturment class printing IR before execution of each passes: .. GENERATED FROM PYTHON SOURCE LINES 252-271 .. code-block:: default @tvm.instrument.pass_instrument class PrintIR: """Print the name of the pass, the IR, only before passes execute.""" def run_before_pass(self, mod, info): print("Running pass: {}", info) print(mod) with tvm.transform.PassContext(opt_level=3, instruments=[PrintIR()]): with tvm.target.Target("llvm"): # Perform the optimizations. mod = seq(mod) print(mod) print("done") .. rst-class:: sphx-glr-script-out .. code-block:: none Running pass: {} The meta data of the pass - pass name: sequential, opt_level: 0, required passes: [] def @main(%x: Tensor[(1, 64, 56, 56), float32], %weight: Tensor[(64, 64, 3, 3), float32]) { %0 = add(meta[relay.Constant][0], meta[relay.Constant][0]); %1 = nn.conv2d(%x, %weight, padding=[0, 0, 0, 0]); %2 = multiply(%0, 2f); %3 = add(%1, %2); %4 = add(%3, meta[relay.Constant][0]); %5 = add(%3, meta[relay.Constant][0]); add(%4, %5) } Running pass: {} The meta data of the pass - pass name: FoldConstant, opt_level: 2, required passes: [] def @main(%x: Tensor[(1, 64, 56, 56), float32], %weight: Tensor[(64, 64, 3, 3), float32]) { %0 = add(meta[relay.Constant][0], meta[relay.Constant][0]); %1 = nn.conv2d(%x, %weight, padding=[0, 0, 0, 0]); %2 = multiply(%0, 2f); %3 = add(%1, %2); %4 = add(%3, meta[relay.Constant][0]); %5 = add(%3, meta[relay.Constant][0]); add(%4, %5) } Running pass: {} The meta data of the pass - pass name: InferType, opt_level: 0, required passes: [] def @main(%x: Tensor[(1, 64, 56, 56), float32], %weight: Tensor[(64, 64, 3, 3), float32]) { %0 = nn.conv2d(%x, %weight, padding=[0, 0, 0, 0]); %1 = add(%0, meta[relay.Constant][0]); %2 = add(%1, meta[relay.Constant][1]); %3 = add(%1, meta[relay.Constant][1]); add(%2, %3) } Running pass: {} The meta data of the pass - pass name: PrintIR, opt_level: 0, required passes: [] def @main(%x: Tensor[(1, 64, 56, 56), float32] /* ty=Tensor[(1, 64, 56, 56), float32] */, %weight: Tensor[(64, 64, 3, 3), float32] /* ty=Tensor[(64, 64, 3, 3), float32] */) -> Tensor[(1, 64, 54, 54), float32] { %0 = nn.conv2d(%x, %weight, padding=[0, 0, 0, 0]) /* ty=Tensor[(1, 64, 54, 54), float32] */; %1 = add(%0, meta[relay.Constant][0] /* ty=Tensor[(1, 64, 54, 54), float32] */) /* ty=Tensor[(1, 64, 54, 54), float32] */; %2 = add(%1, meta[relay.Constant][1] /* ty=Tensor[(1, 64, 54, 54), float32] */) /* ty=Tensor[(1, 64, 54, 54), float32] */; %3 = add(%1, meta[relay.Constant][1] /* ty=Tensor[(1, 64, 54, 54), float32] */) /* ty=Tensor[(1, 64, 54, 54), float32] */; add(%2, %3) /* ty=Tensor[(1, 64, 54, 54), float32] */ } Running pass: {} The meta data of the pass - pass name: InferType, opt_level: 0, required passes: [] def @main(%x: Tensor[(1, 64, 56, 56), float32] /* ty=Tensor[(1, 64, 56, 56), float32] */, %weight: Tensor[(64, 64, 3, 3), float32] /* ty=Tensor[(64, 64, 3, 3), float32] */) -> Tensor[(1, 64, 54, 54), float32] { %0 = nn.conv2d(%x, %weight, padding=[0, 0, 0, 0]) /* ty=Tensor[(1, 64, 54, 54), float32] */; %1 = add(%0, meta[relay.Constant][0] /* ty=Tensor[(1, 64, 54, 54), float32] */) /* ty=Tensor[(1, 64, 54, 54), float32] */; %2 = add(%1, meta[relay.Constant][1] /* ty=Tensor[(1, 64, 54, 54), float32] */) /* ty=Tensor[(1, 64, 54, 54), float32] */; %3 = add(%1, meta[relay.Constant][1] /* ty=Tensor[(1, 64, 54, 54), float32] */) /* ty=Tensor[(1, 64, 54, 54), float32] */; add(%2, %3) /* ty=Tensor[(1, 64, 54, 54), float32] */ } Running pass: {} The meta data of the pass - pass name: EliminateCommonSubexpr, opt_level: 3, required passes: [ InferType, ] def @main(%x: Tensor[(1, 64, 56, 56), float32] /* ty=Tensor[(1, 64, 56, 56), float32] */, %weight: Tensor[(64, 64, 3, 3), float32] /* ty=Tensor[(64, 64, 3, 3), float32] */) -> Tensor[(1, 64, 54, 54), float32] { %0 = nn.conv2d(%x, %weight, padding=[0, 0, 0, 0]) /* ty=Tensor[(1, 64, 54, 54), float32] */; %1 = add(%0, meta[relay.Constant][0] /* ty=Tensor[(1, 64, 54, 54), float32] */) /* ty=Tensor[(1, 64, 54, 54), float32] */; %2 = add(%1, meta[relay.Constant][1] /* ty=Tensor[(1, 64, 54, 54), float32] */) /* ty=Tensor[(1, 64, 54, 54), float32] */; %3 = add(%1, meta[relay.Constant][1] /* ty=Tensor[(1, 64, 54, 54), float32] */) /* ty=Tensor[(1, 64, 54, 54), float32] */; add(%2, %3) /* ty=Tensor[(1, 64, 54, 54), float32] */ } Running pass: {} The meta data of the pass - pass name: InferType, opt_level: 0, required passes: [] def @main(%x: Tensor[(1, 64, 56, 56), float32] /* ty=Tensor[(1, 64, 56, 56), float32] */, %weight: Tensor[(64, 64, 3, 3), float32] /* ty=Tensor[(64, 64, 3, 3), float32] */) -> Tensor[(1, 64, 54, 54), float32] { %0 = nn.conv2d(%x, %weight, padding=[0, 0, 0, 0]) /* ty=Tensor[(1, 64, 54, 54), float32] */; %1 = add(%0, meta[relay.Constant][0] /* ty=Tensor[(1, 64, 54, 54), float32] */) /* ty=Tensor[(1, 64, 54, 54), float32] */; %2 = add(%1, meta[relay.Constant][1] /* ty=Tensor[(1, 64, 54, 54), float32] */) /* ty=Tensor[(1, 64, 54, 54), float32] */; add(%2, %2) /* ty=Tensor[(1, 64, 54, 54), float32] */ } Running pass: {} The meta data of the pass - pass name: InferType, opt_level: 0, required passes: [] def @main(%x: Tensor[(1, 64, 56, 56), float32] /* ty=Tensor[(1, 64, 56, 56), float32] */, %weight: Tensor[(64, 64, 3, 3), float32] /* ty=Tensor[(64, 64, 3, 3), float32] */) -> Tensor[(1, 64, 54, 54), float32] { %0 = nn.conv2d(%x, %weight, padding=[0, 0, 0, 0]) /* ty=Tensor[(1, 64, 54, 54), float32] */; %1 = add(%0, meta[relay.Constant][0] /* ty=Tensor[(1, 64, 54, 54), float32] */) /* ty=Tensor[(1, 64, 54, 54), float32] */; %2 = add(%1, meta[relay.Constant][1] /* ty=Tensor[(1, 64, 54, 54), float32] */) /* ty=Tensor[(1, 64, 54, 54), float32] */; add(%2, %2) /* ty=Tensor[(1, 64, 54, 54), float32] */ } Running pass: {} The meta data of the pass - pass name: FuseOps, opt_level: 0, required passes: [ InferType, ] def @main(%x: Tensor[(1, 64, 56, 56), float32] /* ty=Tensor[(1, 64, 56, 56), float32] */, %weight: Tensor[(64, 64, 3, 3), float32] /* ty=Tensor[(64, 64, 3, 3), float32] */) -> Tensor[(1, 64, 54, 54), float32] { %0 = nn.conv2d(%x, %weight, padding=[0, 0, 0, 0]) /* ty=Tensor[(1, 64, 54, 54), float32] */; %1 = add(%0, meta[relay.Constant][0] /* ty=Tensor[(1, 64, 54, 54), float32] */) /* ty=Tensor[(1, 64, 54, 54), float32] */; %2 = add(%1, meta[relay.Constant][1] /* ty=Tensor[(1, 64, 54, 54), float32] */) /* ty=Tensor[(1, 64, 54, 54), float32] */; add(%2, %2) /* ty=Tensor[(1, 64, 54, 54), float32] */ } Running pass: {} The meta data of the pass - pass name: InferType, opt_level: 0, required passes: [] def @main(%x: Tensor[(1, 64, 56, 56), float32] /* ty=Tensor[(1, 64, 56, 56), float32] */, %weight: Tensor[(64, 64, 3, 3), float32] /* ty=Tensor[(64, 64, 3, 3), float32] */) -> Tensor[(1, 64, 54, 54), float32] { %3 = fn (%p0: Tensor[(1, 64, 56, 56), float32], %p1: Tensor[(64, 64, 3, 3), float32], %p2: Tensor[(1, 64, 54, 54), float32], %p3: Tensor[(1, 64, 54, 54), float32], Primitive=1) -> Tensor[(1, 64, 54, 54), float32] { %0 = nn.conv2d(%p0, %p1, padding=[0, 0, 0, 0]); %1 = add(%0, %p2); %2 = add(%1, %p3); add(%2, %2) }; %3(%x, %weight, meta[relay.Constant][0] /* ty=Tensor[(1, 64, 54, 54), float32] */, meta[relay.Constant][1] /* ty=Tensor[(1, 64, 54, 54), float32] */) } def @main(%x: Tensor[(1, 64, 56, 56), float32] /* ty=Tensor[(1, 64, 56, 56), float32] */, %weight: Tensor[(64, 64, 3, 3), float32] /* ty=Tensor[(64, 64, 3, 3), float32] */) -> Tensor[(1, 64, 54, 54), float32] { %3 = fn (%p0: Tensor[(1, 64, 56, 56), float32] /* ty=Tensor[(1, 64, 56, 56), float32] */, %p1: Tensor[(64, 64, 3, 3), float32] /* ty=Tensor[(64, 64, 3, 3), float32] */, %p2: Tensor[(1, 64, 54, 54), float32] /* ty=Tensor[(1, 64, 54, 54), float32] */, %p3: Tensor[(1, 64, 54, 54), float32] /* ty=Tensor[(1, 64, 54, 54), float32] */, Primitive=1) -> Tensor[(1, 64, 54, 54), float32] { %0 = nn.conv2d(%p0, %p1, padding=[0, 0, 0, 0]) /* ty=Tensor[(1, 64, 54, 54), float32] */; %1 = add(%0, %p2) /* ty=Tensor[(1, 64, 54, 54), float32] */; %2 = add(%1, %p3) /* ty=Tensor[(1, 64, 54, 54), float32] */; add(%2, %2) /* ty=Tensor[(1, 64, 54, 54), float32] */ } /* ty=fn (Tensor[(1, 64, 56, 56), float32], Tensor[(64, 64, 3, 3), float32], Tensor[(1, 64, 54, 54), float32], Tensor[(1, 64, 54, 54), float32]) -> Tensor[(1, 64, 54, 54), float32] */; %3(%x, %weight, meta[relay.Constant][0] /* ty=Tensor[(1, 64, 54, 54), float32] */, meta[relay.Constant][1] /* ty=Tensor[(1, 64, 54, 54), float32] */) /* ty=Tensor[(1, 64, 54, 54), float32] */ } done .. GENERATED FROM PYTHON SOURCE LINES 272-280 Summary ------- This tutorial has covered how we can write and invoke passes in TVM more conveniently using the pass infra. Different ways of invoking a pass are also discussed. Using :py:class:`tvm.transform.Sequential` can largely help users to ease the work of handling multiple optimization passes and their dependencies. In addition, an example is provided to illustrate how we can debug a pass using the ``PrintIR`` and tracing. .. _sphx_glr_download_how_to_extend_tvm_use_pass_infra.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: use_pass_infra.py ` .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: use_pass_infra.ipynb ` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_