[Bug] schedule tuning for sparse_dense with sparse_lhs=True crashes

[usamec]

Thanks for participating in the TVM community! We use https://discuss.tvm.ai for any general usage questions and discussions. The issue tracker is used for actionable items such as feature proposals discussion, roadmaps, and bug tracking. You are always welcomed to post on the forum first 😸

Issues that are inactive for a period of time may get closed. We adopt this policy so that we won't lose track of actionable issues that may fall at the bottom of the pile. Feel free to reopen a new one if you feel there is an additional problem that needs attention when an old one gets closed.

Expected behavior

Sparse matmul tuning using https://tvm.apache.org/docs/how_to/tune_with_autoscheduler/tune_sparse_x86.html#sphx-glr-how-to-tune-with-autoscheduler-tune-sparse-x86-py and sparse_lhs=True and N=M=K=256 would work.

Actual behavior

--> 498 _ffi_api.AutoSchedule(search_policy, tuning_options)

File tvm/_ffi/_cython/./packed_func.pxi:331, in tvm._ffi._cy3.core.PackedFuncBase.__call__()

File tvm/_ffi/_cython/./packed_func.pxi:262, in tvm._ffi._cy3.core.FuncCall()

File tvm/_ffi/_cython/./packed_func.pxi:251, in tvm._ffi._cy3.core.FuncCall3()

File tvm/_ffi/_cython/./base.pxi:181, in tvm._ffi._cy3.core.CHECK_CALL()

TVMError: Traceback (most recent call last):
  7: TVMFuncCall
  6: _ZN3tvm7runtime13PackedFuncObj9ExtractorINS0_16PackedFuncSubObjIZNS0_15TypedPackedFuncIFNS0_5ArrayINS0_9ObjectRefEvEENS_14auto_scheduler12SearchPolicyENS8_13TuningOptionsEEE17AssignTypedLambdaINS8_L9__mk_TVM3MUlS9_SA_E_EEEvT_SsEUlRKNS0_7TVMArgsEPNS0_11TVMRetValueEE_EEE4CallEPKS1_SG_SK_
  5: tvm::auto_scheduler::AutoSchedule(tvm::auto_scheduler::SearchPolicy, tvm::auto_scheduler::TuningOptions)
  4: tvm::auto_scheduler::SketchPolicyNode::Search(int, int, int, tvm::auto_scheduler::ProgramMeasurer)
  3: tvm::auto_scheduler::SketchPolicyNode::SearchOneRound(int, tvm::runtime::Array<tvm::auto_scheduler::State, void>*)
  2: tvm::auto_scheduler::SketchPolicyNode::SampleInitPopulation(tvm::runtime::Array<tvm::auto_scheduler::State, void> const&)
  1: tvm::support::parallel_for(int, int, std::function<void (int)> const&, int, std::function<std::vector<std::vector<int, std::allocator<int> >, std::allocator<std::vector<int, std::allocator<int> > > > (int, int, int, int)>) [clone .cold]
  0: _ZN3tvm7runtime6deta
  10: 0xffffffffffffffff
  9: __clone
  8: start_thread
        at /build/glibc-sMfBJT/glibc-2.31/nptl/pthread_create.c:477
  7: execute_native_thread_routine
  6: std::thread::_State_impl<std::thread::_Invoker<std::tuple<std::packaged_task<void (std::vector<int, std::allocator<int> > const&, std::function<void (int)> const&)>, std::vector<int, std::allocator<int> >, std::function<void (int)> > > >::_M_run()
  5: __pthread_once_slow
        at /build/glibc-sMfBJT/glibc-2.31/nptl/pthread_once.c:116
  4: _ZNSt13__future_base13_State_baseV29_M_do_setEPSt8functionIFSt10unique_ptrINS_12_Result_baseEN
  3: _ZNSt17_Function_handlerIFSt10unique_ptrINSt13__future_base12_Result_baseENS2_8_DeleterEEvENS1_12_Task_setterIS0_INS1_7_ResultIvEES3_EZNS
  2: std::_Function_handler<void (int), tvm::auto_scheduler::SketchPolicyNode::SampleInitPopulation(tvm::runtime::Array<tvm::auto_scheduler::State, void> const&)::{lambda(int)#1}>::_M_invoke(std::_Any_data const&, int&&)
  1: tvm::auto_scheduler::InitFillTileSize::Apply(tvm::auto_scheduler::SketchPolicyNode*, tvm::auto_scheduler::State*, std::mersenne_twister_engine<unsigned long, 32ul, 624ul, 397ul, 31ul, 2567483615ul, 11ul, 4294967295ul, 7ul, 2636928640ul, 15ul, 4022730752ul, 18ul, 1812433253ul>*) const
  0: _ZN3tvm7runtime6deta
  File "/workspace/tvm/src/support/parallel_for.cc", line 92
TVMError: Parallel_for error with [18:20:06] /workspace/tvm/src/auto_scheduler/search_policy/sketch_policy_rules.cc:515: 
---------------------------------------------------------------
An error occurred during the execution of TVM.
For more information, please see: https://tvm.apache.org/docs/errors.html
---------------------------------------------------------------
  Check failed: (ps->extent) is false: 

Environment

pip freeze | grep tvm
apache-tvm==0.9.0

Running Linux (ubuntu 20.04) on Cascade Lake arch.

Steps to reproduce

Take https://tvm.apache.org/docs/how_to/tune_with_autoscheduler/tune_sparse_x86.html#sphx-glr-how-to-tune-with-autoscheduler-tune-sparse-x86-py and change sparse_lhs=True and N=M=K=256