Inconsistent result of the same pass optimization

Hello community, I am studying pass optimization recently. When I use tvm.transform.Sequential() to wrapper passes, it works well. But if I use the pass to optimize the module separately, it will throw an error. I’m curious the reason why there exists an inconsistency? I think the two methods should have the same result while actually not. I’d appreciate any help :smiley:

from tvm import relay
import tvm
from tvm.relay import testing

def example(batch_dim=1):
    out_channels = 32

    data = relay.var("data", relay.TensorType((batch_dim, 3, 224, 224), "float32"))
    weight = relay.var("weight")
    bn_gamma = relay.var("bn_gamma")
    bn_beta = relay.var("bn_beta")
    bn_mmean = relay.var("bn_mean")
    bn_mvar = relay.var("bn_var")

    simple_net = relay.nn.conv2d(
        data=data, weight=weight, kernel_size=(5, 5), channels=out_channels, padding=(1, 1)
    )
    simple_net = relay.nn.batch_norm(simple_net, bn_gamma, bn_beta, bn_mmean, bn_mvar)[0]
    simple_net = relay.nn.relu(simple_net)
    simple_net = relay.Function(relay.analysis.free_vars(simple_net), simple_net)

    return testing.create_workload(simple_net)

def test1():
    module, params = example()
    target = tvm.target.Target('llvm')

    with tvm.transform.PassContext(opt_level=0):
        module, params = relay.optimize(module, target, params)

    with tvm.transform.PassContext(opt_level=4):
        with target:
            tvm.transform.Sequential(
                [
                    relay.transform.ToBasicBlockNormalForm(),
                    relay.transform.FuseOps()
                ], opt_level = 4
            )(module)

    print('test1 run successfully')

def test2():
    module, params = example()
    target = tvm.target.Target('llvm')

    with tvm.transform.PassContext(opt_level=0):
        module, params = relay.optimize(module, target, params)

    with tvm.transform.PassContext(opt_level=4):
        with target:
            module2 = relay.transform.ToBasicBlockNormalForm()(module)
            module2 = relay.transform.FuseOps()(module2)

        print('test2 run successfully')
test1()
# test1 execute successfully
test2()
# test2 fail

Here is the error message:

> Traceback (most recent call last):
  File "test.py", line 57, in <module>
    test2()
  File "test.py", line 52, in test2
    module2 = relay.transform.FuseOps()(module2)
  File "/home/syang/tvm-cov/python/tvm/ir/transform.py", line 161, in __call__
    return _ffi_transform_api.RunPass(self, mod)
  File "/home/syang/tvm-cov/python/tvm/_ffi/_ctypes/packed_func.py", line 237, in __call__
    raise get_last_ffi_error()
