It seems like Softmax has a CPU schedule defined in Python at topi/python/topi/x86/nn.py. However, there is no corresponding schedule in C++. topi.cc just seems to call topi::x86::default_schedule. I searched, and there is no implementation of schedule_softmax for x86 in C++. Is this expected? If not, I’d be happy to add the C++ schedule.
Yes. That part is missing. You are more than welcome to contribute!
@kevinthesun, I have modified Softmax’s compute and schedule on CPU in a way that I believe is more performant. What are the proper steps to send it out for review? Does it need an RFC?
You can just send out a PR. RFC is not necessary for this.
@haichen or @kevinthesun would you be able to provide some guidance here?
I submitted a PR to update compute and schedule here: https://github.com/dmlc/tvm/pull/3680
However, log_softmax is now broken since softmax and log_softmax were written to use the same schedule. I am thinking the best fix is to introduce a new schedule just for log_softmax. What do you think?
@jonso You can take a look at related python changes here: https://github.com/kevinthesun/tvm/commit/52d57e19520f80c4915892b2e65546806a5f184f
@kevinthesun Do you think that creating a separate schedule for log_softmax is the best way to go here? With my change, log_softmax and softmax now have different inputs to the normalize stage, so you can’t use the same schedule.
I just tested the schedule in your fork, and my change has more perf improvements.
You can still use the same schedule if the only difference is _compute_exp stage is missing in log_softmax. You can use branch to deal with different input op tag.
Thanks a lot. I’ve updated my PR
Thank you! I’ll have a look.