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

Author: asglover

openequivariance/openequivariance/benchmark/problems.py

@@ -196,6 +196,56 @@ def __init__(
         )
 
 
+class NequIPTPP(TPProblem):
+    """
+    Taken from nequip.nn.interaction_block.InteractionBlock:
+    https://github.com/mir-group/nequip/blob/27d9d2182da918ab7be0017d8300e53278f5e00e/nequip/nn/interaction_block.py#L89-L116
+    """
+
+    def __init__(
+        self,
+        feature_irreps_in: Irreps,
+        irreps_edge_attr: Irreps,
+        feature_irreps_out: Irreps,
+        label: Optional[str] = None,
+        irrep_dtype=np.float32,
+        weight_dtype=np.float32,
+    ):
+        feature_irreps_in = Irreps(feature_irreps_in)
+        irreps_edge_attr = Irreps(irreps_edge_attr)
+        feature_irreps_out = Irreps(feature_irreps_out)
+
+        irreps_mid = []
+        instructions = []
+        for i, (mul, ir_in) in enumerate(feature_irreps_in):
+            for j, (_, ir_edge) in enumerate(irreps_edge_attr):
+                for ir_out in ir_in * ir_edge:
+                    if ir_out in feature_irreps_out:
+                        k = len(irreps_mid)
+                        irreps_mid.append((mul, ir_out))
+                        instructions.append((i, j, k, "uvu", True))
+
+        irreps_mid = Irreps(irreps_mid)
+        irreps_mid, p, _ = irreps_mid.sort()
+
+        instructions = [
+            (i_in1, i_in2, p[i_out], mode, train)
+            for i_in1, i_in2, i_out, mode, train in instructions
+        ]
+
+        super().__init__(
+            feature_irreps_in,
+            irreps_edge_attr,
+            irreps_mid,
+            instructions,
+            internal_weights=False,
+            shared_weights=False,
+            label=label,
+            irrep_dtype=irrep_dtype,
+            weight_dtype=weight_dtype,
+        )
+
+
 FCTPP = FullyConnectedTPProblem
 CTPP = ChannelwiseTPP
 
@@ -347,6 +397,17 @@ def nequip_problems():
     ]
 
 
+# https://github.com/mir-group/nequip/blob/27d9d2182da918ab7be0017d8300e53278f5e00e/nequip/model/nequip_models.py#L30-L58
+def nequip_oam_problems():
+    sh = "1x0e+1x1o+1x2e+1x3o"
+    hidden = "128x0e+64x1o+32x2e+32x3o"
+    return [
+        NequIPTPP("32x0e", sh, hidden, "nequip-oam-l-first-layer"),
+        NequIPTPP(hidden, sh, hidden, "nequip-oam-l-main-layers"),
+        NequIPTPP(hidden, sh, "128x0e", "nequip-oam-l-last-layer"),
+    ]
+
+
 # https://github.com/atomicarchitects/nequix/blob/main/configs/nequix-mp-1.yml
 def nequix_problems():
     return [

openequivariance/openequivariance/templates/jinja_utils.py

@@ -26,7 +26,11 @@ def get_jinja_environment(is_hip=False):
     env.globals["enumerate"] = enumerate
 
     env.globals["is_hip"] = is_hip
-    env.globals["syncwarp"] = "__threadfence_block()" if is_hip else "__syncwarp()"
+    env.globals["syncwarp"] = (
+        '__builtin_amdgcn_fence(__ATOMIC_RELEASE, "wavefront");__builtin_amdgcn_wave_barrier();__builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "wavefront");'
+        if is_hip
+        else "__syncwarp()"
+    )
     env.globals["atomic_add"] = "unsafeAtomicAdd" if is_hip else "atomicAdd"
 
     if is_hip:

openequivariance/pyproject.toml

@@ -56,7 +56,6 @@ dev = [
   "pytest",
   "pytest-check",
   "pytest-subtests",
-  "mace-torch",
   "torch_geometric",
   "cmake"
 ]

tests/batch_test.py

@@ -14,6 +14,7 @@
     diffdock_problems,
     e3nn_torch_tetris_poly_problems,
     mace_problems,
+    nequip_oam_problems,
     nequip_problems,
 )
 from pytest_check import check
@@ -131,6 +132,7 @@ class TestProductionModels(TPCorrectness):
     production_model_tpps = (
         mace_problems()
         + nequip_problems()
+        + nequip_oam_problems()
         + e3nn_torch_tetris_poly_problems()
         + diffdock_problems()
     )

tests/conv_test.py

@@ -19,6 +19,7 @@
     mace_problems,
     diffdock_problems,
     e3tools_problems,
+    nequip_oam_problems,
 )
 
 
@@ -172,7 +173,10 @@ def test_tp_triple_bwd(self, conv_object, graph, with_jax):
 
 class TestProductionModels(ConvCorrectness):
     production_model_tpps = (
-        mace_problems() + diffdock_problems() + [e3tools_problems()[0]]
+        mace_problems()
+        + diffdock_problems()
+        + [e3tools_problems()[0]]
+        + nequip_oam_problems()
     )
 
     @pytest.fixture(params=production_model_tpps, ids=lambda x: x.label, scope="class")