[Bug] MetaSchedule PerStoreFeature crashes when processing zero-extent loops #17994
Description
Actual behavior
Traceback (most recent call last):
File "/data/qshenaf/remote_pc/TirFuzz/bugs/05-19_15-37/topi.take_2.py", line 12, in <module>
database = ms.tir_integration.tune_tir(mod=sch.mod, target='llvm --num-cores=16', work_dir='./tune_tmp', max_trials_global=1, num_trials_per_iter=1)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/data/qshenaf/envs/tvm/python/tvm/meta_schedule/tir_integration.py", line 146, in tune_tir
return tune_tasks(
^^^^^^^^^^^
File "/data/qshenaf/envs/tvm/python/tvm/meta_schedule/tune.py", line 122, in tune_tasks
task_scheduler.tune(
File "/data/qshenaf/envs/tvm/python/tvm/meta_schedule/task_scheduler/task_scheduler.py", line 132, in tune
_ffi_api.TaskSchedulerTune( # type: ignore # pylint: disable=no-member
File "tvm/_ffi/_cython/./packed_func.pxi", line 339, in tvm._ffi._cy3.core.PackedFuncBase.__call__
File "tvm/_ffi/_cython/./packed_func.pxi", line 284, in tvm._ffi._cy3.core.FuncCall
File "tvm/_ffi/_cython/./base.pxi", line 185, in tvm._ffi._cy3.core.CHECK_CALL
File "/data/qshenaf/envs/tvm/python/tvm/_ffi/base.py", line 468, in raise_last_ffi_error
raise py_err
File "/data/qshenaf/envs/tvm/src/meta_schedule/task_scheduler/gradient_based.cc", line 54, in 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>)
TaskSchedulerNode::Tune(tasks, task_weights, max_trials_global, max_trials_per_task,
^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/data/qshenaf/envs/tvm/src/meta_schedule/task_scheduler/task_scheduler.cc", line 205, in 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>)
JoinRunningTask(task_id);
^^^^^^^^^^^^^^^^
File "/data/qshenaf/envs/tvm/src/meta_schedule/task_scheduler/gradient_based.cc", line 126, in tvm::meta_schedule::GradientBasedNode::JoinRunningTask(int)
Array<RunnerResult> results = TaskSchedulerNode::JoinRunningTask(task_id);
^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/data/qshenaf/envs/tvm/src/meta_schedule/task_scheduler/task_scheduler.cc", line 232, in tvm::meta_schedule::TaskSchedulerNode::JoinRunningTask(int)
callback->Apply(GetRef<TaskScheduler>(this), task_id, task->measure_candidates.value(),
^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/data/qshenaf/envs/tvm/src/meta_schedule/measure_callback/update_cost_model.cc", line 53, in tvm::meta_schedule::UpdateCostModelNode::Apply(tvm::meta_schedule::TaskScheduler const&, int, tvm::runtime::Array<tvm::meta_schedule::MeasureCandidate, void> const&, tvm::runtime::Array<tvm::meta_schedule::BuilderResult, void> const&, tvm::runtime::Array<tvm::meta_schedule::RunnerResult, void> const&)
cost_model->Update(task->ctx, pruned_candidate, pruned_runner_result);
^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/data/qshenaf/envs/tvm/src/meta_schedule/cost_model/cost_model.cc", line 37, in tvm::meta_schedule::PyCostModelNode::Update(tvm::meta_schedule::TuneContext const&, tvm::runtime::Array<tvm::meta_schedule::MeasureCandidate, void> const&, tvm::runtime::Array<tvm::meta_schedule::RunnerResult, void> const&)
f_update(context, candidates, results);
^^^^^^^^^^^^^^^^^^^^^^^^
File "tvm/_ffi/_cython/./packed_func.pxi", line 56, in tvm._ffi._cy3.core.tvm_callback
File "/data/qshenaf/envs/tvm/python/tvm/meta_schedule/utils.py", line 76, in method
return getattr(inst, name)(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/data/qshenaf/envs/tvm/python/tvm/meta_schedule/cost_model/xgb_model.py", line 495, in update
new_features = [_feature(x) for x in self.extractor.extract_from(context, candidates)]
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/data/qshenaf/envs/tvm/python/tvm/meta_schedule/feature_extractor/feature_extractor.py", line 58, in extract_from
result = _ffi_api.FeatureExtractorExtractFrom( # type: ignore # pylint: disable=no-member
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "tvm/_ffi/_cython/./packed_func.pxi", line 339, in tvm._ffi._cy3.core.PackedFuncBase.__call__
File "tvm/_ffi/_cython/./packed_func.pxi", line 270, in tvm._ffi._cy3.core.FuncCall
File "tvm/_ffi/_cython/./packed_func.pxi", line 259, in tvm._ffi._cy3.core.FuncCall3
File "tvm/_ffi/_cython/./base.pxi", line 185, in tvm._ffi._cy3.core.CHECK_CALL
File "/data/qshenaf/envs/tvm/src/meta_schedule/feature_extractor/per_store_feature.cc", line 1414, in tvm::meta_schedule::PerStoreFeatureNode::ExtractFrom(tvm::meta_schedule::TuneContext const&, tvm::runtime::Array<tvm::meta_schedule::MeasureCandidate, void> const&)
support::parallel_for_dynamic(0, candidates.size(), tune_context->num_threads, f);
^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/data/qshenaf/envs/tvm/src/support/parallel_for.cc", line 139, in tvm::support::parallel_for_dynamic(int, int, int, std::function<void (int, int)> const&)
LOG(FATAL) << "RuntimeError: parallel_for_dynamic error with " << e.what();
^^^^^^^^^^^^^^^^^^^^^^^^^^^
tvm._ffi.base.TVMError: Traceback (most recent call last):
1: tvm::meta_schedule::PerStoreFeatureNode::ExtractFrom(tvm::meta_schedule::TuneContext const&, tvm::runtime::Array<tvm::meta_schedule::MeasureCandidate, void> const&)
at /data/qshenaf/envs/tvm/src/meta_schedule/feature_extractor/per_store_feature.cc:1414
0: tvm::support::parallel_for_dynamic(int, int, int, std::function<void (int, int)> const&)
at /data/qshenaf/envs/tvm/src/support/parallel_for.cc:139
20: 0xffffffffffffffff
19: clone3
18: start_thread
17: execute_native_thread_routine
at /opt/conda/conda-bld/gcc-compiler_1654084175708/work/gcc/libstdc++-v3/src/c++11/thread.cc:82
16: __pthread_once_slow
15: operator()
at /data/qshenaf/envs/tvm/src/support/parallel_for.cc:113
14: tvm::meta_schedule::PerStoreFeatureNode::ExtractFrom(tvm::meta_schedule::TuneContext const&, tvm::runtime::Array<tvm::meta_schedule::MeasureCandidate, void> const&)::{lambda(int, int)#1}::operator()(int, int) const
at /data/qshenaf/envs/tvm/src/meta_schedule/feature_extractor/per_store_feature.cc:1406
13: tvm::meta_schedule::PerStoreFeatureNode::ExtractSingle(tvm::IRModule, bool, std::vector<std::vector<double, std::allocator<double> >, std::allocator<std::vector<double, std::allocator<double> > > >*)
at /data/qshenaf/envs/tvm/src/meta_schedule/feature_extractor/per_store_feature.cc:1378
12: tvm::tir::PerStoreFeatureCollector::Collect(bool, long, long, tvm::IRModule const&)
at /data/qshenaf/envs/tvm/src/meta_schedule/feature_extractor/per_store_feature.cc:1272
11: tvm::tir::StmtVisitor::VisitStmt_(tvm::tir::BlockRealizeNode const*)
at /data/qshenaf/envs/tvm/src/tir/ir/stmt_functor.cc:150
10: tvm::tir::PerStoreFeatureCollector::VisitStmt_(tvm::tir::BlockNode const*)
at /data/qshenaf/envs/tvm/src/meta_schedule/feature_extractor/per_store_feature.cc:1326
9: tvm::tir::StmtVisitor::VisitStmt_(tvm::tir::BlockNode const*)
at /data/qshenaf/envs/tvm/src/tir/ir/stmt_functor.cc:144
8: tvm::tir::PerStoreFeatureCollector::VisitStmt_(tvm::tir::ForNode const*)
at /data/qshenaf/envs/tvm/src/meta_schedule/feature_extractor/per_store_feature.cc:1301
7: tvm::tir::StmtVisitor::VisitStmt_(tvm::tir::ForNode const*)
at /data/qshenaf/envs/tvm/src/tir/ir/stmt_functor.cc:48
6: tvm::tir::PerStoreFeatureCollector::VisitStmt_(tvm::tir::ForNode const*)
at /data/qshenaf/envs/tvm/src/meta_schedule/feature_extractor/per_store_feature.cc:1301
5: tvm::tir::StmtVisitor::VisitStmt_(tvm::tir::ForNode const*)
at /data/qshenaf/envs/tvm/src/tir/ir/stmt_functor.cc:48
4: tvm::tir::StmtVisitor::VisitStmt_(tvm::tir::BlockRealizeNode const*)
at /data/qshenaf/envs/tvm/src/tir/ir/stmt_functor.cc:150
3: tvm::tir::PerStoreFeatureCollector::VisitStmt_(tvm::tir::BlockNode const*)
at /data/qshenaf/envs/tvm/src/meta_schedule/feature_extractor/per_store_feature.cc:1326
2: tvm::tir::StmtVisitor::VisitStmt_(tvm::tir::BlockNode const*)
at /data/qshenaf/envs/tvm/src/tir/ir/stmt_functor.cc:144
1: tvm::tir::PerStoreFeatureCollector::VisitStmt_(tvm::tir::BufferStoreNode const*)
at /data/qshenaf/envs/tvm/src/meta_schedule/feature_extractor/per_store_feature.cc:1320
0: tvm::tir::group3::Feature::Feature(int, tvm::tir::LoopNest const&, std::vector<long, std::allocator<long> > const&, tvm::tir::group1::Feature::ArithOps const&)
at /data/qshenaf/envs/tvm/src/meta_schedule/feature_extractor/per_store_feature.cc:1111
File "/data/qshenaf/envs/tvm/src/support/parallel_for.cc", line 139
RuntimeError: parallel_for_dynamic error with [13:55:31] /data/qshenaf/envs/tvm/src/meta_schedule/feature_extractor/per_store_feature.cc:1111: Check failed: p < n_loops (2 vs. 2) :
Environment
tvm-0.21.dev0
Steps to reproduce
a = te.placeholder([8], dtype='float64', name='a')
indices = te.placeholder([0, 2, 7], dtype='int32', name='indices')
op_output = topi.take(a, indices, axis=0)
sch = tir.Schedule(te.create_prim_func([a, indices, op_output]).with_attr('target', tvm.target.Target('llvm')))
print(sch.mod)
tvm.build(sch.mod)
database = ms.tir_integration.tune_tir(mod=sch.mod, target='llvm --num-cores=16', work_dir='./tune_tmp', max_trials_global=1, num_trials_per_iter=1)
Triage
- needs-triage
- tune:meta_schedule