Add multiplatform variants

CMakeLists.txt

@@ -15,6 +15,10 @@ include("${PROJECT_MODULE_PATH}/FetchDLPack.cmake")
 
 # Create the library
 add_library(${PROJECT_NAME} SHARED "")
+# Also create a CUDA variant when possible
+if(BUILD_CUDA_LIB)
+    add_library(${PROJECT_NAME}CUDA SHARED "")
+endif()
 
 add_subdirectory(openmmapi)
 add_subdirectory(platforms)
@@ -23,6 +27,13 @@ target_compile_features(${PROJECT_NAME} PRIVATE cxx_std_11)
 target_include_directories(${PROJECT_NAME} PUBLIC include)
 target_link_libraries(${PROJECT_NAME} PUBLIC OpenMM::OpenMM dlpack::dlpack)
 
+if(BUILD_CUDA_LIB)
+    target_compile_definitions(${PROJECT_NAME}CUDA PUBLIC OPENMM_BUILD_CUDA_LIB)
+    target_link_libraries(${PROJECT_NAME}CUDA PUBLIC
+        ${PROJECT_NAME} ${OpenMM_CUDA_LIBRARY}
+    )
+endif()
+
 # Install
 if(CMAKE_INSTALL_PREFIX_INITIALIZED_TO_DEFAULT)
     set(CMAKE_INSTALL_PREFIX ${OpenMM_INSTALL_PREFIX} CACHE PATH "" FORCE)
@@ -42,5 +53,11 @@ install(DIRECTORY include/
     FILES_MATCHING PATTERN "*.h"
 )
 
+if(BUILD_CUDA_LIB)
+    install(TARGETS ${PROJECT_NAME}CUDA
+        DESTINATION ${OpenMM_INSTALL_LIBDIR}
+    )
+endif()
+
 # Build wrappers
 add_subdirectory(wrappers)

cmake/Modules/FindOpenMM.cmake

@@ -82,16 +82,15 @@ if(OpenMM_FOUND AND NOT TARGET OpenMM::OpenMM)
         endif()
 
         if(CUDA_FOUND OR CUDAToolkit_FOUND)
-            target_compile_definitions(OpenMM::OpenMM INTERFACE OPENMM_BUILD_CUDA_LIB)
+            set(BUILD_CUDA_LIB ON CACHE INTERNAL "Build CUDA variant")
             target_include_directories(OpenMM::OpenMM SYSTEM INTERFACE
                 "${CUDAToolkit_INCLUDE_DIRS}"
             )
             target_include_directories(OpenMM::OpenMM INTERFACE
                 "${OpenMM_INCLUDE_DIR}/openmm/cuda"
             )
-            target_link_libraries(OpenMM::OpenMM INTERFACE
-                ${OpenMM_CUDA_LIBRARY}
-            )
+        else()
+            set(BUILD_CUDA_LIB OFF CACHE INTERNAL "Build CUDA variant")
         endif()
     endif(OpenMM_CUDA_LIBRARY)
 endif()

cmake/Modules/OpenMMTools.cmake

@@ -18,26 +18,37 @@ for m in ('openmm', 'simtk.openmm'):
         continue
 if openmm is not None:
     libpath = openmm.version.openmm_library_path
-    print(os.path.normpath(os.path.join(libpath, os.pardir)), end='')
+    print(os.path.normpath(os.path.join(libpath, os.pardir)))
+    print(openmm.__version__, end='')
 else:
-    print('', end='')"
+    print()"
     )
     execute_process(
         COMMAND ${Python_EXECUTABLE} -c "${FIND_OpenMM_SCRIPT}"
-        OUTPUT_VARIABLE OpenMM_PATH
+        OUTPUT_VARIABLE OpenMM_PATH_AND_VER
     )
+    string(REGEX MATCH "(.*)\n(.*)" _ "${OpenMM_PATH_AND_VER}")
+    set(OpenMM_PATH ${CMAKE_MATCH_1})
+    set(OpenMM_VERSION ${CMAKE_MATCH_2})
     if(OpenMM_PATH)
         if(NOT OpenMM_ROOT)
             set(OpenMM_ROOT ${OpenMM_PATH} PARENT_SCOPE)
         endif()
         set(OpenMM_Python_EXECUTABLE ${Python_EXECUTABLE} PARENT_SCOPE)
     endif()
