diff --git a/openequivariance/openequivariance/benchmark/problems.py b/openequivariance/openequivariance/benchmark/problems.py index b486941c..1dcdb173 100644 --- a/openequivariance/openequivariance/benchmark/problems.py +++ b/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 [ diff --git a/openequivariance/openequivariance/templates/jinja_utils.py b/openequivariance/openequivariance/templates/jinja_utils.py index 609ba3c1..b4ede597 100644 --- a/openequivariance/openequivariance/templates/jinja_utils.py +++ b/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: diff --git a/openequivariance/pyproject.toml b/openequivariance/pyproject.toml index 1ea7c89b..089ebdf4 100644 --- a/openequivariance/pyproject.toml +++ b/openequivariance/pyproject.toml @@ -56,7 +56,6 @@ dev = [ "pytest", "pytest-check", "pytest-subtests", - "mace-torch", "torch_geometric", "cmake" ] diff --git a/tests/batch_test.py b/tests/batch_test.py index e5c0bf7b..77715028 100644 --- a/tests/batch_test.py +++ b/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() ) diff --git a/tests/conv_test.py b/tests/conv_test.py index cd8228cf..1adb5553 100644 --- a/tests/conv_test.py +++ b/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")