Skip to content

Fix AMD wave barrier in HIP kernels; add NequIP-OAM-L test problems - #217

Merged
vbharadwaj-bk merged 5 commits into
mainfrom
amd-wave-barrier-fix
Sep 10, 2026
Merged

Fix AMD wave barrier in HIP kernels; add NequIP-OAM-L test problems#217
vbharadwaj-bk merged 5 commits into
mainfrom
amd-wave-barrier-fix

Conversation

@asglover

@asglover asglover commented Sep 8, 2026

Copy link
Copy Markdown
Collaborator

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's syncwarp used 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!

@asglover

Copy link
Copy Markdown
Collaborator Author

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.

@asglover
asglover marked this pull request as ready for review September 10, 2026 04:33
@asglover asglover self-assigned this Sep 10, 2026
@asglover asglover added bug Something isn't working ci-ready Triggers CI checks for a pull request labels Sep 10, 2026
@vbharadwaj-bk

Copy link
Copy Markdown
Member

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

@asglover asglover added ci-ready Triggers CI checks for a pull request and removed ci-ready Triggers CI checks for a pull request labels Sep 10, 2026
@asglover

Copy link
Copy Markdown
Collaborator Author

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?

@asglover

Copy link
Copy Markdown
Collaborator Author

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 vbharadwaj-bk left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good, test suite almost done on my end, but this should be good to go.

@vbharadwaj-bk vbharadwaj-bk added ci-ready Triggers CI checks for a pull request and removed ci-ready Triggers CI checks for a pull request labels Sep 10, 2026
@vbharadwaj-bk
vbharadwaj-bk merged commit 8b3e58a into main Sep 10, 2026
2 checks passed
@asglover
asglover deleted the amd-wave-barrier-fix branch September 11, 2026 05:10
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working ci-ready Triggers CI checks for a pull request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants