Fix AMD wave barrier in HIP kernels; add NequIP-OAM-L test problems - #217
Conversation
|
I used AMD Dev Cloud to reproduce the bug and verify the fix. So I think the only opens should be lints / style corrections. I'm ready for a review. |
|
Thanks, I'm running the test suite on my end as well; so, gonna request a couple of things, although I know this is a small fix: can you see if this patch matches CONTRIBUTING.md and correct anything as needed? In particular, would prefer the PR description be handwritten, at least (it can be empty, that's fine). |
|
No worries, happy to maintain the standard of human written comments. Thanks for the review. I Please let me know if you'd like to see any additional changes? |
|
This PR also now drops the mace-torch dev dependency that pulls in the old e3nn version. In a separate PR, I'll full remove the dependency. |
vbharadwaj-bk
left a comment
There was a problem hiding this comment.
Looks good, test suite almost done on my end, but this should be good to go.
Problem
Initial report was in issue #211. Thank you @iansteeg for the report!
Root cause was found by @kavanase here: mir-group/nequip#610
Huge thanks to you for the work in root causing this!
Fix
jinja_utils.py: on HIP, CUDA'ssyncwarpused to be replaced with AMD's threadfence block. Now it's replaced with the expanded version of AMD's__syncwarp(). The reason we didn't just use__syncwarp()is because it was added somewhat recently (ROCM 7.0) and we wanted it to work for older ROCMs too.Test
The failing tensor product was not in the test suite. This PR adds:
A nequip TPP builder (it's a different instruction / irrep order than mace, this interacts with the greedy memory planning algorithm)
The TPP's from the model that had the issue reported.
These TPPs were added to the batch and conv production model lists.
Also, somewhat unrelated, this removes the mace-torch test dependency which pulls e3nn <4.4
Status
Reproduced bug and fix verified on AMD MI-300. Thank you AMD for the Dev Cloud Access!