+    if(OpenMM_VERSION)
+        set(OpenMM_VERSION ${OpenMM_VERSION} PARENT_SCOPE)
+        if(${OpenMM_VERSION} VERSION_LESS "7.6")
+            set(OpenMM_SIMTK_API TRUE PARENT_SCOPE)
+        endif()
+    endif()
 endfunction()
 
 function(check_python_and_openmm)
     if(OpenMM_ROOT AND NOT OpenMM_Python_EXECUTABLE)
-        if(NOT Python_EXECUTABLE)
-            set(Python_EXECUTABLE "${OpenMM_ROOT}/bin/python")
+        set(Python_EXECUTABLE_GUESS "${OpenMM_ROOT}/bin/python")
+        if(NOT Python_EXECUTABLE AND EXISTS ${Python_EXECUTABLE_GUESS})
+            set(Python_EXECUTABLE ${Python_EXECUTABLE_GUESS})
         endif()
         find_openmm_with_python()
     endif()
@@ -63,5 +74,5 @@ print(os.path.normpath(os.path.join(openmm.__file__, os.pardir, os.pardir)), end
         COMMAND ${Python_EXECUTABLE} -c "${FIND_OpenMM_SCRIPT}"
         OUTPUT_VARIABLE OpenMM_Python_PATH
     )
-    set(OpenMMDLExt_Python_PATH "${OpenMM_Python_PATH}/openmm_dlext" PARENT_SCOPE)
+    set(OpenMMDLExt_Python_PATH "${OpenMM_Python_PATH}" PARENT_SCOPE)
 endfunction()

openmmapi/CMakeLists.txt

@@ -2,6 +2,10 @@ add_subdirectory(src)
 
 target_include_directories(${PROJECT_NAME} PUBLIC include)
 
