diff --git a/src/openfe/protocols/openmm_utils/omm_compute.py b/src/openfe/protocols/openmm_utils/omm_compute.py index 1779af2bf..a09420040 100644 --- a/src/openfe/protocols/openmm_utils/omm_compute.py +++ b/src/openfe/protocols/openmm_utils/omm_compute.py @@ -52,6 +52,7 @@ def get_openmm_platform( "cpu": "CPU", "opencl": "OpenCL", "cuda": "CUDA", + "hip": "HIP", }[str(platform_name).lower()] except KeyError: pass @@ -61,18 +62,18 @@ def get_openmm_platform( platform = Platform.getPlatformByName(platform_name) # Set precision and properties name = platform.getName() - if name in ["CUDA", "OpenCL"]: + if name in ["CUDA", "OpenCL", "HIP"]: platform.setPropertyDefaultValue("Precision", "mixed") if gpu_device_index is not None: index_list = ",".join(str(i) for i in gpu_device_index) platform.setPropertyDefaultValue("DeviceIndex", index_list) - if name == "CUDA": + if name in ["CUDA", "HIP"]: platform.setPropertyDefaultValue("DeterministicForces", "true") - if name != "CUDA": + if name not in ["CUDA", "HIP"]: wmsg = ( - f"Non-CUDA platform selected: {name}, this may significantly " + f"Either a non GPU platform or OpenCL selected: {name}, this may significantly " "impact simulation performance" ) warnings.warn(wmsg) diff --git a/src/openfe/protocols/openmm_utils/omm_settings.py b/src/openfe/protocols/openmm_utils/omm_settings.py index 97c7ad8ce..1fce1a612 100644 --- a/src/openfe/protocols/openmm_utils/omm_settings.py +++ b/src/openfe/protocols/openmm_utils/omm_settings.py @@ -329,9 +329,12 @@ class OpenMMEngineSettings(SettingsBaseModel): """ OpenMM compute platform to perform MD integration with. If ``None``, will choose fastest available platform. - Allowed platforms are; ``cuda``, ``opencl``, ``cpu``. + Allowed platforms are; ``cuda``, ``opencl``, ``cpu``, and ``hip``. Default ``cuda``. + Notes + ----- + To use the ``hip`` platform, you will need to install ``openmm-hip``. """ gpu_device_index: Optional[list[int]] = None """ @@ -348,7 +351,7 @@ class OpenMMEngineSettings(SettingsBaseModel): @field_validator("compute_platform") def supported_sampler(cls, v): - supported = ["cpu", "opencl", "cuda"] + supported = ["cpu", "opencl", "cuda", "hip"] if v is not None and v.lower() not in supported: errmsg = f"Only the following OpenMM compute backends are supported: {supported}" raise ValueError(errmsg)