tsu-bin opened a new pull request, #15431:
URL: https://github.com/apache/tvm/pull/15431
To reproduce the issue, just construct a simple softmax op and run
auto-scheduler, below is the code snippet
```
@auto_scheduler.register_workload
def create_softmax_op(*shape):
logits = te.placeholder(shape, name="logits")
softmax = topi.nn.softmax(logits)
return [logits, softmax]
input_shape = (20, 10)
task = auto_scheduler.SearchTask(create_softmax_op, input_shape,
target='llvm -mcpu=skylake-avx512')
task.tune(auto_scheduler.TuningOptions(
num_measure_trials=20,
measure_callbacks=[auto_scheduler.RecordToFile('./softmax.json')],
verbose=2
))
```
And during the auto-scheduling process, a lot error logs will show up, below
is some screen shot
```
[20:40:35]
/hostShare/tvm_all/tvm_latest/src/auto_scheduler/compute_dag.cc:1377: Warning:
InferBound fails on the state:
Placeholder: logits
parallel i0 (0,20)
T_softmax_expsum auto_unroll: 64
for i0 (None)
for k (None)
T_softmax_expsum = ...
T_softmax_maxelem auto_unroll: 64
for i0 (None)
for k (None)
T_softmax_maxelem = ...
for i1 (0,10)
for i0 (None)
vectorize i1 (None)
T_softmax_exp = ...
T_softmax_norm = ...
with: [20:40:35] /hostShare/tvm_all/tvm_latest/src/te/schedule/bound.cc:175:
InternalError: Check failed: (found_attach || stage_attach.size() == 0) is
false: Invalid Schedule, cannot find the producer compute(T_softmax_exp,
body=[T.exp(logits[i0, i1] - T_softmax_maxelem[i0])], axis=[T.iter_var(i0,
T.Range(0, 20), "DataPar", ""), T.iter_var(i1, T.Range(0, 10), "DataPar", "")],
reduce_axis=[], tag=softmax_output, attrs={}) along the loop nest specified by
compute_at of consumer compute(T_softmax_expsum,
body=[T.reduce(T.comm_reducer(lambda x, y: x + y, [T.float32(0)]),
source=[T_softmax_exp[i0, k]], init=[], axis=[T.iter_var(k, T.Range(0, 10),
"CommReduce", "")], condition=T.bool(True), value_index=0)],
axis=[T.iter_var(i0, T.Range(0, 20), "DataPar", "")],
reduce_axis=[T.iter_var(k, T.Range(0, 10), "CommReduce", "")],
tag=softmax_output, attrs={})
Stack trace:
0: tvm::te::InferRootBound(tvm::te::Stage const&, tvm::te::GraphContext
const&, std::unordered_map<tvm::tir::IterVar, tvm::Range,
std::hash<tvm::tir::IterVar>, std::equal_to<tvm::tir::IterVar>,
std::allocator<std::pair<tvm::tir::IterVar const, tvm::Range> > >*)
at /hostShare/tvm_all/tvm_latest/src/te/schedule/bound.cc:175
1: tvm::te::InferBound(tvm::te::Schedule const&)
at /hostShare/tvm_all/tvm_latest/src/te/schedule/bound.cc:233
2: tvm::auto_scheduler::ComputeDAG::InferBound(tvm::auto_scheduler::State
const&) const
at
/hostShare/tvm_all/tvm_latest/src/auto_scheduler/compute_dag.cc:1336
3:
tvm::auto_scheduler::ComputeDAG::InferBound(tvm::runtime::Array<tvm::auto_scheduler::State,
void> const&) const::$_8::operator()(int) const
at
/hostShare/tvm_all/tvm_latest/src/auto_scheduler/compute_dag.cc:1375
4: 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)>)::$_0::operator()(std::vector<int, std::allocator<int> > const&,
std::function<void (int)> const&) const
at /hostShare/tvm_all/tvm_latest/src/support/parallel_for.cc:72
5: __pthread_once_slow
6: __gthread_once(int*, void (*)())
at
/usr/bin/../lib/gcc/x86_64-linux-gnu/9/../../../../include/x86_64-linux-gnu/c++/9/bits/gthr-default.h:700
7: execute_native_thread_routine
at
/opt/conda/conda-bld/gcc-compiler_1654084175708/work/gcc/libstdc++-v3/src/c++11/thread.cc:82
8: start_thread
9: __clone
10: 0xffffffffffffffff
```
When I use LLDB to attach to the python ansor tuning process, I found the
root cause is:
The schedule attaches the stage 'T_softmax_exp' to the leaf IterVar 'i' of
stage 'T_softmax_norm', but the stage 'T_softmax_exp' has two consumer stages,
one is 'T_softmax_norm' and the other is 'T_softmax_expsum'. Inside
'InferRootBound', the OPs of these two consumer stages are stored into an
unordered_set `std::unordered_set<Operation> consumers`, so there is a chance
that 'T_softmax_expsum' is processed before 'T_softmax_norm', so the
`Array<IterVar> stage_attach` is not empty (there are two elements in it) but
the IterVar of 'T_softmax_expsum' can not match the `stage_attach` and
`found_attach` is false, then the `ICHECH` raises an exception just as the
error log shows above, and finally the exception is caught by
`SketchPolicyNode::EvolutionarySearch`.
The consequence is `SketchPolicyNode::EvolutionarySearch` missed some
potential good seeds, and the final schedule may not be the best one.
The change is rather simple, just check if the current OP's stage is the
attached stage.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
