MacOS-only version with MPS support

.github/workflows/Coverage.yml

@@ -23,8 +23,8 @@ jobs:
       continue-on-error: true # The one test will most likely fail, but if we skip we lose a lot of coverage.
       run: mkdir build && cd build && cmake -DBUILD_COVERAGE=1 .. && make -j8 && make install
 
-    - name: Run tests 
-      run: cd build && ctest -T Test -T Coverage
+    - name: Run tests
+      run: cd build && ctest -T Test -T Coverage --output-on-failure
 
     - name: Upload code coverage
       run: curl -Os https://uploader.codecov.io/latest/linux/codecov && chmod +x codecov && ./codecov

.gitignore

@@ -1,5 +1,6 @@
 build*/
 Debug*/
+.claude/
 
 # Ignore metadata generated by Jekyll
 _site/

CMakeLists.txt

@@ -1,5 +1,5 @@
 cmake_minimum_required (VERSION 3.18)
-project(LibMolGrid LANGUAGES C CXX CUDA)
+project(LibMolGrid LANGUAGES C CXX)
 
 #version number
 set (VERSION_MAJOR 0)
@@ -17,64 +17,106 @@ set(CMAKE_CXX_USE_RESPONSE_FILE_FOR_INCLUDES OFF)
 option(BUILD_SHARED "Build shared library" ON)
 option(BUILD_STATIC "Build static library" ON)
 option(BUILD_COVERAGE "Build with code coverage" OFF)
+option(LIBMOLGRID_ENABLE_CUDA "Build the upstream CUDA backend when available" ON)
+option(LIBMOLGRID_ENABLE_METAL "Build the Apple Metal backend when available" ${APPLE})
 
-if(${BUILD_CONTAINED_PYTHON}) 
-  #if make a mostly statically linked python module, use static boost
-   set(BUILD_STATIC ON)
-   set(Boost_USE_STATIC_LIBS ON)
+if(${BUILD_CONTAINED_PYTHON})
+  set(BUILD_STATIC ON)
+  set(Boost_USE_STATIC_LIBS ON)
 endif()
 
 list(APPEND CMAKE_MODULE_PATH "${PROJECT_SOURCE_DIR}/cmake/Modules;${PROJECT_SOURCE_DIR}/cmake")
-# get git hash
 include(cmake/git_revision.cmake)
-
-# guard against in-source builds and bad build-type strings
 include(cmake/safeguards.cmake)
-
-# place binaries and libraries according to GNU standards
 include(GNUInstallDirs)
+include(CheckLanguage)
+
 set(CMAKE_ARCHIVE_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/${CMAKE_INSTALL_LIBDIR})
 set(CMAKE_LIBRARY_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/${CMAKE_INSTALL_LIBDIR})
 set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/${CMAKE_INSTALL_BINDIR})
 
 #dependencies
-find_package(CUDA REQUIRED)
-find_package(Boost REQUIRED COMPONENTS regex unit_test_framework program_options system filesystem iostreams serialization)
+set(LIBMOLGRID_USE_CUDA OFF)
+set(LIBMOLGRID_USE_METAL OFF)
+
+if(LIBMOLGRID_ENABLE_CUDA)
+  check_language(CUDA)
+  if(CMAKE_CUDA_COMPILER)
+    enable_language(CUDA)
+    find_package(CUDA QUIET)
+    if(CUDA_FOUND)
+      set(LIBMOLGRID_USE_CUDA ON)
+    elseif(NOT APPLE OR NOT LIBMOLGRID_ENABLE_METAL)
+      message(FATAL_ERROR "CUDA was enabled, but the CUDA toolkit was not found.")
+    else()
+      message(STATUS "CUDA toolkit not found. Falling back to the Metal backend.")
+    endif()
+  elseif(NOT APPLE OR NOT LIBMOLGRID_ENABLE_METAL)
+    message(FATAL_ERROR "CUDA was enabled, but no CUDA compiler was found.")
+  else()
+    message(STATUS "CUDA compiler not found. Falling back to the Metal backend.")
+  endif()
+endif()
+
+if(NOT LIBMOLGRID_USE_CUDA AND LIBMOLGRID_ENABLE_METAL)
+  if(NOT APPLE)
+    message(FATAL_ERROR "The Metal backend is only available on Apple platforms.")
+  endif()
+  enable_language(OBJCXX)
+  set(LIBMOLGRID_USE_METAL ON)
+endif()
+
+if(NOT LIBMOLGRID_USE_CUDA AND NOT LIBMOLGRID_USE_METAL)
+  message(FATAL_ERROR "No supported GPU backend is enabled. Enable CUDA or Metal.")
+endif()
+
+find_package(Boost REQUIRED COMPONENTS regex unit_test_framework program_options filesystem iostreams serialization)
 find_package(OpenBabel3 REQUIRED)
 include_directories(SYSTEM ${OPENBABEL3_INCLUDE_DIR})
 find_package(ZLIB)
 