tvm._ffi.base.TVMError: Traceback (most recent call last):
  18: TVMFuncCall
  17: std::_Function_handler<void (tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*), tvm::runtime::TypedPackedFunc<tvm::IRModule (tvm::transform::Pass, tvm::IRModule)>::AssignTypedLambda<tvm::transform::$_6>(tvm::transform::$_6, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >)::{lambda(tvm::runtime::TVMArgs const&, tvm::runtime::TVMRetValue*)#1}>::_M_invoke(std::_Any_data const&, tvm::runtime::TVMArgs&&, tvm::runtime::TVMRetValue*&&)
  16: tvm::transform::Pass::operator()(tvm::IRModule) const
  15: tvm::transform::Pass::operator()(tvm::IRModule, tvm::transform::PassContext const&) const
  14: tvm::relay::transform::FunctionPassNode::operator()(tvm::IRModule, tvm::transform::PassContext const&) const
  13: std::_Function_handler<void (tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*), tvm::runtime::TypedPackedFunc<tvm::relay::Function (tvm::relay::Function, tvm::IRModule, tvm::transform::PassContext)>::AssignTypedLambda<tvm::relay::transform::FuseOps(int)::$_0>(tvm::relay::transform::FuseOps(int)::$_0)::{lambda(tvm::runtime::TVMArgs const&, tvm::runtime::TVMRetValue*)#1}>::_M_invoke(std::_Any_data const&, tvm::runtime::TVMArgs&&, tvm::runtime::TVMRetValue*&&)
  12: tvm::relay::FuseOps(tvm::RelayExpr const&, int, unsigned long, tvm::IRModule const&)
  11: tvm::relay::FuseMutator::Transform(tvm::RelayExpr const&, int, unsigned long)
  10: tvm::relay::IndexedForwardGraph::Create(tvm::support::GenericArena<tvm::support::SimplePageAllocator>*, tvm::RelayExpr const&)
  9: tvm::relay::ExprVisitor::VisitExpr(tvm::RelayExpr const&)
  8: tvm::relay::ExprFunctor<void (tvm::RelayExpr const&)>::VisitExpr(tvm::RelayExpr const&)
  7: tvm::NodeFunctor<void (tvm::runtime::ObjectRef const&, tvm::relay::ExprFunctor<void (tvm::RelayExpr const&)>*)>::operator()(tvm::runtime::ObjectRef const&, tvm::relay::ExprFunctor<void (tvm::RelayExpr const&)>*) const
  6: tvm::relay::IndexedForwardGraph::Creator::VisitExpr_(tvm::relay::FunctionNode const*)
  5: tvm::relay::ExprVisitor::VisitExpr_(tvm::relay::FunctionNode const*)
  4: tvm::relay::ExprVisitor::VisitExpr(tvm::RelayExpr const&)
  3: tvm::relay::ExprFunctor<void (tvm::RelayExpr const&)>::VisitExpr(tvm::RelayExpr const&)
  2: tvm::NodeFunctor<void (tvm::runtime::ObjectRef const&, tvm::relay::ExprFunctor<void (tvm::RelayExpr const&)>*)>::operator()(tvm::runtime::ObjectRef const&, tvm::relay::ExprFunctor<void (tvm::RelayExpr const&)>*) const
  1: tvm::relay::IndexedForwardGraph::Creator::VisitExpr_(tvm::relay::CallNode const*)
  0: tvm::RelayExprNode::checked_type() const
  File "/home/syang/tvm-cov/include/tvm/ir/expr.h", line 475
TVMError: 
---------------------------------------------------------------
An error occurred during the execution of TVM.
For more information, please see: https://tvm.apache.org/docs/errors.html
---------------------------------------------------------------
  Check failed: (checked_type_.defined()) is false: internal error: the type checker has not populated the checked_type field for CallNode(fn (%p0: Tensor[(1, 32, 222, 222), float32], Primitive=1, hash="548642771bf3d9e6") -> Tensor[(1, 32, 222, 222), float32] {
  nn.relu(%p0) /* ty=Tensor[(1, 32, 222, 222), float32] */
}, [CallNode(fn (%p0: Tensor[(1, 32, 222, 222), float32], %p1: Tensor[(32, 1, 1), float32], Primitive=1, hash="e07496501583cc09") -> Tensor[(1, 32, 222, 222), float32] {
  add(%p0, %p1) /* ty=Tensor[(1, 32, 222, 222), float32] */
}, [CallNode(fn (%p0: Tensor[(1, 32, 222, 222), float32], %p1: Tensor[(32, 1, 1), float32], Primitive=1, hash="7bc711a654d34e5b") -> Tensor[(1, 32, 222, 222), float32] {
  multiply(%p0, %p1) /* ty=Tensor[(1, 32, 222, 222), float32] */
}, [CallNode(fn (%p0: Tensor[(1, 3, 224, 224), float32], %p1: Tensor[(32, 3, 5, 5), float32], hash="7a783d96ca6c3b54", out_layout="", kernel_layout="OIHW", Primitive=1, data_layout="NCHW") -> Tensor[(1, 32, 222, 222), float32] {
  nn.conv2d(%p0, %p1, padding=[1, 1, 1, 1], channels=32, kernel_size=[5, 5]) /* ty=Tensor[(1, 32, 222, 222), float32] */
}, [Var(data, ty=TensorType([1, 3, 224, 224], float32))
......

Because pass sequential will not only run the passes you specified, but also run some other passes are marked as required for the pass (e.g., InferType). Specifically, your error message shows that the module doesn’t run InferType before running FuseOps. If you run the pass sequential, it actually runs ToBasicBlockNormalFormInferTypeFuseOps, as indicated here:

Thanks for your answer! It helps me understand pass optimization better. I thought the tvm.transform.Sequential() is just a wrapper before :rofl: