[Bug] Check failed when using MetaSchedule

[Fichimio]

Expected behavior

The code ran successfully on TVM 0.9.0 with following message printed.

2023-02-14 23:25:47.417 INFO [Updated] Task #0: "main"
 ID | Name |    FLOP | Weight | Speed (GFLOPS) | Latency (us) | Weighted Latency (us) | Trials | Terminated 
------------------------------------------------------------------------------------------------------------
  0 | main | 4194304 |      1 |       108.7843 |      38.5562 |               38.5562 |     64 |          
------------------------------------------------------------------------------------------------------------
Total trials: 64
Total latency (us): 38.5562

2023-02-14 23:25:47.419 INFO Task #0 has finished. Remaining task(s): 0
2023-02-14 23:25:47.514 INFO Saved XGBModel to ./tune_tmp/cost_model.xgb

Actual behavior

The code failed to run with following error message

  tvm._ffi.base.TVMError: Traceback (most recent call last):
  7: TVMFuncCall
  6: tvm::runtime::PackedFuncObj::Extractor<tvm::runtime::PackedFuncSubObj<tvm::runtime::TypedPackedFunc<void (tvm::meta_schedule::TaskScheduler)>::AssignTypedLambda<tvm::runtime::Registry::set_body_method<tvm::meta_schedule::TaskScheduler, tvm::meta_schedule::TaskSchedulerNode, void, , void>(void (tvm::meta_schedule::TaskSchedulerNode::*)())::{lambda(tvm::meta_schedule::TaskScheduler)#1}>(tvm::runtime::Registry::set_body_method<tvm::meta_schedule::TaskScheduler, tvm::meta_schedule::TaskSchedulerNode, void, , void>(void (tvm::meta_schedule::TaskSchedulerNode::*)())::{lambda(tvm::meta_schedule::TaskScheduler)#1}, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >)::{lambda(tvm::runtime::TVMArgs const&, tvm::runtime::TVMRetValue*)#1}> >::Call(tvm::runtime::PackedFuncObj const*, tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*)
  5: tvm::meta_schedule::TaskSchedulerNode::Tune()
  4: tvm::meta_schedule::TaskSchedulerNode::InitializeTask(int)
  3: tvm::tir::TraceNode::AsPython(bool) const
  2: _ZN3tvm7runtime13PackedFuncObj
  1: tvm::runtime::TypedPackedFunc<tvm::runtime::String (tvm::runtime::Array<tvm::runtime::ObjectRef, void> const&, tvm::runtime::Array<tvm::runtime::ObjectRef, void> const&, tvm::runtime::Optional<tvm::runtime::ObjectRef> const&, tvm::runtime::Array<tvm::runtime::String, void> const&)>::AssignTypedLambda<tvm::runtime::String (*)(tvm::runtime::Array<tvm::runtime::ObjectRef, void> const&, tvm::runtime::Array<tvm::runtime::ObjectRef, void> const&, tvm::runtime::Optional<tvm::runtime::ObjectRef> const&, tvm::runtime::Array<tvm::runtime::String, void> const&)>(tvm::runtime::String (*)(tvm::runtime::Array<tvm::runtime::ObjectRef, void> const&, tvm::runtime::Array<tvm::runtime::ObjectRef, void> const&, tvm::runtime::Optional<tvm::runtime::ObjectRef> const&, tvm::runtime::Array<tvm::runtime::String, void> const&))::{lambda(tvm::runtime::TVMArgs const&, tvm::runtime::TVMRetValue*)#1}::operator()(tvm::runtime::TVMArgs const&, tvm::runtime::TVMRetValue*) const
  0: tvm::tir::UnpackedInstTraits<tvm::tir::ReorderTraits>::AsPython(tvm::runtime::Array<tvm::runtime::ObjectRef, void> const&, tvm::runtime::Array<tvm::runtime::ObjectRef, void> const&, tvm::runtime::Optional<tvm::runtime::ObjectRef> const&, tvm::runtime::Array<tvm::runtime::String, void> const&)
  File "/sparsetir/src/tir/schedule/primitive/.././instruction_traits.h", line 372
TVMError:
---------------------------------------------------------------
An error occurred during the execution of TVM.
For more information, please see: https://tvm.apache.org/docs/errors.html
---------------------------------------------------------------

  Check failed: kNumInputs == inputs.size() (1 vs. 10) : ValueError: Incorrect kNumInputs for instruction: Reorder

Environment

Ubuntu 20.04.1

This error occurs on SparseTIR but doesn't occur on TVM 0.9.0, 0.10.0, 0.11.dev0

Steps to reproduce

import tvm.testing
import tvm.tir as tir
from tvm.script import tir as T
from tvm import meta_schedule as ms

print(tvm.__version__)

@tvm.script.ir_module
class MyModule:
    @T.prim_func
    def main(
        A: T.Buffer[(128, 128), "float32"],
        B: T.Buffer[(128, 128), "float32"],
        C: T.Buffer[(128, 128), "float32"],
    ):
        T.func_attr({"global_symbol": "main", "tir.noalias": True})
        for i, j, k in T.grid(128, 128, 128):
            with T.block("C"):
                vi, vj, vk = T.axis.remap("SSR", [i, j, k])
                with T.init():
                    C[vi, vj] = 0.0
                C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]

sch_tuned = ms.tune_tir(
    mod=MyModule,
    target="llvm --num-cores=1",
    config=ms.TuneConfig(
      max_trials_global=64,
      num_trials_per_iter=64,
    ),
    work_dir="./tune_tmp",
    task_name="main",
)
print(sch_tuned.trace)

Above code comes from a note of MLC course

[yzh119]

Hi @Fichimio, we haven't tested compatibility with MetaSchedule and currently only support searching for the best parameters for hand-crafted schedule templates.

Our codebase is forked from TVM before the release of TVM 0.9 and we don't guarantee all features in TVM 0.9 are supported.