(segment|grouped)_matmul implementation using MKL BLAS

[DamianSzwichtenberg]

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

Same benchmark with time in log2 scale:

We can see that segment_matmul execution path starts shining from num_node_types=32 onward.

[rusty1s]

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?

[DamianSzwichtenberg]

Can we get a better review in than mine from @pyg-team/intel-team?

Sure, I'll ping someone.

Can we enable it in our CPP tests to ensure that it is working correct?

Do you mean to enable it in the CI? 😉

[rusty1s]

Do you mean to enable it in the CI? 😉

Yes.

[yanbing-j]

LGTM. Thanks for your hard work!
I have question about the improvement in E2E, do you verify this in a real benchmark?

[codecov-commenter]

Codecov Report

Merging #146 (32db269) into master (2eab973) will decrease coverage by 2.01%.
The diff coverage is 81.81%.

❗ Current head 32db269 differs from pull request most recent head a6048a4. Consider uploading reports for the commit a6048a4 to get more accurate results

@@            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     

Impacted Files	Coverage Δ
pyg_lib/csrc/ops/cpu/matmul_kernel.cpp	83.33% <81.81%> (-16.67%)	⬇️
pyg_lib/csrc/ops/matmul.h	50.00% <0.00%> (-50.00%)	⬇️

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

[DamianSzwichtenberg]

Do you mean to enable it in the CI? 😉

Yes.

@rusty1s Just added. Please take another look. 😉

[DamianSzwichtenberg]

LGTM. Thanks for your hard work! I have question about the improvement in E2E, do you verify this in a real benchmark?

Training and inference benchmarks, unfortunately, do not cover the segment_matmul execution path in HeteroLinear or RGCNConv.

[mingfeima]

Do we have test cases for it?

[DamianSzwichtenberg]

Do we have test cases for it?

Yes, here and here.

[DamianSzwichtenberg]

@mingfeima Please take another look. 😉