Repository navigation
Move bmm calls to use PyTorch bmm_out via the stable op dispatcher API - #220
Conversation
vbharadwaj-bk
left a comment
vbharadwaj-bk
left a comment
There was a problem hiding this comment.
Code looks fine, just one or two questions. Perf-wise, do you know how this compares to the original CUDA version?
|
And this is a great change, thanks - good to get the dependence on BLAS eliminated. |
64af319 to
4b89132
Compare
…est it through symmetric tensor product. The test that we don't rely on blas is over testing.
x[permutation] keeps the stride order of non-contiguous node features, so x.view fails when there is more than one channel. reshape is still a view for contiguous inputs, and the non-contiguous tensor still reaches group_gemm.
asglover
changed the title
It introduces more CPU overhead. On expected workloads, it's not measurable in overall model execution time. For really small models, with ~30 atoms, it can be ~3% slower. With only ~8 atoms, the worse case was ~5% slower. But this is somewhat expected. I think because the performance difference is negligible during large workloads it should be acceptable? It also enables SYCL support for symmetric contractions and they use MACE. And so they will be very happy about this. Please weigh in on this. If the performance impact is unacceptable, then I'm blocked for a little bit. The ways out would be:
|
|
Ok, that overhead seems minimal, continuing with this strategy should be fine. |
69e4728 to
fe75c96
Compare
|
Okay, all the tests on cuda and on sycl have come back green. I'm going to go over this in detail, probably polish a few things up. I'll be out of town traveling for two days, will pick it up on Wednesday. |
|
Okay, I ran this on CUDA, HIP and SYCL and it's looking good. I'm ready for a review @vbharadwaj-bk. Thank you for your patience on this one. |
This PR removes our dependency on BLAS. Instead we move to calling PyTorch's bmm_out through the dispatcher. It makes it so that we can support SYCL for symmetric contractions.
It has a slight performance regression when the amount of work done per BMM is very small, however it is not significant during our benchmark / expected workloads.
This PR contributes to enabling multi device support