Softmax CPU Schedule

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 :slight_smile:

Thank you! I’ll have a look.