+if(LIBMOLGRID_USE_CUDA)
+  include_directories(SYSTEM ${CUDA_INCLUDE_DIRS})
+endif()
+
+if(LIBMOLGRID_USE_METAL)
+  find_library(METAL_FRAMEWORK Metal REQUIRED)
+  find_library(FOUNDATION_FRAMEWORK Foundation REQUIRED)
+endif()
+
 # configure a header file to pass some of the CMake settings
-# to the source code
 configure_file (
   "${PROJECT_SOURCE_DIR}/cmake/config.h.in"
   "${PROJECT_BINARY_DIR}/include/libmolgrid/config.h"
   )
 
 set(CMAKE_CXX_STANDARD 17)
-set(CMAKE_CUDA_STANDARD 17)
+if(LIBMOLGRID_USE_CUDA)
+  set(CMAKE_CUDA_STANDARD 17)
+  set(CMAKE_CUDA_ARCHITECTURES "all-major")
+  set(CMAKE_CUDA_FLAGS_RELEASE "-O3 -g -lineinfo")
+  set(CMAKE_CUDA_FLAGS_DEBUG "-O0 -g -G")
+endif()
+if(LIBMOLGRID_USE_METAL)
+  set(CMAKE_OBJCXX_STANDARD 17)
+endif()
 
