Request: AMD HIP Platform Support

[HiteSit]

Hi all,

I was wondering if is possible to add the support for the AMD HIP platform.
Specifically OpenMM already is able to work on HIP platforms installing a specific hack with Conda:

mamba install jaimergp/label/unsupported-cudatoolkit-shim::cudatoolkit=11.2.2 && mamba install streamhpc::openmm-hip=8.0.0

But the problem is that in your code, specifically under openmmutils/utils.py there is an assert that allows only ["OpenCL", "CUDA"], probably the fix is easy since as said OpenMM is able to see HIP as platform as well as give it the "right" speed.

The fix would have a huge impact since the LUMI HPC (most powerful in europe) only supports HIP platform.

[mikemhenry]

@HiteSit Thank you for raising this issue!

"We" (not the OpenFE team but other orgs I am apart of) are working to get ROCm/HIP onto conda-forge so that no hacks will be needed to install openmm.

Can you link the line of code that has this assert? Also do you know how OpenMM reports the HIP platform string-wise? I would be happy to get this working.

[HiteSit]

So,
the utils file is under:
/mambaforge/envs/cheminf_3_11/lib/python3.11/site-packages/openmmtools/utils/utils.py

def platform_supports_precision(platform, precision):
    """Determine whether the specified OpenMM Platform supports the specified minimum precision.

    Parameters
    ----------
    platform : str or openmm.Platform
        The platform or platform name to check
    precision : str
        One of ['single', 'mixed', 'double']

    Returns
    -------
    is_supported : bool
        True if the platform supports the specified precision; False otherwise
    """
    SUPPORTED_PRECISIONS = ['single', 'mixed', 'double']
    assert precision in SUPPORTED_PRECISIONS, f"Precision {precision} must be one of {SUPPORTED_PRECISIONS}"

    if isinstance(platform, str):
        # Get the actual Platform object if the platform_name was specified
        platform = openmm.Platform.getPlatformByName(platform)

    if platform.getName() == 'Reference':
        # Reference is double precision
        return (precision == 'double')

    if platform.getName() == 'CPU':
        return precision in ['mixed']

    if platform.getName() in ['CUDA', 'OpenCL']:
        properties = { 'Precision' : precision }
        system = openmm.System()
        system.addParticle(1.0) # Cannot create Context on a system with no particles
        integrator = openmm.VerletIntegrator(0.001)
        try:
            context = openmm.Context(system, integrator, platform, properties)
            del context, integrator
            return True
        except Exception as e:
            return False

    raise Exception(f"Platform {platform.getName()} unknown")

def get_available_platforms(minimum_precision='mixed'):
    """Return a list of the available OpenMM Platforms that can satisfy the requested minimum precision.

    Parameters
    ----------
    minimum_precision : str, optional, default='mixed'
        One of [None, 'single', 'mixed', 'double']
        If None, all available platforms will be returned.

    Returns
    -------
    platforms : list of openmm.Platform
        Platforms that support specified minimumprecision
    """
    platforms = [openmm.Platform.getPlatform(i) for i in range(openmm.Platform.getNumPlatforms())]

    if minimum_precision is not None:
        # Filter based on precision support
        platforms = [ platform for platform in platforms if platform_supports_precision(platform, minimum_precision) ]

    return platforms

def get_fastest_platform(minimum_precision='mixed'):
    """Return the fastest available platform.

    This relies on the hardcoded speed values in Platform.getSpeed().

    Parameters
    ----------
    minimum_precision : str, optional, default='mixed'
        One of ['single', 'mixed', 'double']

    Returns
    -------
    platform : openmm.Platform
       The fastest available platform.

    """
    platforms = get_available_platforms(minimum_precision=minimum_precision)
    fastest_platform = max(platforms, key=lambda x: x.getSpeed())
    return fastest_platform

If I run:

platforms = [openmm.Platform.getPlatform(i) for i in range(openmm.Platform.getNumPlatforms())]
for platform in platforms:
     name = platform.getName()
     print(name)

>>> Something like "CPU", "OpenCL", "HIP"

Wrong timing, this week LUMI HPC is down for maintenance, I will edit this message with the output from the print. But I'm quite sure that the only problem is that assert. I would like to edit the code by myself and report if just adding to the assert also ["HIP"] would work, but LUMI install Conda (OpenFF, OpenMM, OpenFreeEnergy and so on) using Singularity enviroments (read-only) and I could not figure it out yet how to use it with --sandbox, If I figure it out I will edit this post.

[IAlibay]

Support for the HIP platform would be nice, however I would warn that it definitely needs validation prior to use. Untested platforms tend to be prone to odd behaviour in the alchemical world. Validation would require at least an HFE validation test & a couple of RBFE test cases.

Is this something you'd be willing to take on @HiteSit ?

[HiteSit]

@IAlibay
Surely I would like to contribute.
First I have to resolve the problem with the Singularity Enviroment. Beside that I can code quite well but I do not have a enough experience with alchemical transformation so I need to be guided.

[IAlibay]

@HiteSit - we need to discuss this internally first, but the requirement here would be mostly to run a suitably large set of alchemical simulations to verify that the results are reasonable. This would mostly require access to suitable AMD HIP compute resources to do such a validation, which unfortunately we do not have :(

[HiteSit]

@IAlibay
Yes I understood you need access to the platform. I will try to grab some computational time (node/hours) for free, but anyway if you do not need an astonishing amount of computational time (maybe try to give me more or less a range of node/hours) I'm willing to share my computational time without problem. It's my pleasure to contribute.

You can contact me on riccardo.fusco@upol.cz

[IAlibay]

@HiteSit - I'm re-opening this issue if it's ok, there definitely needs to be some kind of update to our compute platform selection to allow for HIP. My question was more of a "once this is one, someone will need to check it works".

[HiteSit]

@IAlibay
I can check, the only constraint is the computational time, I have a limited amount of computational time, but as rule of thumb if the testing is around let's say 10 proteins each of them with 30 ligands should not be a problem. If it's more I can figure it out a way to grab more computational time.

[mikemhenry]

Thanks for pointing to the code file! I will raise this as a separate issue on the openmmtools side of things since there isn't really any reason why we couldn't add support for HIP there, but as @IAlibay said, when it comes to using it in an openfe workflow, we will need it validated.

[mikemhenry]

@HiteSit I've made a branch on openmmtools with the changes that I think are needed. To play around with it, to test it you will need to run:

(you can also use mamba or conda to do this)

# Create an env that has openmm with the hip platform
$ micromamba create -n openmm82beta-openfe -c conda-forge/label/openmm_rc -c conda-forge "openmm-hip==8.2.0beta" "openfe==1.1.0"
# Activate env
$ micromamba activate openmm82beta-openfe
# Install openmmtools branch 
$ pip install git+https://github.com/choderalab/openmmtools.git@feat/add-hip-platform

Can you let me know if that works for you? I don't have an AMD card handy but I can spin up some cloud resources to test if needed.