ptrendx opened a new pull request #15545: Softmax fwd optimization for GPU URL: https://github.com/apache/incubator-mxnet/pull/15545 ## Description ## This PR optimizes Softmax implementation (fwd only for now, I will add bwd in the future PR) for cases where stride is 1 and the leading dimension is small (up to 20kB of data in that dimension). There are 2 optimizations in this kernel compared to the previous one: - using of the longer datatypes for loading/writing data (so using up to 8B per read/write instead of e.g. 2B in case of fp16 I/O) - using persistent storage to reduce the number of memory accesses (previous implementation used 3 reads and 1 write, this implementation uses 1 read and 1 write). Compared to the previous implementation on fp16 I/O the new kernel is up to 4x faster. @eric-haibin-lin ## Checklist ## ### Essentials ### Please feel free to remove inapplicable items for your PR. - [x] Changes are complete (i.e. I finished coding on this PR) - [x] All changes have test coverage: - [x] To the my best knowledge, examples are either not affected by this change, or have been fixed to be compatible with this change
---------------------------------------------------------------- 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. For queries about this service, please contact Infrastructure at: [email protected] With regards, Apache Git Services
