[TOPI] Update softmax compute and CPU schedule #3680
Merged
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
@kevinthesun @vinx13 can you please review and add any other reviewers you think are necessary? I am currently modifying log_softmax, and it seems worthwhile to create a new generic schedule for it, since the inputs to tvm.compute are now different for softmax and log_softmax. What do you think? |
|
Thank you @soiferj , can you check the CI problem? |
|
Yeah, I'm taking a look at the CI failure now. It seems to be an issue in the CUDA schedule. I will work on it. |
|
The CI issue is fixed. |
|
@kevinthesun feel free to merge the PR given you are managing it |
|
Thank you for contributing! |
wweic
pushed a commit
to wweic/tvm
that referenced
this pull request
Aug 9, 2019
* Update Softmax compute and CPU schedule * Add C++ compute * Fix schedule * Update CUDA and OpenGL schedules * Fix log_softmax * Fix hls and opengl schedules * Fix CUDA schedule
wweic
pushed a commit
to neo-ai/tvm
that referenced
this pull request
Sep 6, 2019
* Update Softmax compute and CPU schedule * Add C++ compute * Fix schedule * Update CUDA and OpenGL schedules * Fix log_softmax * Fix hls and opengl schedules * Fix CUDA schedule
|
Another suggestion - https://discuss.tvm.ai/t/softmax-sequence-of-relay-ops/5686 |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This change improves performance for softmax by simplifying the computation and writing a schedule that supports better parallelization.
Compute: Currently,
exp(input - max)is computed twice: once in the_compute_expsumstage and once in the_normalizestage. This change adds an extra stage to compute this tensor once. It is then re-used in the_compute_expsumand_normalizestages.Schedule: Currently, the schedule only parallelizes the
_normalizestage of the computation. This change puts all stages of computation under a common root and parallelizes the outer dimensions.The following results are with a tensor of shape
(1,12,128,128)andaxis=-1. This simulates the softmax in BERT base. The CPU is Intel Xeon E5-2650, and the Relay target string isllvm -mcpu=core-avx2.