-
Notifications
You must be signed in to change notification settings - Fork 40
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
(segment|grouped)_matmul
implementation using MKL BLAS
#146
(segment|grouped)_matmul
implementation using MKL BLAS
#146
Conversation
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think this looks great. Thanks for the effort!
Two questions:
- Can we get a better review in than mine from @pyg-team/intel-team?
- Can we enable it in our CPP tests to ensure that it is working correct?
Sure, I'll ping someone.
Do you mean to enable it in the CI? 😉 |
Yes. |
LGTM. Thanks for your hard work! |
64a3b43
to
2e38ea5
Compare
Codecov Report
@@ Coverage Diff @@
## master #146 +/- ##
==========================================
- Coverage 94.01% 92.00% -2.02%
==========================================
Files 20 20
Lines 535 588 +53
==========================================
+ Hits 503 541 +38
- Misses 32 47 +15
📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more |
@rusty1s Just added. Please take another look. 😉 |
Training and inference benchmarks, unfortunately, do not cover the |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Do we have test cases for it?
@mingfeima Please take another look. 😉 |
32db269
to
13c6526
Compare
for more information, see https://pre-commit.ci
13c6526
to
6b76eb8
Compare
6b76eb8
to
a6048a4
Compare
By default PyTorch is built with MKL BLAS support, we can take advantage of that and use gemm_batch to implement
(segment|grouped)_matmul
.segment_matmul
C++ benchmarks (24 core machine) showed that the new version is on average 3.5 times faster than the previous one (with maximum speedup up to 6 times - note that chosen shapes are quite small, so speedup will be more impressive if a problem becomes more complex).Additionally, HeteroLinear benchmarks were performed (thanks to @puririshi98 for sharing the benchmarks code):
![segment_matmul_bench](https://user-images.githubusercontent.com/42009345/202194324-abdf8802-3207-4c0f-9c1e-3d84dc7d2e43.PNG)
![segment_linear_bench_log2](https://user-images.githubusercontent.com/42009345/202194625-8b0f1b7b-fae8-48fd-bcda-37492195c1ca.PNG)
Same benchmark with time in log2 scale:
We can see that
segment_matmul
execution path starts shining fromnum_node_types=32
onward.