+if(BUILD_CUDA_LIB)
+    target_include_directories(${PROJECT_NAME}CUDA PUBLIC include)
+endif()
+
 install(DIRECTORY include/
     DESTINATION "${CMAKE_INSTALL_INCLUDEDIR}/openmm-dlext"
     FILES_MATCHING PATTERN "*.h"

openmmapi/src/CMakeLists.txt

@@ -1,3 +1,9 @@
 target_sources(${PROJECT_NAME}
     PRIVATE ContextView.cpp DLExtForce.cpp DLExtForceImpl.cpp
 )
+
+if(BUILD_CUDA_LIB)
+    target_sources(${PROJECT_NAME}CUDA
+        PRIVATE ContextView.cpp DLExtForce.cpp DLExtForceImpl.cpp
+    )
+endif()

platforms/common/CMakeLists.txt

@@ -2,6 +2,10 @@ add_subdirectory(src)
 
 target_include_directories(${PROJECT_NAME} PUBLIC include)
 
+if(BUILD_CUDA_LIB)
+    target_include_directories(${PROJECT_NAME}CUDA PUBLIC include)
+endif()
+
 install(DIRECTORY include/
     DESTINATION "${CMAKE_INSTALL_INCLUDEDIR}/openmm-dlext/common"
     FILES_MATCHING PATTERN "*.h"

platforms/common/src/CMakeLists.txt

@@ -1 +1,5 @@
 target_sources(${PROJECT_NAME} PRIVATE DLExtKernelFactory.cpp)
+
+if(BUILD_CUDA_LIB)
+    target_sources(${PROJECT_NAME}CUDA PRIVATE DLExtKernelFactory.cpp)
+endif()

wrappers/python/CMakeLists.txt

@@ -11,14 +11,17 @@ cmake_policy(SET CMP0086 NEW)
 include(UseSWIG)
 include(${PROJECT_MODULE_PATH}/Fetchpybind11.cmake)
 
+install(DIRECTORY dlext/
+    DESTINATION ${OpenMMDLExt_Python_PATH}/openmm/dlext
+)
 install(DIRECTORY openmm_dlext/
-    DESTINATION ${OpenMMDLExt_Python_PATH}
+    DESTINATION ${OpenMMDLExt_Python_PATH}/openmm_dlext
 )
 
 # We build just a minimal SWIG wrapper for the DLExt::Force class
 # to make it compatible with the SWIG wrapper of OpenMM::Force
 # The name of the module in the SWIG headers should be "_${PySWIG_MODULE_NAME}"
-set(PySWIG_MODULE_NAME "dlpack_extension_swig")
+set(PySWIG_MODULE_NAME "api_swig")
 
 set_property(SOURCE DLExtForce.i PROPERTY CPLUSPLUS ON)
 swig_add_library(${PySWIG_MODULE_NAME} TYPE MODULE LANGUAGE python SOURCES DLExtForce.i)
@@ -27,28 +30,44 @@ target_include_directories(${PySWIG_MODULE_NAME} SYSTEM PRIVATE ${Python_INCLUDE
 target_compile_features(${PySWIG_MODULE_NAME} PRIVATE cxx_std_11)
 target_link_libraries(${PySWIG_MODULE_NAME} PRIVATE ${PROJECT_NAME})
 
+if(OpenMM_SIMTK_API)
+    target_compile_definitions(${PySWIG_MODULE_NAME} PRIVATE SIMTK_API)
+endif()
+
 if(APPLE)
     target_link_options(${PySWIG_MODULE_NAME} PRIVATE "LINKER:-undefined,dynamic_lookup")
 endif()
 
 install(TARGETS ${PySWIG_MODULE_NAME}
-    DESTINATION ${OpenMMDLExt_Python_PATH}
+    DESTINATION ${OpenMMDLExt_Python_PATH}/openmm/dlext
 )
 
 install(FILES "${CMAKE_CURRENT_BINARY_DIR}/${PySWIG_MODULE_NAME}.py"
-    DESTINATION ${OpenMMDLExt_Python_PATH}
+    DESTINATION ${OpenMMDLExt_Python_PATH}/openmm/dlext
 )
 
 # We then build a Python API for the rest of the main library with pybind11
-# The name of the module should match the one declared within the sources
-set(pybind11_MODULE_NAME "dlpack_extension")
+set(pybind11_MODULE_NAME "_api")
 
-pybind11_add_module(${pybind11_MODULE_NAME} MODULE "")
+pybind11_add_module(${pybind11_MODULE_NAME}_fallback MODULE "")
 
-target_sources(${pybind11_MODULE_NAME} PRIVATE PyDLExt.cpp)
-target_compile_features(${pybind11_MODULE_NAME} PRIVATE cxx_std_11)
-target_link_libraries(${pybind11_MODULE_NAME} PRIVATE ${PROJECT_NAME})
+target_sources(${pybind11_MODULE_NAME}_fallback PRIVATE PyDLExt.cpp)
+target_compile_features(${pybind11_MODULE_NAME}_fallback PRIVATE cxx_std_11)
+target_link_libraries(${pybind11_MODULE_NAME}_fallback PRIVATE ${PROJECT_NAME})
 
-install(TARGETS ${pybind11_MODULE_NAME}
-    DESTINATION ${OpenMMDLExt_Python_PATH}
+install(TARGETS ${pybind11_MODULE_NAME}_fallback
+    DESTINATION ${OpenMMDLExt_Python_PATH}/openmm/dlext
 )
+
+# Also build a CUDA variant when possible
+if(BUILD_CUDA_LIB)
+    pybind11_add_module(${pybind11_MODULE_NAME} MODULE "")
+
+    target_sources(${pybind11_MODULE_NAME} PRIVATE PyDLExt.cpp)
+    target_compile_features(${pybind11_MODULE_NAME} PRIVATE cxx_std_11)
+    target_link_libraries(${pybind11_MODULE_NAME} PRIVATE ${PROJECT_NAME}CUDA)
+
+    install(TARGETS ${pybind11_MODULE_NAME}
+        DESTINATION ${OpenMMDLExt_Python_PATH}/openmm/dlext
+    )
+endif()

wrappers/python/DLExtForce.i

@@ -1,4 +1,4 @@
-%module dlpack_extension_swig
+%module api_swig
 
 %pythoncode
 %{
@@ -9,7 +9,11 @@ except:
     openmm = import_module("simtk.openmm")
 %}
 
+#ifdef SIMTK_API
+%import "OpenMMForce.i"
+#else
 %import(module = "openmm") "OpenMMForce.i"
+#endif
 %include "std_string.i"
 
 %{

wrappers/python/PyDLExt.cpp

@@ -91,8 +91,13 @@ void export_Force(py::module& m)
         });
 }
 
-PYBIND11_MODULE(dlpack_extension, m)
+PYBIND11_MODULE(_api, m)
 {
+    // We want to display the members of the module as `openmm.dlext.x`
+    // instead of `openmm.dlext._api.x`.
+    py::str module_name = m.attr("__name__");
+    m.attr("__name__") = "openmm.dlext";
+
     // Enums
     py::enum_<DLDeviceType>(m, "DeviceType").value("CPU", kDLCPU).value("GPU", kDLCUDA);
     py::enum_<IdsOrdering>(m, "IdsOrdering")
@@ -111,4 +116,7 @@ PYBIND11_MODULE(dlpack_extension, m)
     m.def("inverse_masses", encapsulate<&inverseMasses>);
     m.def("positions", encapsulate<&positions>);
     m.def("velocities", encapsulate<&velocities>);
+
+    // Set back the module_name to its original value
+    m.attr("__name__") = module_name;
 }

wrappers/python/dlext/__init__.py

@@ -0,0 +1,71 @@
+# SPDX-License-Identifier: MIT
+# See LICENSE.md and CONTRIBUTORS.md at https://github.com/SSAGESLabs/openmm-dlext
+
+# flake8: noqa:F401
+
+import importlib.util
+import pathlib
+
+from .api_swig import Force, _to_capsule
+
+
+def _load_api(pattern):
+    path = next(pathlib.Path(__file__).parent.glob(pattern))
+    spec = importlib.util.spec_from_file_location("_api", path)
+    mod = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(mod)
+    return mod
+
+
+try:
+    _api = _load_api("_api.*")
+except:
+    _api = _load_api("_api_fallback.*")
+
+ContextView = _api.ContextView
+DeviceType = _api.DeviceType
+IdsOrdering = _api.IdsOrdering
+atom_ids = _api.atom_ids
+forces = _api.forces
+inverse_masses = _api.inverse_masses
+positions = _api.positions
+register_plugin = _api.register_plugin
+velocities = _api.velocities
+
+
+class Force(Force):
+    def __init__(self):
+        super(Force, self).__init__()
+        self.__alt__ = _api.Force(_to_capsule(self))
+
+    def is_present_in(self, system):
+        return self.__alt__.is_present_in(_to_capsule(system))
+
+    def add_to(self, context):
+        system = context.getSystem()
+        if not self.thisown and not self.is_present_in(system):
+            raise RuntimeError("Force already added to another system")
+        self.__alt__.add_to(_to_capsule(context), _to_capsule(system))
+        self.thisown = False
+
+    def set_callback_in(self, context, callback):
+        self.__alt__.set_callback_in(_to_capsule(context), callback)
+
+    def view(self, context):
+        return self.__alt__.view(_to_capsule(context))
+
+
+register_plugin()
+
+
+# Clean up namespace
+importlib = None
+pathlib = None
+api_swig = None
+
+del _load_api
+del register_plugin
+
+del api_swig
+del pathlib
+del importlib

wrappers/python/openmm_dlext/__init__.py

@@ -1,51 +1,11 @@
-# flake8: noqa:F401
+# SPDX-License-Identifier: MIT
+# See LICENSE.md and CONTRIBUTORS.md at https://github.com/SSAGESLabs/openmm-dlext
 
-from .dlpack_extension import ContextView, DeviceType
-from .dlpack_extension import Force as _Force
-from .dlpack_extension import (
-    IdsOrdering,
-    atom_ids,
-    forces,
-    inverse_masses,
-    positions,
-    register_plugin,
-    velocities,
-)
-from .dlpack_extension_swig import Force, _to_capsule
+# pylint: disable=unused-import,relative-beyond-top-level
+# flake8: noqa F401
 
+from warnings import warn
 
-class Force(Force):
-    def __init__(self):
-        super(Force, self).__init__()
-        self.__alt__ = _Force(_to_capsule(self))
+from openmm.dlext import *
 
-    def is_present_in(self, system):
-        return self.__alt__.is_present_in(_to_capsule(system))
-
-    def add_to(self, context):
-        system = context.getSystem()
-        if not self.thisown and not self.is_present_in(system):
-            raise RuntimeError("Force already added to another system")
-        self.__alt__.add_to(_to_capsule(context), _to_capsule(system))
-        self.thisown = False
-
-    def set_callback_in(self, context, callback):
-        self.__alt__.set_callback_in(_to_capsule(context), callback)
-
-    def view(self, context):
-        return self.__alt__.view(_to_capsule(context))
-
-
-register_plugin()
-
-
-dlpack_extension = None
-dlpack_extension_swig = None
-_dlpack_extension = None
-_dlpack_extension_swig = None
-
-del register_plugin
-del dlpack_extension
-del dlpack_extension_swig
-del _dlpack_extension
-del _dlpack_extension_swig
+warn("Importing `openmm_dlext` is deprecated. Import `openmm.dlext` instead")