[Bug] [Relay] [MetaSchedule] ValueError: The block no longer exists in the IRModule

[Civitasv]

I'm trying to use MetaSchedule to tune vae_decoder model. But it failed to extract tasks.

Expected behavior

It should successful.

Actual behavior

  File "/home/xxx/xxx/tvm_tunning/tunningx/tunning/meta_schedule.py", line 67, in do_all_tune
    database = ms.relay_integration.tune_relay(
  File "/home/xxx/xxx/tvm/python/tvm/meta_schedule/relay_integration.py", line 352, in tune_relay
    return tune_tasks(
  File "/home/xxx/xxx/tvm/python/tvm/meta_schedule/tune.py", line 118, in tune_tasks
    task_scheduler.tune(
  File "/home/xxx/xxx/tvm/python/tvm/meta_schedule/task_scheduler/task_scheduler.py", line 132, in tune
    _ffi_api.TaskSchedulerTune(  # type: ignore # pylint: disable=no-member
  File "/home/xxx/xxx/tvm/python/tvm/_ffi/_ctypes/packed_func.py", line 238, in __call__
    raise get_last_ffi_error()
ValueError: Traceback (most recent call last):
  11: TVMFuncCall
  10: _ZN3tvm7runtime13PackedFunc
  9: tvm::runtime::TypedPackedFunc<void (tvm::meta_schedule::TaskScheduler, tvm::runtime::Array<tvm::meta_schedule::TuneContext, void>, tvm::runtime::Array<tvm::FloatImm, void>, int, int, int, tvm::meta_schedule::Builder, tvm::meta_schedule::Runner, tvm::runtime::Array<tvm::meta_schedule::MeasureCallback, void>, tvm::runtime::Optional<tvm::meta_schedule::Database>, tvm::runtime::Optional<tvm::meta_schedule::CostModel>)>::AssignTypedLambda<tvm::runtime::Registry::set_body_method<tvm::meta_schedule::TaskScheduler, tvm::meta_schedule::TaskSchedulerNode, void, tvm::runtime::Array<tvm::meta_schedule::TuneContext, void>, tvm::runtime::Array<tvm::FloatImm, void>, int, int, int, tvm::meta_schedule::Builder, tvm::meta_schedule::Runner, tvm::runtime::Array<tvm::meta_schedule::MeasureCallback, void>, tvm::runtime::Optional<tvm::meta_schedule::Database>, tvm::runtime::Optional<tvm::meta_schedule::CostModel>, void>(void (tvm::meta_schedule::TaskSchedulerNode::*)(tvm::runtime::Array<tvm::meta_schedule::TuneContext, void>, tvm::runtime::Array<tvm::FloatImm, void>, int, int, int, tvm::meta_schedule::Builder, tvm::meta_schedule::Runner, tvm::runtime::Array<tvm::meta_schedule::MeasureCallback, void>, tvm::runtime::Optional<tvm::meta_schedule::Database>, tvm::runtime::Optional<tvm::meta_schedule::CostModel>))::{lambda(tvm::meta_schedule::TaskScheduler, tvm::runtime::Array<tvm::meta_schedule::TuneContext, void>, tvm::runtime::Array<tvm::FloatImm, void>, int, int, int, tvm::meta_schedule::Builder, tvm::meta_schedule::Runner, tvm::runtime::Array<tvm::meta_schedule::MeasureCallback, void>, tvm::runtime::Optional<tvm::meta_schedule::Database>, tvm::runtime::Optional<tvm::meta_schedule::CostModel>)#1}>(tvm::runtime::Registry::set_body_method<tvm::meta_schedule::TaskScheduler, tvm::meta_schedule::TaskSchedulerNode, void, tvm::runtime::Array<tvm::meta_schedule::TuneContext, void>, tvm::runtime::Array<tvm::FloatImm, void>, int, int, int, tvm::meta_schedule::Builder, tvm::meta_schedule::Runner, tvm::runtime::Array<tvm::meta_schedule::MeasureCallback, void>, tvm::runtime::Optional<tvm::meta_schedule::Database>, tvm::runtime::Optional<tvm::meta_schedule::CostModel>, void>(void (tvm::meta_schedule::TaskSchedulerNode::*)(tvm::runtime::Array<tvm::meta_schedule::TuneContext, void>, tvm::runtime::Array<tvm::FloatImm, void>, int, int, int, tvm::meta_schedule::Builder, tvm::meta_schedule::Runner, tvm::runtime::Array<tvm::meta_schedule::MeasureCallback, void>, tvm::runtime::Optional<tvm::meta_schedule::Database>, tvm::runtime::Optional<tvm::meta_schedule::CostModel>))::{lambda(tvm::meta_schedule::TaskScheduler, tvm::runtime::Array<tvm::meta_schedule::TuneContext, void>, tvm::runtime::Array<tvm::FloatImm, void>, int, int, int, tvm::meta_schedule::Builder, tvm::meta_schedule::Runner, tvm::runtime::Array<tvm::meta_schedule::MeasureCallback, void>, tvm::runtime::Optional<tvm::meta_schedule::Database>, tvm::runtime::Optional<tvm::meta_schedule::CostModel>)#1}, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >)::{lambda(tvm::runtime::TVMArgs const&, tvm::runtime::TVMRetValue*)#1}::operator()(tvm::runtime::TVMArgs const&, tvm::runtime::TVMRetValue*) const [clone .isra.0]
  8: tvm::meta_schedule::GradientBasedNode::Tune(tvm::runtime::Array<tvm::meta_schedule::TuneContext, void>, tvm::runtime::Array<tvm::FloatImm, void>, int, int, int, tvm::meta_schedule::Builder, tvm::meta_schedule::Runner, tvm::runtime::Array<tvm::meta_schedule::MeasureCallback, void>, tvm::runtime::Optional<tvm::meta_schedule::Database>, tvm::runtime::Optional<tvm::meta_schedule::CostModel>)
  7: tvm::meta_schedule::TaskSchedulerNode::Tune(tvm::runtime::Array<tvm::meta_schedule::TuneContext, void>, tvm::runtime::Array<tvm::FloatImm, void>, int, int, int, tvm::meta_schedule::Builder, tvm::meta_schedule::Runner, tvm::runtime::Array<tvm::meta_schedule::MeasureCallback, void>, tvm::runtime::Optional<tvm::meta_schedule::Database>, tvm::runtime::Optional<tvm::meta_schedule::CostModel>)
  6: tvm::meta_schedule::PostOrderApplyNode::GenerateDesignSpace(tvm::IRModule const&)
  5: tvm::meta_schedule::MultiLevelTilingTensorCoreNode::Apply(tvm::tir::Schedule const&, tvm::tir::BlockRV const&)
  4: tvm::meta_schedule::MultiLevelTilingTensorCoreNode::ApplySubRules(std::vector<tvm::meta_schedule::State, std::allocator<tvm::meta_schedule::State> >)
  3: tvm::meta_schedule::MultiLevelTilingTensorCoreNode::AddWriteReuseTensorCore(tvm::meta_schedule::TensorCoreState) const
  2: tvm::tir::TracedScheduleNode::ReverseComputeInline(tvm::tir::BlockRV const&)
  1: tvm::tir::ConcreteScheduleNode::ReverseComputeInline(tvm::tir::BlockRV const&)
  0: tvm::tir::ConcreteScheduleNode::GetSRef(tvm::tir::BlockRV const&) const
  File "/home/xxx/xxx/tvm/src/tir/schedule/./concrete_schedule.h", line 285
ValueError: The block no longer exists in the IRModule

I think ValueError: The block no longer exists in the IRModule is the reason.

Environment

TVM Version: Unity Latest
OS: Ubuntu 20.04

Triage

• core:ffi

cc @ibsidorenko @shingjan

[nautasolva]

I managed to reproduce this issue:

import tempfile

import tvm
import tvm.relay as relay

# cin = 17 = 16 + 1, so padding is necessary for fitting in a WMMA fragment
# bug is not present when cin = 16
cin, cout = 17, 16
h, w = 1, 1

inp = relay.var("input", shape=(1, cin, h, w), dtype="float16")
w = relay.var("weight", shape=(cout, cin, 1, 1), dtype="float16")
b = relay.var("bias", shape=(cout,), dtype="float16")

out = relay.nn.bias_add(relay.nn.conv2d(inp, w), b)

mod = tvm.IRModule.from_expr(out)

import tvm.meta_schedule
target = tvm.target.Target({
        "kind": "cuda",
        "arch": "sm_80",
        "max_threads_per_block": 1024,
        "max_shared_memory_per_block": 49152,
})

with tempfile.TemporaryDirectory() as tuning_dir:
    tasks = tvm.meta_schedule.relay_integration.tune_relay(tvm.IRModule.from_expr(out), params={}, target=target, work_dir=tuning_dir, max_trials_global=32)

The error is triggered when the MultiLevelTilingTensorCore rule tries to inline a block that was previously inlined (hence does not exist anymore). For example, when trying to add padding so that we can fit a WMMA fragment, we inline consumers of the newly created padding blocks, effectively consuming the block.
The solution I found was to guard the inlining operation.

Please let met know if this fix is acceptable, I can submit a pull request.

[Civitasv]

Can confirm this fix my problem, much appreciate!