add HIP platform to openmmtools (#753)

* add HIP platform to openmmtools * bump ci * just test openmm 8.1.2 and 8.2.0 * cleanup testing matrix + add missing dep for re-running tests
choderalab · Nov 14, 2024 · cba4ed5 · cba4ed5
1 parent 8352b38
commit cba4ed5
Show file tree

Hide file tree

Showing 6 changed files with 14 additions and 22 deletions.
diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
@@ -29,24 +29,21 @@ jobs:
       fail-fast: false
       matrix:
         python-version: ["3.10", "3.11", "3.12"]
-        openmm: ["8.0", "8.1.1", "8.1.2"]
+        openmm: ["8.1.2", "8.2.0"]
         os: [macOS-latest, ubuntu-latest, macOS-12]
         pymbar-version: ["4"]
         include:
           # Test newest python, openmm, and pymbar we support on windows
-          - openmm: "8.0"
-            python-version: "3.10"
+          - openmm: "8.2"
+            python-version: "3.12"
             os: windows-latest
             pymbar-version: "4"
           # Have one job test pymbar 3 support
-          - openmm: "8.0"
-            python-version: "3.10"
+          - openmm: "8.2"
+            python-version: "3.12"
             os: ubuntu-latest
             pymbar-version: "3"
-        exclude:
-          # no py312 builds of openmm 8.0
-          - openmm: "8.0"
-            python-version: "3.12"
+
 
     steps:
       - uses: actions/checkout@v3

diff --git a/openmmtools/scripts/test_openmm_platforms.py b/openmmtools/scripts/test_openmm_platforms.py
@@ -406,10 +406,7 @@ def main():
 
                 for precision_model in precision_models:
                     # Set precision.
-                    if platform_name == 'CUDA':
-                        platform.setPropertyDefaultValue('CudaPrecision', precision_model)
-                    if platform_name == 'OpenCL':
-                        platform.setPropertyDefaultValue('OpenCLPrecision', precision_model)
+                    platform.setPropertyDefaultValue('Precision', precision_model)
 
                     # Compute potential and force.
                     [platform_potential, platform_force] = compute_potential_and_force(system, positions, platform)
@@ -505,10 +502,7 @@ def main():
 
                         for precision_model in precision_models:
                             # Set precision.
-                            if platform_name == 'CUDA':
-                                platform.setPropertyDefaultValue('CudaPrecision', precision_model)
-                            if platform_name == 'OpenCL':
-                                platform.setPropertyDefaultValue('OpenCLPrecision', precision_model)
+                            platform.setPropertyDefaultValue('Precision', precision_model)
 
                             # Compute potential and force.
                             [platform_potential, platform_force] = compute_potential_and_force_by_force_group(system, positions, platform, force_group)

diff --git a/openmmtools/tests/test_integrators.py b/openmmtools/tests/test_integrators.py
@@ -448,7 +448,7 @@ def create_system(
 
         # Create the context
         platform = openmm.Platform.getPlatformByName(platform_name)
-        if platform_name in ["CPU", "CUDA"]:
+        if platform_name in ["CPU", "CUDA", "HIP"]:
             try:
                 platform.setPropertyDefaultValue("DeterministicForces", "true")
             except Exception as e:

diff --git a/openmmtools/tests/test_utils.py b/openmmtools/tests/test_utils.py
@@ -37,6 +37,7 @@
 def test_platform_supports_precision():
     """Test that platform_supports_precision works correctly."""
 
+    #TODO Add HIP
     for platform_index in range(openmm.Platform.getNumPlatforms()):
         platform = openmm.Platform.getPlatform(platform_index)
         platform_name = platform.getName()
@@ -649,7 +650,7 @@ def test_gentle_equilibration_setup(self):
                 save_box_vectors=False,
             )
 
-    # TODO: Marking as not a test until we solve our GPU CI
+    # TODO: Make a test for HIP as well
     @pytest.mark.slow
     def test_gentle_equilibration_cuda(self):
         """

diff --git a/openmmtools/utils/equilibration.py b/openmmtools/utils/equilibration.py
@@ -104,9 +104,9 @@ def run_gentle_equilibration(topology, positions, system, stages, filename, plat
 
         # Set up context
         platform = openmm.Platform.getPlatformByName(platform_name)
-        if platform_name in ['CUDA', 'OpenCL']:
+        if platform_name in ['CUDA', 'OpenCL', 'HIP']:
             platform.setPropertyDefaultValue('Precision', 'mixed')
-        if platform_name in ['CUDA']:
+        if platform_name in ['CUDA', 'HIP']:
             platform.setPropertyDefaultValue('DeterministicForces', 'true')
 
         context = openmm.Context(system_copy, integrator, platform)

diff --git a/openmmtools/utils/utils.py b/openmmtools/utils/utils.py
@@ -544,7 +544,7 @@ def platform_supports_precision(platform, precision):
     if platform.getName() == 'CPU':
         return precision in ['mixed']
 
-    if platform.getName() in ['CUDA', 'OpenCL']:
+    if platform.getName() in ['CUDA', 'OpenCL', 'HIP']:
         properties = { 'Precision' : precision }
         system = openmm.System()
         system.addParticle(1.0) # Cannot create Context on a system with no particles