-if(CMAKE_CXX_COMPILER_ID MATCHES GNU)
-    set(CMAKE_CXX_FLAGS         "${CMAKE_CXX_FLAGS} -Wall -Wno-unknown-pragmas -Werror")
+if(CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang|AppleClang")
+    # OpenBabel 3.1.1 uses std::binary_function which was removed in C++17.
+    # Restore it via the compatibility shim in libc++.
+    set(CMAKE_CXX_FLAGS         "${CMAKE_CXX_FLAGS} -Wall -Wno-unknown-pragmas -D_LIBCPP_ENABLE_CXX17_REMOVED_UNARY_BINARY_FUNCTION")
     set(CMAKE_CXX_FLAGS_DEBUG   "-O0 -g3")
     if (BUILD_COVERAGE)
-	set(CMAKE_CXX_FLAGS_RELEASE "-O3 -g -fprofile-arcs -ftest-coverage")
+        set(CMAKE_CXX_FLAGS_RELEASE "-O3 -g -fprofile-arcs -ftest-coverage")
     else()
         set(CMAKE_CXX_FLAGS_RELEASE "-O3 -g")
-    endif()	
+    endif()
 endif()
 
-set(CMAKE_CUDA_ARCHITECTURES "all-major")
-
-set(CMAKE_CUDA_FLAGS_RELEASE "-O3 -g -lineinfo")
-set(CMAKE_CUDA_FLAGS_DEBUG "-O0 -g -G")
-
 # add the binary tree to the search path for include files
-# so that we will find libmolgrid_config.h
 include_directories("${PROJECT_BINARY_DIR}")
 include_directories(BEFORE ${CMAKE_SOURCE_DIR}/include)
 include_directories(${PROJECT_BINARY_DIR}/include)
-include_directories(SYSTEM ${CUDA_INCLUDE_DIRS})
 include_directories(SYSTEM ${Boost_INCLUDE_DIR})
 
 add_subdirectory(src)
@@ -84,5 +126,4 @@ add_subdirectory(docs)
 # enable testing
 include(CTest)
 enable_testing()
-# define tests
 add_subdirectory(test)

README.md

@@ -1,9 +1,15 @@
-libmolgrid
-==========
-[![Github CI](https://github.com/gnina/libmolgrid/actions/workflows/CI.yml/badge.svg)](https://github.com/gnina/libmolgrid/actions)
-[![codecov](https://codecov.io/gh/gnina/libmolgrid/branch/master/graph/badge.svg?token=gUFisf34rf)](https://codecov.io/gh/gnina/libmolgrid)
-
-libmolgrid is under active development, but should be suitable for use by early adopters.
+libmolgrid fork with CUDA and Metal backends
+============================================
+
+> This fork keeps the upstream CUDA backend available and adds an Apple Silicon
+> Metal backend for macOS builds.
+>
+> Backend selection is handled at CMake configure time:
+> `CUDA` is used when enabled and available, otherwise Apple builds can fall
+> back to `Metal`.
+>
+> Upstream project information:
+> [https://gnina.github.io/libmolgrid/](https://gnina.github.io/libmolgrid/)
 
 If you use libmolgrid in your research, please cite:
 
@@ -22,22 +28,24 @@ If you use libmolgrid in your research, please cite:
 }
 ```
 
-## Documentation
-
-[https://gnina.github.io/libmolgrid/](https://gnina.github.io/libmolgrid/)
-
 ## Installation
 
 ### PIP
 
-```pip install molgrid```
+```bash
+pip install molgrid
+```
 
 ### conda
-```conda install -c gnina molgrid```
+
+```bash
+conda install -c gnina molgrid
+```
 
 ### Build from Source
 
-```apt install git build-essential libboost-all-dev python3-pip rapidjson-dev
+```bash
+apt install git build-essential libboost-all-dev python3-pip rapidjson-dev
 pip3 install numpy pytest pyquaternion
 ```
 
@@ -49,7 +57,8 @@ pip3 install numpy pytest pyquaternion
 
 `apt install libeigen3-dev libboost-all-dev`
 
-```git clone https://github.com/gnina/libmolgrid.git
+```bash
+git clone https://github.com/gnina/libmolgrid.git
 cd libmolgrid
 mkdir build
 cd build
@@ -58,13 +67,69 @@ make -j8
 sudo make install
 ```
 
+### macOS / Apple Silicon
 
+For this fork, macOS builds can use the Metal backend when CUDA is not available.
+Install dependencies via [Homebrew](https://brew.sh):
 
+```bash
+brew install cmake boost open-babel rapidjson eigen python
+pip3 install numpy pytest pyquaternion
+```
 
-## Example
+Xcode Command Line Tools are required for the Metal toolchain:
+
+```bash
+xcode-select --install
+```
+
+Build with Metal explicitly enabled:
+
+```bash
+git clone https://github.com/fnachon/libmolgrid.git
+cd libmolgrid
+mkdir build
+cd build
+cmake .. -DCMAKE_BUILD_TYPE=Release -DLIBMOLGRID_ENABLE_CUDA=OFF
+make -j$(sysctl -n hw.logicalcpu)
+sudo make install
 ```
+
+To build with the upstream CUDA backend instead, disable Metal and use a CUDA-enabled toolchain:
+
+```bash
+cmake .. -DCMAKE_BUILD_TYPE=Release -DLIBMOLGRID_ENABLE_METAL=OFF
+```
+
+When the Metal backend is selected, the shaders in `src/metal/shaders.metal`
+are compiled automatically into a `.metallib`.
+
+On Apple Silicon, GPU and CPU share unified memory — no explicit host↔device
+copies occur at runtime.
+
+### Run Tests
+
+Tests must be run from the `build/bin/` directory so that relative data paths resolve correctly:
+
+```bash
+cd build/bin
+./test_gridmaker_cpp
+./test_gridmaker_mps_mm   # Metal backend
+./test_grid_mps_mm        # Metal backend
+./test_mgrid_mps_mm       # Metal backend
+./test_transform_mps_mm   # Metal backend
+
+# or, on CUDA builds
+./test_gridmaker_cu
+./test_grid_cu
+./test_mgrid_cu
+./test_transform_cu
+```
+
+## Example
+
+```python
 import molgrid
-import pytest
 import numpy as np
 import torch
 import torch.nn as nn
@@ -110,36 +175,32 @@ def test_train_torch_cnn():
         if isinstance(m, nn.Conv3d) or isinstance(m, nn.Linear):
             init.xavier_uniform_(m.weight.data)
 
-    batch_size = 50
-    e = molgrid.ExampleProvider(data_root=datadir+"/structs",balanced=True,shuffle=True)
+    e = molgrid.ExampleProvider(data_root=datadir+"/structs", balanced=True, shuffle=True)
     e.populate(fname)
 
     gmaker = molgrid.GridMaker()
     dims = gmaker.grid_dimensions(e.num_types())
-    tensor_shape = (batch_size,)+dims
+    tensor_shape = (batch_size,) + dims
 
-    model = Net(dims).to('cuda')
+    model = Net(dims).to('mps')
     model.apply(weights_init)
 
     optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
 
-    input_tensor = torch.zeros(tensor_shape, dtype=torch.float32, device='cuda')
+    input_tensor = torch.zeros(tensor_shape, dtype=torch.float32, device='mps')
     float_labels = torch.zeros(batch_size, dtype=torch.float32)
 
     losses = []
     for iteration in range(100):
-        #load data
         batch = e.next_batch(batch_size)
-        gmaker.forward(batch, input_tensor, 0, random_rotation=False) #not rotating since convergence is faster this way
+        gmaker.forward(batch, input_tensor, 0, random_rotation=False)
         batch.extract_label(0, float_labels)
-        labels = float_labels.long().to('cuda')
+        labels = float_labels.long().to('mps')
 
         optimizer.zero_grad()
         output = model(input_tensor)
-        loss = F.cross_entropy(output,labels)
+        loss = F.cross_entropy(output, labels)
         loss.backward()
         optimizer.step()
         losses.append(float(loss))
-
 ```
-

cmake/Modules/FindOpenBabel3.cmake

@@ -36,6 +36,8 @@ else()
         /usr/local/openbabel/include
         /usr/local/openbabel3/include/openbabel3
         /usr/local/openbabel3/include
+        /opt/homebrew/include/openbabel3
+        /opt/homebrew/include
         ~/include/openbabel3
         ~/include
       )
@@ -57,6 +59,7 @@ else()
       $ENV{OPENBABEL_BASE}/lib
       /usr/lib
       /usr/local/lib
+      /opt/homebrew/lib
       ~/lib
       $ENV{LD_LIBRARY_PATH}
     )
@@ -72,4 +75,3 @@ message("Setting openbabel found ${OPENBABEL3_FOUND}")
 
   mark_as_advanced(OPENBABEL3_INCLUDE_DIR OPENBABEL3_LIBRARIES)
 endif()
-

cmake/config.h.in

@@ -1,7 +1,15 @@
+#ifndef LIBMOLGRID_CONFIG_H_
+#define LIBMOLGRID_CONFIG_H_
+
 #define VERSION_MAJOR @VERSION_MAJOR@
 #define VERSION_MINOR @VERSION_MINOR@
 #define VERSION_PATCH @VERSION_PATCH@
 
+#cmakedefine01 LIBMOLGRID_USE_CUDA
+#cmakedefine01 LIBMOLGRID_USE_METAL
+
 namespace libmolgrid {
-    const char *GIT_REVISION = "@GIT_REVISION@";
+    inline const char *GIT_REVISION = "@GIT_REVISION@";
 }
+
+#endif /* LIBMOLGRID_CONFIG_H_ */