Inconsistent result of the same pass optimization

syang · August 18, 2021, 12:23am

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

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))
......

comaniac · August 18, 2021, 12:41am

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 ToBasicBlockNormalForm → InferType → FuseOps, as indicated here:

github.com

apache/tvm/blob/main/src/relay/transforms/fuse_ops.cc#L1031




namespace transform {


Pass FuseOps(int fuse_opt_level) {
  runtime::TypedPackedFunc<Function(Function, IRModule, PassContext)> pass_func =
      [=](Function f, IRModule m, PassContext pc) {
        int opt_level = fuse_opt_level == -1 ? pc->opt_level : fuse_opt_level;
        auto max_fuse_depth = pc->GetConfig("relay.FuseOps.max_depth", Integer(kMaxFusedOps));
        return Downcast<Function>(FuseOps(f, opt_level, max_fuse_depth.value(), m));
      };
  return CreateFunctionPass(pass_func, 1, "FuseOps", {"InferType"});
}


TVM_REGISTER_GLOBAL("relay._transform.FuseOps").set_body_typed(FuseOps);


}  // namespace transform


}  // namespace relay
}  // namespace tvm

syang · August 18, 2021, 3:07pm

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