diff --git a/.claude/settings.local.json b/.claude/settings.local.json new file mode 100644 index 000000000..ce2d88a7f --- /dev/null +++ b/.claude/settings.local.json @@ -0,0 +1,13 @@ +{ + "permissions": { + "allow": [ + "Bash(/Users/florian/miniconda3/envs/tmol/bin/python3.13 -m pytest tmol/tests/pack/test_pack_rotamers.py -k \"mps and not gradcheck\" -v)", + "Bash(python -m pytest tmol/tests/test_mps.py --no-header)", + "Bash(echo \"---EXIT: $?\")", + "Bash(conda activate:*)", + "Bash(python -m pytest tmol/tests/ -k mps --no-header -q --ignore=tmol/tests/score/common/test_uaid_util.py)", + "Bash(/Users/florian/miniconda3/envs/tmol/bin/python -m pytest tmol/tests/ -k mps --no-header -q --ignore=tmol/tests/score/common/test_uaid_util.py)", + "Bash(/Users/florian/miniconda3/envs/tmol/bin/python -c ':*)" + ] + } +} diff --git a/tmol/kinematics/compiled/compiled_inverse_kin.py b/.cmake/api/v1/query/cache-v2 similarity index 100% rename from tmol/kinematics/compiled/compiled_inverse_kin.py rename to .cmake/api/v1/query/cache-v2 diff --git a/.cmake/api/v1/query/cmakeFiles-v1 b/.cmake/api/v1/query/cmakeFiles-v1 new file mode 100644 index 000000000..e69de29bb diff --git a/.cmake/api/v1/query/codemodel-v2 b/.cmake/api/v1/query/codemodel-v2 new file mode 100644 index 000000000..e69de29bb diff --git a/.cmake/api/v1/query/toolchains-v1 b/.cmake/api/v1/query/toolchains-v1 new file mode 100644 index 000000000..e69de29bb diff --git a/.gitignore b/.gitignore index 3fa142cc7..ce9b5b34c 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,21 @@ .torch_extensions +# scikit-build-core / CMake in-source build artifacts +.cmake/api/v1/reply/ +.ninja_deps +.ninja_log +.skbuild-info.json +CMakeCache.txt +CMakeFiles/ +CMakeInit.txt +Makefile +build.ninja +cmake_install.cmake +metal_air/ + +# Output PDB files from pack_rotamers runs +pack_rotamers_*.pdb + # Conda environment .conda # Rendered environment definitions diff --git a/CMakeLists.txt b/CMakeLists.txt index 175c39b5e..c48385868 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -89,6 +89,29 @@ else() message(STATUS "tmol: CUDA not found or disabled; building CPU-only extensions") endif() +# ═══════════════════════════════════════════════════════════════════════════════ +# MPS (Apple Metal) detection +# Enabled automatically on macOS when xcrun and the Metal SDK are present. +# Can be forced on/off with -DTMOL_BUILD_MPS=ON/OFF. +# ═══════════════════════════════════════════════════════════════════════════════ +if(APPLE) + find_program(XCRUN_EXECUTABLE xcrun) + if(XCRUN_EXECUTABLE) + execute_process( + COMMAND ${XCRUN_EXECUTABLE} --sdk macosx --show-sdk-path + OUTPUT_VARIABLE _MACOSX_SDK_PATH + OUTPUT_STRIP_TRAILING_WHITESPACE + RESULT_VARIABLE _SDK_RESULT + ) + if(_SDK_RESULT EQUAL 0 AND EXISTS "${_MACOSX_SDK_PATH}/System/Library/Frameworks/Metal.framework") + set(_TMOL_MPS_AVAILABLE TRUE) + endif() + endif() +endif() + +option(TMOL_BUILD_MPS "Build Apple Metal/MPS backend" ${_TMOL_MPS_AVAILABLE}) +message(STATUS "tmol: TMOL_BUILD_MPS = ${TMOL_BUILD_MPS}") + # ═══════════════════════════════════════════════════════════════════════════════ # Dependencies # ═══════════════════════════════════════════════════════════════════════════════ @@ -109,8 +132,12 @@ message(STATUS "tmol: Python include dirs = ${_TMOL_PYTHON_INCLUDE_DIRS}") # PyTorch 2.13 and newer require C++20 for third-party extensions. Keep the # older release lanes on C++17 so their established CUDA toolchains remain # unchanged. +# KMP_DUPLICATE_LIB_OK=TRUE is required on macOS when multiple copies of the +# OpenMP runtime are loaded (e.g. conda torch + Homebrew clang-rt), otherwise +# Python aborts with "OMP Error #15" before printing any output. execute_process( - COMMAND ${Python_EXECUTABLE} -c "import torch; v=torch.__version__.split('.'); print(f'{v[0]};{v[1]}')" + COMMAND ${CMAKE_COMMAND} -E env KMP_DUPLICATE_LIB_OK=TRUE + ${Python_EXECUTABLE} -c "import torch; v=torch.__version__.split('.'); print(f'{v[0]};{v[1]}')" OUTPUT_VARIABLE _TORCH_VERSION OUTPUT_STRIP_TRAILING_WHITESPACE ) @@ -158,6 +185,38 @@ if(TMOL_HAS_CUDA) message(STATUS "tmol: Setting TORCH_CUDA_ARCH_LIST=${_TORCH_ARCH_STR} for find_package(Torch)") endif() +# Prepend the active Python's torch cmake directory to CMAKE_PREFIX_PATH so +# that find_package(Torch) resolves to the conda/pip installation rather than +# any system or Homebrew libtorch that may appear in CMake's default search +# paths (e.g. /opt/homebrew/lib on macOS). +execute_process( + COMMAND ${CMAKE_COMMAND} -E env KMP_DUPLICATE_LIB_OK=TRUE + ${Python_EXECUTABLE} -c "import torch; print(torch.utils.cmake_prefix_path)" + OUTPUT_VARIABLE _TORCH_CMAKE_PREFIX + OUTPUT_STRIP_TRAILING_WHITESPACE + RESULT_VARIABLE _TORCH_CMAKE_RESULT +) +if(_TORCH_CMAKE_RESULT EQUAL 0 AND _TORCH_CMAKE_PREFIX) + list(PREPEND CMAKE_PREFIX_PATH "${_TORCH_CMAKE_PREFIX}") + message(STATUS "tmol: Prepending torch cmake prefix: ${_TORCH_CMAKE_PREFIX}") +else() + message(WARNING "tmol: Could not determine torch cmake prefix path from Python — " + "find_package(Torch) may pick up Homebrew or system libtorch") +endif() + +# Similarly discover pybind11's cmake dir from the active Python environment. +execute_process( + COMMAND ${CMAKE_COMMAND} -E env KMP_DUPLICATE_LIB_OK=TRUE + ${Python_EXECUTABLE} -c "import pybind11; print(pybind11.get_cmake_dir())" + OUTPUT_VARIABLE _PYBIND11_CMAKE_DIR + OUTPUT_STRIP_TRAILING_WHITESPACE + RESULT_VARIABLE _PYBIND11_CMAKE_RESULT +) +if(_PYBIND11_CMAKE_RESULT EQUAL 0 AND _PYBIND11_CMAKE_DIR) + list(PREPEND CMAKE_PREFIX_PATH "${_PYBIND11_CMAKE_DIR}") + message(STATUS "tmol: Prepending pybind11 cmake dir: ${_PYBIND11_CMAKE_DIR}") +endif() + find_package(Torch REQUIRED) find_package(pybind11 CONFIG REQUIRED) @@ -303,6 +362,57 @@ function(tmol_add_pybind_cpp_ext TARGET_NAME INSTALL_DIR EXT_NAME) install(TARGETS ${TARGET_NAME} DESTINATION "${INSTALL_DIR}") endfunction() +# Common compile options for MPS (Objective-C++ + Metal) extensions +function(tmol_set_mps_flags TARGET) + target_compile_options(${TARGET} PRIVATE + $<$:-O3 -w -DWITH_MPS> + # ObjC++ files (.mm) — enable ARC and pass Metal SDK headers + $<$: + -O3 -w -DWITH_MPS + -fobjc-arc + -fmodules + -fcxx-modules + > + ) + target_include_directories(${TARGET} PRIVATE ${TMOL_INCLUDE_DIRS}) + target_link_libraries(${TARGET} PRIVATE + ${TORCH_LIBRARIES} + "-framework Metal" + "-framework Foundation" + ) +endfunction() + +# Compile Metal shaders: → tmol_primitives.metallib +# The metallib is placed in CMAKE_SOURCE_DIR/tmol/ alongside _C.so. +function(tmol_compile_metal_shaders) + set(_METAL_SOURCES ${ARGN}) + set(_AIR_FILES "") + + foreach(_SRC ${_METAL_SOURCES}) + get_filename_component(_BASE "${_SRC}" NAME_WE) + set(_AIR "${CMAKE_BINARY_DIR}/metal_air/${_BASE}.air") + add_custom_command( + OUTPUT "${_AIR}" + COMMAND ${CMAKE_COMMAND} -E make_directory "${CMAKE_BINARY_DIR}/metal_air" + COMMAND ${XCRUN_EXECUTABLE} -sdk macosx metal + -std=metal3.0 -O2 -c "${_SRC}" -o "${_AIR}" + DEPENDS "${_SRC}" + COMMENT "Compiling Metal shader ${_SRC}" + ) + list(APPEND _AIR_FILES "${_AIR}") + endforeach() + + set(_METALLIB "${CMAKE_SOURCE_DIR}/tmol/tmol_primitives.metallib") + add_custom_command( + OUTPUT "${_METALLIB}" + COMMAND ${XCRUN_EXECUTABLE} -sdk macosx metallib ${_AIR_FILES} -o "${_METALLIB}" + DEPENDS ${_AIR_FILES} + COMMENT "Linking tmol_primitives.metallib" + ) + add_custom_target(tmol_metallib ALL DEPENDS "${_METALLIB}") + install(FILES "${_METALLIB}" DESTINATION tmol) +endfunction() + # ═══════════════════════════════════════════════════════════════════════════════ # 1. tmol._C — monolithic TORCH_LIBRARY extension (14 op namespaces) # ═══════════════════════════════════════════════════════════════════════════════ @@ -330,11 +440,9 @@ set(_C_SOURCES # score/cartbonded tmol/score/cartbonded/potentials/compiled.ops.cpp tmol/score/cartbonded/potentials/cartbonded_pose_score.cpu.cpp - tmol/score/cartbonded/potentials/cartbonded_pose_score.cuda.cu # score/genbonded tmol/score/genbonded/potentials/compiled.ops.cpp tmol/score/genbonded/potentials/genbonded_pose_score.cpu.cpp - tmol/score/genbonded/potentials/genbonded_pose_score.cuda.cu # score/constraint tmol/score/constraint/potentials/compiled.ops.cpp tmol/score/constraint/potentials/constraint_score.cpu.cpp @@ -360,6 +468,7 @@ set(_C_SOURCES tmol/score/lk_ball/potentials/gen_pose_waters.cpu.cpp ) +# ── CUDA sources (Linux / NVIDIA GPU only) ───────────────────────────────── if(TMOL_HAS_CUDA) list(APPEND _C_SOURCES tmol/io/details/compiled/gen_pose_leaf_atoms.cuda.cu @@ -370,6 +479,7 @@ if(TMOL_HAS_CUDA) tmol/pose/compiled/apsp.cuda.cu tmol/score/backbone_torsion/potentials/backbone_torsion_pose_score.cuda.cu tmol/score/cartbonded/potentials/cartbonded_pose_score.cuda.cu + tmol/score/genbonded/potentials/genbonded_pose_score.cuda.cu tmol/score/constraint/potentials/constraint_score.cuda.cu tmol/score/disulfide/potentials/disulfide_pose_score.cuda.cu tmol/score/dunbrack/potentials/dunbrack_pose_score.cuda.cu @@ -382,6 +492,42 @@ if(TMOL_HAS_CUDA) ) endif() +# ── MPS (Apple Metal) sources ────────────────────────────────────────────── +if(TMOL_BUILD_MPS) + list(APPEND _C_SOURCES + # Metal context ObjC++ bridge + tmol/score/common/metal_context.mm + # io + tmol/io/details/compiled/gen_pose_leaf_atoms.mps.mm + tmol/io/details/compiled/resolve_his_taut.mps.mm + # kinematics + tmol/kinematics/compiled/compiled.mps.mm + # pack + tmol/pack/compiled/compiled.mps.mm + tmol/pack/rotamer/dunbrack/compiled.mps.mm + # pose + tmol/pose/compiled/apsp.mps.mm + # score terms + tmol/score/backbone_torsion/potentials/backbone_torsion_pose_score.mps.mm + tmol/score/cartbonded/potentials/cartbonded_pose_score.mps.mm + tmol/score/genbonded/potentials/genbonded_pose_score.mps.mm + tmol/score/constraint/potentials/constraint_score.mps.mm + tmol/score/disulfide/potentials/disulfide_pose_score.mps.mm + tmol/score/dunbrack/potentials/dunbrack_pose_score.mps.mm + tmol/score/elec/potentials/elec_pose_score.mps.mm + tmol/score/hbond/potentials/hbond_pose_score.mps.mm + tmol/score/hbond/potentials/gen_hbond_bases.mps.mm + tmol/score/ljlk/potentials/ljlk_pose_score.mps.mm + tmol/score/lk_ball/potentials/lk_ball_pose_score.mps.mm + tmol/score/lk_ball/potentials/gen_pose_waters.mps.mm + ) + + # Compile Metal shader library + tmol_compile_metal_shaders( + ${CMAKE_SOURCE_DIR}/tmol/score/common/metal_primitives.metal + ) +endif() + add_library(_C MODULE ${_C_SOURCES}) # _C is a MODULE library loaded by torch, not a regular Python extension. @@ -403,6 +549,17 @@ if(TMOL_HAS_CUDA) else() tmol_set_cpp_flags(_C) endif() + +if(TMOL_BUILD_MPS) + # Apply MPS-specific compile options (ObjC++ ARC, Metal framework, -DWITH_MPS) + tmol_set_mps_flags(_C) + # Also define WITH_MPS for the plain C++ sources compiled into _C + target_compile_definitions(_C PRIVATE WITH_MPS) + # Enable Objective-C++ language support for .mm sources + set_property(TARGET _C PROPERTY OBJCXX_STANDARD 17) + set_property(TARGET _C PROPERTY OBJCXX_STANDARD_REQUIRED ON) +endif() + install(TARGETS _C DESTINATION tmol) diff --git a/DEVELOPMENT.md b/DEVELOPMENT.md index 1d00ed4d7..e9bcdfc31 100644 --- a/DEVELOPMENT.md +++ b/DEVELOPMENT.md @@ -15,25 +15,41 @@ This document covers building, testing, and contributing to tmol. ## Local Setup +**NVIDIA GPU (Linux):** ```bash git clone https://github.com/uw-ipd/tmol.git && cd tmol -pip install -e ".[dev]" # builds C++/CUDA extensions via CMake +pip install -e ".[dev]" # builds C++/CUDA extensions via CMake ``` Requirements: Python 3.11+, PyTorch 2.8+, C++17 compiler, CMake 3.18+. CUDA toolkit (`nvcc`) is optional — without it, only CPU extensions are built. Pre-built wheels are published for Python `cp311`-`cp314`. +**Apple Silicon Mac (macOS):** + +> [!IMPORTANT] +> MPS support lives in the **[fnachon/tmol](https://github.com/fnachon/tmol)** fork. +> Clone that repository for Apple Silicon development. + +```bash +git clone https://github.com/fnachon/tmol.git && cd tmol +pip install -e ".[dev,mps]" # builds C++/Metal extensions via CMake +``` +Requirements: Python 3.10+, PyTorch 2.5+, macOS 13+, Xcode Command Line Tools (`xcode-select --install`), CMake 3.18+. No `nvcc` needed. + ## Building Extensions -tmol ships custom C++/CUDA kernels that are compiled via CMake (using scikit-build-core as the build backend). `pip install -e .` handles compilation automatically. +tmol ships custom C++/CUDA/Metal kernels compiled via CMake (using scikit-build-core as the build backend). `pip install -e .` handles compilation automatically. ```bash -# Full build (production extensions only) +# Full build — NVIDIA GPU pip install -e . +# Full build — Apple Silicon (MPS backend) +pip install -e ".[mps]" + # Build with test extensions pip install -e . -Ccmake.define.TMOL_BUILD_TESTS=ON -# Target specific GPU architectures (default: "80;86;89;90") +# Target specific CUDA GPU architectures (default: "80;86;89;90") pip install -e . -Ccmake.define.CMAKE_CUDA_ARCHITECTURES="80;90" # Control parallelism @@ -44,12 +60,30 @@ CMake build options: | Variable | Default | Description | |----------|---------|-------------| -| `CMAKE_CUDA_ARCHITECTURES` | `80;86;89;90` | GPU compute capabilities to compile for | +| `CMAKE_CUDA_ARCHITECTURES` | `80;86;89;90` | CUDA GPU compute capabilities to compile for | | `TMOL_BUILD_TESTS` | `OFF` | Build test-only C++/CUDA extensions | +| `TMOL_BUILD_MPS` | auto-detected | Build MPS (Metal) backend; auto-enabled when `xcrun` and Metal SDK are found | | `TMOL_NVCC_THREADS` | `4` | Threads per nvcc invocation | | `TMOL_ENABLE_CUDA` | `ON` | Set to `OFF` for CPU-only build (no `nvcc` needed) | | `MAX_JOBS` | auto | Max parallel compilation jobs | +### MPS / Metal build notes + +The MPS backend is enabled automatically on macOS when `xcrun` and the Metal SDK are present (they ship with Xcode Command Line Tools). The build compiles: + +- Objective-C++ (`.mm`) bridge files that call Metal API +- Metal Shading Language kernels (`.metal` → `.air` → `tmol_primitives.metallib`) for GPU-accelerated scan, reduce, and segmented scan primitives + +To explicitly enable or disable the MPS backend: + +```bash +# Force-enable (will fail if Metal SDK is absent) +pip install -e . -Ccmake.define.TMOL_BUILD_MPS=ON + +# Force-disable (CPU-only build on macOS) +pip install -e . -Ccmake.define.TMOL_BUILD_MPS=OFF +``` + ## Extension Loading: AOT vs JIT tmol's C++/CUDA kernels can be loaded in two ways: @@ -105,6 +139,10 @@ JIT mode requires `nvcc` and CUDA headers. You can either: pip install .[cuda] ``` +### MPS / Metal and JIT mode + +The MPS backend does not use JIT compilation — Metal shaders are always compiled ahead-of-time at build time via `xcrun metal`. Setting `TMOL_USE_JIT=1` on macOS still compiles the C++/Objective-C++ bridge code via `torch.utils.cpp_extension`, but the `.metallib` binary is loaded from disk. No additional environment variables are needed for MPS. + ## Running Tests ```bash @@ -114,8 +152,11 @@ pytest tmol/tests/ -v # Specific test file pytest tmol/tests/score/test_score_function.py -v -# Only CPU tests (skip cuda-parametrized tests) -pytest tmol/tests/ -v -k "not cuda" +# Only CPU tests (skip cuda- and mps-parametrized tests) +pytest tmol/tests/ -v -k "not cuda and not mps" + +# Only MPS tests (Apple Silicon) +pytest tmol/tests/test_mps.py -v # With coverage pytest tmol/tests/ --cov=./tmol --junitxml=results.xml @@ -124,6 +165,23 @@ pytest tmol/tests/ --cov=./tmol --junitxml=results.xml pytest --benchmark-enable --benchmark-only --benchmark-max-time=.1 ``` +### MPS test suite + +> [!NOTE] +> MPS tests require the [fnachon/tmol](https://github.com/fnachon/tmol) fork — the upstream repository does not include MPS patches. + +`tmol/tests/test_mps.py` contains a five-layer smoke test for the Apple Silicon backend: + +| Layer | What it checks | +|-------|---------------| +| 1 — Tensor plumbing | MPS availability, creation, matmul, autograd | +| 2 — Primitives | cumsum, reduce, elementwise ops via PyTorch wrappers | +| 3 — Dispatch macro | Pose stack construction on MPS (exercises compiled ops) | +| 4 — Forward pass | CartBonded, Elec, LJLK, HBond, full beta2016 score function | +| 5 — CPU consistency | MPS scores and gradients match CPU within float32 tolerance | + +All tests are automatically skipped on non-Apple-Silicon machines via the `@requires_mps` mark. + ### Ligand charges Partial charges come exclusively from the SMILES -> OpenBabel MMFF94 mol2 step and @@ -137,7 +195,7 @@ parameter-generation parity is the guanfeng DUD-80 SMILES suite ### Testing a specific release ```bash -# Install matching PyTorch first (example: x86_64 manylinux cu128/torch2.10) +# CUDA/Linux: install matching PyTorch first (example: x86_64 manylinux cu128/torch2.10) pip install "torch==2.10.*" --index-url https://download.pytorch.org/whl/cu128 # Install a release wheel from GitHub @@ -146,6 +204,9 @@ pip install https://github.com/uw-ipd/tmol/releases/download/vX.Y.Z/tmol-X.Y.Z+c # Or install a specific branch/tag from source pip install git+https://github.com/uw-ipd/tmol.git@vX.Y.Z +# MPS/macOS: install a specific branch/tag from the MPS fork +pip install git+https://github.com/fnachon/tmol.git@master + # Run tests against it pytest --pyargs tmol.tests -v ``` @@ -186,6 +247,9 @@ tmol uses GitHub Actions for all CI: | `wheel-smoke.yml` | Push to wheel feature branches, manual | Builds and installs the complete 32-wheel manylinux matrix, checks auditwheel metadata and glibc-2.28 portability, and loads a representative wheel on the self-hosted GPU runner. | | `publish.yml` | Push `v*` tag, manual | Builds manylinux wheels (GPU + CPU) + sdist, uploads sdist to PyPI, uploads wheels to a GitHub Release. | +> [!NOTE] +> MPS tests (`tmol/tests/test_mps.py`) are not yet part of the automated CI pipeline, which runs on Linux GPU runners. Run them locally on an Apple Silicon Mac with `pytest tmol/tests/test_mps.py -v`. + ### CI architecture ``` diff --git a/README.md b/README.md index bd0771b7b..cc6d7f84c 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,8 @@ # Tmol -`tmol` (TensorMol) is a GPU-accelerated reimplementation of the Rosetta molecular modeling energy function (`beta_nov2016_cart`) in PyTorch with custom C++/CUDA kernels. It computes energies and derivatives for protein structures and supports gradient-based minimization, enabling ML models to incorporate biophysical scoring during training or to refine predicted structures with Rosetta's experimentally validated energy function. +`tmol` (TensorMol) is a GPU-accelerated reimplementation of the Rosetta molecular modeling energy function (`beta_nov2016_cart`) in PyTorch with custom C++/CUDA/Metal kernels. It computes energies and derivatives for protein structures and supports gradient-based minimization, enabling ML models to incorporate biophysical scoring during training or to refine predicted structures with Rosetta's experimentally validated energy function. + +tmol runs on **NVIDIA GPUs** (CUDA), **Apple Silicon Macs** (MPS / Metal), and **CPU**. Full documentation: [tmol Wiki](https://github.com/uw-ipd/tmol/wiki/DevHome) @@ -14,9 +16,61 @@ Full documentation: [tmol Wiki](https://github.com/uw-ipd/tmol/wiki/DevHome) ## Installation -### Pre-built wheels (recommended) +### Apple Silicon / MPS (macOS) + +tmol runs natively on Apple Silicon (M1/M2/M3/M4) via PyTorch's Metal Performance Shaders (MPS) backend. No CUDA toolkit or NVIDIA GPU is needed. + +> [!IMPORTANT] +> MPS support is maintained in the **[fnachon/tmol](https://github.com/fnachon/tmol)** fork. +> The upstream [uw-ipd/tmol](https://github.com/uw-ipd/tmol) repository targets NVIDIA GPUs (CUDA/Linux). +> Use `https://github.com/fnachon/tmol` for Apple Silicon. + +**Requirements:** +- macOS 13.0 (Ventura) or later +- Apple Silicon Mac (M-series) +- PyTorch ≥ 2.0 with MPS support (`torch.backends.mps.is_available()` returns `True`) +- Xcode Command Line Tools (`xcode-select --install`) +- Python 3.10+ + +**Install from source (MPS):** + +```bash +# Install PyTorch with MPS support (ships in the standard macOS wheel) +pip install torch + +# Clone the MPS-enabled fork +git clone https://github.com/fnachon/tmol.git && cd tmol +pip install -e ".[dev,mps]" +``` + +**Verify MPS is working:** + +```python +import torch +print(torch.backends.mps.is_available()) # must be True + +import tmol +pose_stack = tmol.pose_stack_from_pdb("1ubq.pdb", device=torch.device("mps")) +sfxn = tmol.beta2016_score_function(torch.device("mps")) +scorer = sfxn.render_whole_pose_scoring_module(pose_stack) +print(scorer(pose_stack.coords)) +``` -Pre-built wheels ship with **ahead-of-time (AOT) compiled** C++/CUDA extensions, so install does not require `nvcc`. +> [!NOTE] +> The MPS backend uses Apple's unified memory architecture — CPU and GPU share the same physical RAM — so there is no host↔device copy overhead. All energy terms, gradients, and minimization work identically to CUDA. + +> [!TIP] +> Run the MPS smoke tests to confirm everything is wired up: +> ```bash +> pytest tmol/tests/test_mps.py -v +> ``` + +--- + +### Pre-built wheels (Linux / NVIDIA GPU only) + +Pre-built wheels ship with **ahead-of-time (AOT) compiled** C++/CUDA extensions — no `nvcc` or CUDA toolkit needed at install time. +MPS users should install [from source](#from-source) using the [fnachon/tmol](https://github.com/fnachon/tmol) fork. tmol uses two channels: @@ -213,6 +267,7 @@ TMOL_DISABLE_WHEEL_FETCH=1 pip install "tmol[dev]" ### From source ```bash +# NVIDIA GPU (upstream repository) git clone https://github.com/uw-ipd/tmol.git && cd tmol pip install -e ".[dev]" # builds extensions via CMake (CUDA auto-detected) ``` @@ -223,10 +278,11 @@ If you don't have a CUDA toolkit, the build automatically falls back to CPU-only pip install -e . -Ccmake.define.TMOL_ENABLE_CUDA=OFF ``` -For macOS, install from source (CPU-only build): +For Apple Silicon (MPS/Metal backend), use the [fnachon/tmol](https://github.com/fnachon/tmol) fork instead: ```bash -pip install -e . -Ccmake.define.TMOL_ENABLE_CUDA=OFF +git clone https://github.com/fnachon/tmol.git && cd tmol +pip install -e ".[dev,mps]" # builds C++/Metal extensions via CMake ``` ## Usage @@ -234,13 +290,18 @@ pip install -e . -Ccmake.define.TMOL_ENABLE_CUDA=OFF ### Quick start ```python +import torch import tmol +# Pick your device: "cpu", "cuda", or "mps" (Apple Silicon) +device = torch.device("mps" if torch.backends.mps.is_available() else + "cuda" if torch.cuda.is_available() else "cpu") + # Load a structure -pose_stack = tmol.pose_stack_from_pdb("1ubq.pdb") +pose_stack = tmol.pose_stack_from_pdb("1ubq.pdb", device=device) # Score it -sfxn = tmol.beta2016_score_function(pose_stack.device) +sfxn = tmol.beta2016_score_function(device) scorer = sfxn.render_whole_pose_scoring_module(pose_stack) print(scorer(pose_stack.coords)) ``` @@ -294,7 +355,7 @@ xyz = tmol.pose_stack_to_rosettafold2(...) ``` > [!NOTE] -> Tested on Ubuntu 20.04. Other platforms should work but are not yet verified. +> Tested on Ubuntu 20.04 (CUDA) and macOS 14+ (MPS). Other platforms should work but are not yet verified. > [!WARNING] > Call `torch.set_grad_enabled(True)` before using the tmol minimizer, since RF2 disables gradients during inference by default. diff --git a/pyproject.toml b/pyproject.toml index 20a4985ec..74fb5805a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -100,9 +100,15 @@ cuda = [ # CCCL headers (nv/target, cub/, thrust/) needed by nvcc "nvidia-cuda-cccl>=13", ] +# MPS (Apple Metal) backend — macOS + Apple Silicon only. +# No pip dependencies: Metal SDK ships with macOS 13+ / Xcode 14+. +# Requires: macOS >= 13.0, Apple Silicon or AMD GPU with Metal 3 support, +# PyTorch >= 2.0 (MPS backend), Xcode Command Line Tools. +mps = [] [project.urls] -repository = "https://github.com/uw-ipd/tmol" +repository = "https://github.com/fnachon/tmol" +"Upstream (CUDA/Linux)" = "https://github.com/uw-ipd/tmol" # Build settings --------------------------------------------------------------- [build-system] diff --git a/tmol/_cpp_lib.py b/tmol/_cpp_lib.py index 0fe6606d1..b037410ab 100644 --- a/tmol/_cpp_lib.py +++ b/tmol/_cpp_lib.py @@ -97,13 +97,17 @@ def extension_load_error_details(exc: OSError) -> str: def _find_extension_library() -> str | None: - """Locate the _C shared library in tmol's package directory.""" - try: - spec = importlib.util.find_spec("tmol._C") - if spec is not None and spec.origin is not None: - return spec.origin - except (ModuleNotFoundError, ValueError): - pass + """Locate the _C shared library path without loading it.""" + import glob + import os + + # Search by filename pattern — avoids triggering importlib to load the .so + # (which would register TORCH_LIBRARY ops and cause a double-registration + # crash if load_library is later called for the same file). + package_dir = os.path.dirname(__file__) + matches = glob.glob(os.path.join(package_dir, "_C.*.so")) + if matches: + return matches[0] return None @@ -118,6 +122,16 @@ def _ensure_loaded() -> None: if _loaded: return + # Check whether the .so is already in the process (loaded by Python's own + # import machinery, e.g. via an editable-install .pth that puts the package + # directory on sys.path and a subsequent `import tmol._C`). + # In that case all TORCH_LIBRARY ops are already registered; calling + # load_library again causes a "Key already registered" C++ abort. + if "tmol._C" in __import__("sys").modules: + _loaded = True + logger.debug("tmol._C already in sys.modules; skipping load_library") + return + lib_path = _find_extension_library() if lib_path is None: raise TmolExtensionNotBuiltError() diff --git a/tmol/io/details/compiled/compiled.ops.cpp b/tmol/io/details/compiled/compiled.ops.cpp index 1a81594c2..7fe0eb527 100644 --- a/tmol/io/details/compiled/compiled.ops.cpp +++ b/tmol/io/details/compiled/compiled.ops.cpp @@ -42,6 +42,9 @@ class PoseLeafAtomGen : public torch::autograd::Function { Tensor block_type_atom_icoors_backup) { at::Tensor new_coords; + // MPS round-trip: TPack allocates on CPU for MPS; move output back after. + c10::Device orig_device = orig_coords.device(); + using Int = int32_t; TMOL_DISPATCH_FLOATING_DEVICE( @@ -72,6 +75,11 @@ class PoseLeafAtomGen : public torch::autograd::Function { new_coords = result.tensor; })); + // Move result back to original device (MPS→CPU via TPack, then back to MPS) + if (orig_device.is_mps()) { + new_coords = new_coords.to(orig_device); + } + ctx->save_for_backward( {orig_coords, new_coords, @@ -111,6 +119,7 @@ class PoseLeafAtomGen : public torch::autograd::Function { at::Tensor dE_d_orig_coords; + c10::Device orig_device = orig_coords.device(); using Int = int32_t; auto dE_d_new_coords = grad_outputs[0]; @@ -144,6 +153,10 @@ class PoseLeafAtomGen : public torch::autograd::Function { dE_d_orig_coords = result.tensor; })); + if (orig_device.is_mps()) { + dE_d_orig_coords = dE_d_orig_coords.to(orig_device); + } + return { dE_d_orig_coords, torch::Tensor(), @@ -200,6 +213,7 @@ Tensor resolve_his_tautomerization( Tensor his_atom_inds, Tensor his_remapping_dst_index) { at::Tensor his_taut; + c10::Device orig_device = coords.device(); TMOL_DISPATCH_FLOATING_DEVICE(coords.options(), "resolve_his_taut", ([&] { using Real = scalar_t; constexpr tmol::Device Dev = device_t; @@ -221,6 +235,9 @@ Tensor resolve_his_tautomerization( his_taut = result.tensor; })); + if (orig_device.is_mps()) { + his_taut = his_taut.to(orig_device); + } return his_taut; } diff --git a/tmol/io/details/compiled/gen_pose_leaf_atoms.mps.mm b/tmol/io/details/compiled/gen_pose_leaf_atoms.mps.mm new file mode 100644 index 000000000..61211e263 --- /dev/null +++ b/tmol/io/details/compiled/gen_pose_leaf_atoms.mps.mm @@ -0,0 +1,23 @@ +#include +#include + +namespace tmol { +namespace io { +namespace details { +namespace compiled { + +template struct GeneratePoseLeafAtoms< + score::common::DeviceOperations, + tmol::Device::MPS, + float, + int>; +template struct GeneratePoseLeafAtoms< + score::common::DeviceOperations, + tmol::Device::MPS, + double, + int>; + +} // namespace compiled +} // namespace details +} // namespace io +} // namespace tmol diff --git a/tmol/io/details/compiled/resolve_his_taut.mps.mm b/tmol/io/details/compiled/resolve_his_taut.mps.mm new file mode 100644 index 000000000..f26b32822 --- /dev/null +++ b/tmol/io/details/compiled/resolve_his_taut.mps.mm @@ -0,0 +1,23 @@ +#include +#include + +namespace tmol { +namespace io { +namespace details { +namespace compiled { + +template struct ResolveHisTaut< + score::common::DeviceOperations, + tmol::Device::MPS, + float, + int>; +template struct ResolveHisTaut< + score::common::DeviceOperations, + tmol::Device::MPS, + double, + int>; + +} // namespace compiled +} // namespace details +} // namespace io +} // namespace tmol diff --git a/tmol/io/details/his_taut_resolution.py b/tmol/io/details/his_taut_resolution.py index bcafccef8..c7355e1b2 100644 --- a/tmol/io/details/his_taut_resolution.py +++ b/tmol/io/details/his_taut_resolution.py @@ -48,11 +48,23 @@ def resolve_his_tautomerization( his_pose_ind, his_res_ind = torch.nonzero( res_types == his_inds.his_co_aa_ind, as_tuple=True ) + orig_device = coords.device + # The compiled kernel modifies res_type_variants and his_remapping_dst_index + # in-place. For MPS, TCAST would create CPU copies so those modifications + # would be lost. Run everything on CPU and move results back. + if orig_device.type == "mps": + coords = coords.cpu() + res_types = res_types.cpu() + res_type_variants = res_type_variants.cpu() + his_pose_ind = his_pose_ind.cpu() + his_res_ind = his_res_ind.cpu() + atom_is_present = atom_is_present.cpu() + his_remapping_dst_index = torch.tile( torch.arange( canonical_ordering.max_n_canonical_atoms, dtype=torch.int64, - device=res_types.device, + device=coords.device, ), (res_types.shape[0], res_types.shape[1], 1), ).reshape( @@ -77,10 +89,10 @@ def resolve_his_tautomerization( ) return ( - his_taut.to(dtype=torch.int32, device=coords.device), - res_type_variants, - resolved_coords, - resolved_atom_is_present, + his_taut.to(dtype=torch.int32, device=orig_device), + res_type_variants.to(orig_device), + resolved_coords.to(orig_device), + resolved_atom_is_present.to(orig_device), ) diff --git a/tmol/kinematics/compiled/__init__.py b/tmol/kinematics/compiled/__init__.py index 027df2519..c57c74d85 100644 --- a/tmol/kinematics/compiled/__init__.py +++ b/tmol/kinematics/compiled/__init__.py @@ -1,3 +1,26 @@ -from .compiled_ops import forward_kin_op, inverse_kin +import torch + +from .compiled_ops import forward_kin_op, inverse_kin as _inverse_kin_dispatch __all__ = ["forward_kin_op", "inverse_kin"] + + +def inverse_kin(*args, **kwargs): + # inverse_kin has no compiled MPS kernel; run on CPU and move the result back. + any_mps = any(isinstance(a, torch.Tensor) and a.device.type == "mps" for a in args) + if not any_mps: + return _inverse_kin_dispatch(*args, **kwargs) + + cpu_args = tuple(a.to("cpu") if isinstance(a, torch.Tensor) else a for a in args) + cpu_kwargs = { + k: v.to("cpu") if isinstance(v, torch.Tensor) else v for k, v in kwargs.items() + } + dtype = cpu_args[0].dtype + result = _inverse_kin_dispatch(*cpu_args, **cpu_kwargs) + # float64 cannot live on MPS — return on CPU; float32 can be moved back. + if dtype == torch.float64: + return result + mps_device = next( + a.device for a in args if isinstance(a, torch.Tensor) and a.device.type == "mps" + ) + return result.to(mps_device) diff --git a/tmol/kinematics/compiled/compiled.mps.mm b/tmol/kinematics/compiled/compiled.mps.mm new file mode 100644 index 000000000..0cae9df26 --- /dev/null +++ b/tmol/kinematics/compiled/compiled.mps.mm @@ -0,0 +1,233 @@ +// compiled.mps.mm — MPS instantiation of kinematic dispatch structs. +// +// The ForwardKinDispatch / InverseKinDispatch / KinDerivDispatch template +// structs are pure C++ (no CUDA intrinsics) and work for any tmol::Device D +// via explicit loops. We copy their definition here (parameterised over D) +// and instantiate for Device::MPS so that compiled.cpp/.ops.cpp can resolve +// MPS calls through TORCH_LIBRARY dispatch. + +#include + +#include +#include + +#include "common.hh" +#include "params.hh" +#include "compiled.impl.hh" + +namespace tmol { +namespace kinematics { + +template +using Vec = Eigen::Matrix; + +#define HomogeneousTransform Eigen::Matrix +#define KintreeDof Eigen::Matrix +#define Coord Eigen::Matrix + +template +struct ForwardKinDispatch { + static auto f( + ContextManager&, + TView dofs, + TView nodes, + TView scans, + TView, 1, tmol::Device::CPU> gens, + TView, 1, D> kintree) + -> std::tuple, TPack > { + auto num_atoms = dofs.size(0); + + auto HTs_t = TPack::empty({num_atoms}); + auto HTs = HTs_t.view; + auto xs_t = TPack::empty({num_atoms}); + auto xs = xs_t.view; + + auto k_dof2ht = ([=](int i) { + DOFtype doftype = (DOFtype)kintree[i].doftype; + if (doftype == ROOT) { + HTs[i] = HomogeneousTransform::Identity(); + } else if (doftype == JUMP) { + HTs[i] = common::jumpTransform(dofs[i]); + } else if (doftype == BOND) { + HTs[i] = common::bondTransform(dofs[i]); + } + }); + + for (int i = 0; i < num_atoms; i++) { + k_dof2ht(i); + } + + auto k_compose = + ([=](int p, int i) { HTs[i] = HTs[i] * HTs[p]; }); + + int ngens = gens.size(0) - 1; + for (int gen = 0; gen < ngens; gen++) { + int scanstart = gens[gen].scan_start; + int scanstop = gens[gen + 1].scan_start; + for (int j = scanstart; j < scanstop; j++) { + int nodestart = gens[gen].node_start + scans[j]; + int nodestop = (j == scanstop - 1) + ? gens[gen + 1].node_start + : (gens[gen].node_start + scans[j + 1]); + for (int k = nodestart; k < nodestop - 1; k++) { + k_compose(nodes[k], nodes[k + 1]); + } + } + } + + auto k_getcoords = ([=](int i) { + xs[i] = HTs[i].block(3, 0, 1, 3).transpose(); + }); + + for (int i = 0; i < num_atoms; i++) { + k_getcoords(i); + } + + return {xs_t, HTs_t}; + } +}; + +template +struct InverseKinDispatch { + static auto f( + ContextManager&, + TView coords, + TView parent, + TView frame_x, + TView frame_y, + TView frame_z, + TView doftype) -> TPack { + auto num_atoms = coords.size(0); + + auto HTs_t = TPack::empty({num_atoms}); + auto HTs = HTs_t.view; + auto dofs_t = TPack::empty({num_atoms}); + auto dofs = dofs_t.view; + + auto k_coords2hts = ([=](int i) { + if (i == 0) { + HTs[i] = HomogeneousTransform::Identity(); + } else { + HTs[i] = common::hts_from_frames( + coords[i], + coords[frame_x[i]], + coords[frame_y[i]], + coords[frame_z[i]]); + } + }); + + for (int i = 0; i < num_atoms; i++) { + k_coords2hts(i); + } + + auto k_hts2dofs = ([=](int i) { + HomogeneousTransform lclHT; + if (doftype[i] == ROOT) { + dofs[i] = KintreeDof::Constant(0); + } else { + lclHT = HTs[i] * common::ht_inv(HTs[parent[i]]); + if (doftype[i] == JUMP) { + dofs[i] = common::invJumpTransform(lclHT); + } else if (doftype[i] == BOND) { + dofs[i] = common::invBondTransform(lclHT); + } + } + }); + + for (int i = 0; i < num_atoms; i++) { + k_hts2dofs(i); + } + + return dofs_t; + } +}; + +template +struct KinDerivDispatch { + static auto f( + ContextManager&, + TView dVdx, + TView hts, + TView dofs, + TView nodes, + TView scans, + TView, 1, tmol::Device::CPU> gens, + TView, 1, D> kintree) -> TPack { + auto num_atoms = dVdx.size(0); + + auto f1f2s_t = TPack, 1, D>::empty({num_atoms}); + auto f1f2s = f1f2s_t.view; + auto dsc_ddofs_t = TPack::empty({num_atoms}); + auto dsc_ddofs = dsc_ddofs_t.view; + + auto k_f1f2s = ([=](int i) { + Coord trans = hts[i].block(3, 0, 1, 3).transpose(); + Coord f1 = dVdx[i].isZero(0) ? dVdx[i] + : trans.cross(trans - dVdx[i]).transpose(); + f1f2s[i].topRows(3) = f1; + f1f2s[i].bottomRows(3) = dVdx[i]; + }); + + for (int i = 0; i < num_atoms; i++) { + k_f1f2s(i); + } + + auto k_compose = ([=](int p, int i) { + f1f2s[i] = f1f2s[i] + f1f2s[p]; + }); + + int ngens = gens.size(0) - 1; + for (int gen = 0; gen < ngens; gen++) { + int scanstart = gens[gen].scan_start; + int scanstop = gens[gen + 1].scan_start; + for (int j = scanstart; j < scanstop; j++) { + int nodestart = gens[gen].node_start + scans[j]; + int nodestop = (j == scanstop - 1) + ? gens[gen + 1].node_start + : (gens[gen].node_start + scans[j + 1]); + for (int k = nodestart; k < nodestop - 1; k++) { + k_compose(nodes[k], nodes[k + 1]); + } + } + } + + auto k_f1f2s2derivs = ([=](int i) { + Vec f1 = f1f2s[i].topRows(3); + Vec f2 = f1f2s[i].bottomRows(3); + if (kintree[i].doftype == ROOT) { + dsc_ddofs[i] = Vec::Constant(0); + } else if (kintree[i].doftype == JUMP) { + dsc_ddofs[i] = common::jumpDerivatives( + dofs[i], hts[i], hts[kintree[i].parent], f1, f2); + } else if (kintree[i].doftype == BOND) { + dsc_ddofs[i] = common::bondDerivatives( + dofs[i], hts[i], hts[kintree[i].parent], f1, f2); + } + }); + + for (int i = 0; i < num_atoms; i++) { + k_f1f2s2derivs(i); + } + + return dsc_ddofs_t; + } +}; + +template struct ForwardKinDispatch; +template struct ForwardKinDispatch; +template struct InverseKinDispatch; +template struct InverseKinDispatch; +template struct KinDerivDispatch; +template struct KinDerivDispatch; + +template struct KinForestFromStencil< + tmol::score::common::DeviceOperations, + tmol::Device::MPS, + int32_t>; + +#undef HomogeneousTransform +#undef KintreeDof +#undef Coord + +} // namespace kinematics +} // namespace tmol diff --git a/tmol/kinematics/compiled/compiled_ops.cpp b/tmol/kinematics/compiled/compiled_ops.cpp index ecae88946..1b9136974 100644 --- a/tmol/kinematics/compiled/compiled_ops.cpp +++ b/tmol/kinematics/compiled/compiled_ops.cpp @@ -25,6 +25,11 @@ using torch::autograd::AutogradContext; using torch::autograd::Function; using torch::autograd::tensor_list; +// MPS round-trip helper: TPack allocates on CPU for MPS inputs; move back. +static inline at::Tensor mps_to_dev(at::Tensor t, c10::Device dev) { + return dev.is_mps() ? t.to(dev) : t; +} + class KinematicOp : public torch::autograd::Function { public: static Tensor forward( @@ -40,6 +45,7 @@ class KinematicOp : public torch::autograd::Function { at::Tensor coords; at::Tensor HTs; + c10::Device orig_device = dofs.device(); using Int = int32_t; TMOL_DISPATCH_FLOATING_DEVICE(dofs.options(), "forward_kin_op", ([&] { @@ -59,6 +65,8 @@ class KinematicOp : public torch::autograd::Function { HTs = std::get<1>(result).tensor; })); + coords = mps_to_dev(coords, orig_device); + HTs = mps_to_dev(HTs, orig_device); ctx->save_for_backward({HTs, dofs, nodes_b, scans_b, gens_b, kintree}); return coords; @@ -75,6 +83,7 @@ class KinematicOp : public torch::autograd::Function { auto kintree = saved[i++]; at::Tensor dV_ddof; + c10::Device orig_device = dofs.device(); using Int = int32_t; auto dVdx = grad_outputs[0]; TMOL_DISPATCH_FLOATING_DEVICE(HTs.options(), "kin_deriv_op", ([&] { @@ -95,6 +104,8 @@ class KinematicOp : public torch::autograd::Function { dV_ddof = result.tensor; })); + dV_ddof = mps_to_dev(dV_ddof, orig_device); + return { dV_ddof, torch::Tensor(), @@ -130,6 +141,7 @@ Tensor forward_only_op( Tensor kintree) { at::Tensor coords; + c10::Device orig_device = dofs.device(); using Int = int32_t; TMOL_DISPATCH_FLOATING_DEVICE(dofs.options(), "forward_kin_only_op", ([&] { @@ -148,7 +160,7 @@ Tensor forward_only_op( coords = std::get<0>(result).tensor; })); - return coords; + return mps_to_dev(coords, orig_device); }; auto get_kfo_indices_for_atoms( @@ -159,12 +171,12 @@ auto get_kfo_indices_for_atoms( at::Tensor block_kfo_offset_tp; at::Tensor kfo_2_orig_mapping_tp; at::Tensor atom_kfo_index; + c10::Device dev = pose_stack_block_coord_offset.device(); TMOL_DISPATCH_INDEX_DEVICE( pose_stack_block_coord_offset.options(), "get_kfo_indices_for_atoms", ([&] { - using Int = int32_t; // ONLY 32-bit integers supported! No atomicAdd - // for signed 64-bit integers in CUDA + using Int = int32_t; constexpr tmol::Device Dev = device_t; auto result = @@ -179,7 +191,10 @@ auto get_kfo_indices_for_atoms( kfo_2_orig_mapping_tp = std::get<1>(result).tensor; atom_kfo_index = std::get<2>(result).tensor; })); - return {block_kfo_offset_tp, kfo_2_orig_mapping_tp, atom_kfo_index}; + return { + mps_to_dev(block_kfo_offset_tp, dev), + mps_to_dev(kfo_2_orig_mapping_tp, dev), + mps_to_dev(atom_kfo_index, dev)}; } auto get_kfo_atom_parents( @@ -195,10 +210,10 @@ auto get_kfo_atom_parents( Tensor block_type_conn_atom) -> tensor_list { at::Tensor kfo_parent_atoms; at::Tensor kfo_grandparent_atoms; + c10::Device dev = pose_stack_block_type.device(); TMOL_DISPATCH_INDEX_DEVICE( pose_stack_block_type.options(), "get_kfo_atom_parents", ([&] { - using Int = int32_t; // ONLY 32-bit integers supported! No atomicAdd - // for signed 64-bit integers in CUDA + using Int = int32_t; constexpr tmol::Device Dev = device_t; auto result = @@ -208,7 +223,6 @@ auto get_kfo_atom_parents( TCAST(pose_stack_block_type), TCAST(pose_stack_inter_residue_connections), TCAST(pose_stack_ff_parent), - // TCAST(pose_stack_ff_conn_to_parent), TCAST(pose_stack_block_in_and_first_out), TCAST(block_type_parents), TCAST(kfo_2_orig_mapping), @@ -220,7 +234,7 @@ auto get_kfo_atom_parents( kfo_parent_atoms = std::get<0>(result).tensor; kfo_grandparent_atoms = std::get<1>(result).tensor; })); - return {kfo_parent_atoms, kfo_grandparent_atoms}; + return {mps_to_dev(kfo_parent_atoms, dev), mps_to_dev(kfo_grandparent_atoms, dev)}; } auto get_children( @@ -235,10 +249,10 @@ auto get_children( at::Tensor child_list; at::Tensor is_atom_jump; + c10::Device dev = pose_stack_block_type.device(); TMOL_DISPATCH_INDEX_DEVICE( pose_stack_block_type.options(), "get_children", ([&] { - using Int = int32_t; // ONLY 32-bit integers supported! No atomicAdd - // for signed 64-bit integers in CUDA + using Int = int32_t; constexpr tmol::Device Dev = device_t; auto result = @@ -256,7 +270,11 @@ auto get_children( child_list = std::get<2>(result).tensor; is_atom_jump = std::get<3>(result).tensor; })); - return {n_children, child_list_span, child_list, is_atom_jump}; + return { + mps_to_dev(n_children, dev), + mps_to_dev(child_list_span, dev), + mps_to_dev(child_list, dev), + mps_to_dev(is_atom_jump, dev)}; } auto get_id_and_frame_xyz( @@ -274,10 +292,10 @@ auto get_id_and_frame_xyz( at::Tensor frame_z; at::Tensor keep_dof_fixed; + c10::Device dev = parents.device(); TMOL_DISPATCH_INDEX_DEVICE( parents.options(), "get_id_and_frame_xyz", ([&] { - using Int = int32_t; // ONLY 32-bit integers supported! No atomicAdd - // for signed 64-bit integers in CUDA + using Int = int32_t; constexpr tmol::Device Dev = device_t; auto result = @@ -298,7 +316,12 @@ auto get_id_and_frame_xyz( frame_z = std::get<3>(result).tensor; keep_dof_fixed = std::get<4>(result).tensor; })); - return {id, frame_x, frame_y, frame_z, keep_dof_fixed}; + return { + mps_to_dev(id, dev), + mps_to_dev(frame_x, dev), + mps_to_dev(frame_y, dev), + mps_to_dev(frame_z, dev), + mps_to_dev(keep_dof_fixed, dev)}; } auto calculate_ff_edge_delays( @@ -320,10 +343,10 @@ auto calculate_ff_edge_delays( Tensor first_child_of_ff_edge; Tensor delay_for_edge; Tensor toposort_index_for_edge; + c10::Device dev = pose_stack_block_type.device(); TMOL_DISPATCH_INDEX_DEVICE( pose_stack_block_type.options(), "calculate_ff_edge_delays", ([&] { - using Int = int32_t; // ONLY 32-bit integers supported! No atomicAdd - // for signed 64-bit integers in CUDA + using Int = int32_t; constexpr tmol::Device Dev = device_t; auto result = @@ -347,15 +370,15 @@ auto calculate_ff_edge_delays( toposort_index_for_edge = std::get<8>(result).tensor; })); return { - dfs_order_of_ff_edges, - n_ff_edges, - ff_edge_parent, - first_ff_edge_for_block_cpu, - pose_stack_ff_parent, - max_gen_depth_of_ff_edge, - first_child_of_ff_edge, - delay_for_edge, - toposort_index_for_edge}; + mps_to_dev(dfs_order_of_ff_edges, dev), + mps_to_dev(n_ff_edges, dev), + mps_to_dev(ff_edge_parent, dev), + mps_to_dev(first_ff_edge_for_block_cpu, dev), + mps_to_dev(pose_stack_ff_parent, dev), + mps_to_dev(max_gen_depth_of_ff_edge, dev), + mps_to_dev(first_child_of_ff_edge, dev), + mps_to_dev(delay_for_edge, dev), + mps_to_dev(toposort_index_for_edge, dev)}; } auto get_jump_atom_indices( @@ -365,10 +388,10 @@ auto get_jump_atom_indices( ) -> tensor_list { Tensor pose_stack_atom_for_jump; Tensor pose_stack_atom_for_root_jump; + c10::Device dev = pose_stack_block_type.device(); TMOL_DISPATCH_INDEX_DEVICE( pose_stack_block_type.options(), "calculate_ff_edge_delays", ([&] { - using Int = int32_t; // ONLY 32-bit integers supported! No atomicAdd - // for signed 64-bit integers in CUDA + using Int = int32_t; constexpr tmol::Device Dev = device_t; auto result = @@ -381,7 +404,9 @@ auto get_jump_atom_indices( pose_stack_atom_for_jump = std::get<0>(result).tensor; pose_stack_atom_for_root_jump = std::get<1>(result).tensor; })); - return {pose_stack_atom_for_jump, pose_stack_atom_for_root_jump}; + return { + mps_to_dev(pose_stack_atom_for_jump, dev), + mps_to_dev(pose_stack_atom_for_root_jump, dev)}; } auto get_block_parent_connectivity_from_toposort( @@ -398,10 +423,10 @@ auto get_block_parent_connectivity_from_toposort( Tensor block_type_n_conn, // T Tensor block_type_polymeric_conn_index) -> Tensor { Tensor pose_stack_block_in_and_first_out; + c10::Device dev = pose_stack_block_type.device(); TMOL_DISPATCH_INDEX_DEVICE( pose_stack_block_type.options(), "calculate_ff_edge_delays", ([&] { - using Int = int32_t; // ONLY 32-bit integers supported! No atomicAdd - // for signed 64-bit integers in CUDA + using Int = int32_t; constexpr tmol::Device Dev = device_t; auto result = @@ -413,18 +438,17 @@ auto get_block_parent_connectivity_from_toposort( pose_stack_inter_residue_connections), // P x L x C x 2 TCAST(pose_stack_ff_parent), TCAST(dfs_order_of_ff_edges), - TCAST(n_ff_edges), // P - TCAST(ff_edges), // P x E x 4 - TCAST(first_ff_edge_for_block), // P x L - // TCAST(max_n_gens_for_ff_edge), // P x E - TCAST(first_child_of_ff_edge), // P x E - TCAST(delay_for_edge), // P x E - TCAST(topo_sort_index_for_edge), // (P*E) - TCAST(block_type_n_conn), // T + TCAST(n_ff_edges), + TCAST(ff_edges), + TCAST(first_ff_edge_for_block), + TCAST(first_child_of_ff_edge), + TCAST(delay_for_edge), + TCAST(topo_sort_index_for_edge), + TCAST(block_type_n_conn), TCAST(block_type_polymeric_conn_index)); pose_stack_block_in_and_first_out = result.tensor; })); - return pose_stack_block_in_and_first_out; + return mps_to_dev(pose_stack_block_in_and_first_out, dev); } auto get_scans2( @@ -462,10 +486,10 @@ auto get_scans2( Tensor nodes_bw; Tensor scans_bw; Tensor gens_bw; + c10::Device dev = pose_stack_block_type.device(); TMOL_DISPATCH_INDEX_DEVICE( pose_stack_block_type.options(), "calculate_ff_edge_delays", ([&] { - using Int = int32_t; // ONLY 32-bit integers supported! No atomicAdd - // for signed 64-bit integers in CUDA + using Int = int32_t; constexpr tmol::Device Dev = device_t; auto result = @@ -504,7 +528,13 @@ auto get_scans2( scans_bw = std::get<4>(result).tensor; gens_bw = std::get<5>(result).tensor; })); - return {nodes_fw, scans_fw, gens_fw, nodes_bw, scans_bw, gens_bw}; + return { + mps_to_dev(nodes_fw, dev), + mps_to_dev(scans_fw, dev), + mps_to_dev(gens_fw, dev), + mps_to_dev(nodes_bw, dev), + mps_to_dev(scans_bw, dev), + mps_to_dev(gens_bw, dev)}; } auto minimizer_map_from_movemap( @@ -552,6 +582,7 @@ auto minimizer_map_from_movemap( Tensor move_atom_dof_mask) -> Tensor { // Minimizer map: a boolean vector of the DOFs that are free Tensor minimizer_map; + c10::Device dev = pose_stack_block_type.device(); TMOL_DISPATCH_INDEX_DEVICE( pose_stack_block_type.options(), "minimizer_map_from_movemap", ([&] { using Int = int32_t; @@ -605,7 +636,7 @@ auto minimizer_map_from_movemap( TCAST(move_atom_dof_mask)); minimizer_map = result.tensor; })); - return minimizer_map; + return mps_to_dev(minimizer_map, dev); } auto inv_kin_dispatch( diff --git a/tmol/kinematics/script_modules.py b/tmol/kinematics/script_modules.py index bf77dff14..85e95cdd5 100644 --- a/tmol/kinematics/script_modules.py +++ b/tmol/kinematics/script_modules.py @@ -40,6 +40,8 @@ def __init__(self, pose_stack: PoseStack, fold_forest: FoldForest): pbt = pose_stack.packed_block_types ff = fold_forest device = pose_stack.device + # Kinematics needs float64, which MPS does not support; use CPU for MPS + kin_device = torch.device("cpu") if device.type == "mps" else device # Setup: initial annotations of block types and packed block types # with the per-block-scan-path segments. @@ -75,14 +77,14 @@ def _tint(ts): ] ), dim=1, - ).to(device) + ).to(kin_device) ) - self.nodes_f = _p(kmd.scan_data_fw.nodes.to(device)) - self.scans_f = _p(kmd.scan_data_fw.scans.to(device)) + self.nodes_f = _p(kmd.scan_data_fw.nodes.to(kin_device)) + self.scans_f = _p(kmd.scan_data_fw.scans.to(kin_device)) self.gens_f = _p(kmd.scan_data_fw.gens) # on cpu - self.nodes_b = _p(kmd.scan_data_bw.nodes.to(device)) - self.scans_b = _p(kmd.scan_data_bw.scans.to(device)) + self.nodes_b = _p(kmd.scan_data_bw.nodes.to(kin_device)) + self.scans_b = _p(kmd.scan_data_bw.scans.to(kin_device)) self.gens_b = _p(kmd.scan_data_bw.gens) # on cpu def forward(self, dofs): diff --git a/tmol/pack/compiled/compiled.cpu.cpp b/tmol/pack/compiled/compiled.cpu.cpp index 86072c820..7abed7e28 100644 --- a/tmol/pack/compiled/compiled.cpu.cpp +++ b/tmol/pack/compiled/compiled.cpu.cpp @@ -92,7 +92,8 @@ auto AnnealerDispatch::forward( int const n_inner_iterations = n_inner_iterations_factor * pose_n_rotamers; for (int traj = 0; traj < n_traj; ++traj) { - // Initial assignment: assign a rotamer to every residue + // Initial assignment: assign a rotamer to every residue. + // Residues with 0 rotamers get -1 (avoids rand() % 0 undefined behavior). for (int i = 0; i < n_res; ++i) { int const i_n_rots = n_rotamers_for_res[pose][i]; if (i_n_rots == 0) { diff --git a/tmol/pack/compiled/compiled.mps.mm b/tmol/pack/compiled/compiled.mps.mm new file mode 100644 index 000000000..2043f3e05 --- /dev/null +++ b/tmol/pack/compiled/compiled.mps.mm @@ -0,0 +1,269 @@ +// compiled.mps.mm — MPS instantiation of pack/compiled dispatch structs. +// +// AnnealerDispatch and InteractionGraphBuilder contain pure C++ logic with +// no CUDA intrinsics. We include the shared impl header and instantiate +// for Device::MPS. + +#include +#include + +#include + +#include "simulated_annealing.hh" +#include "compiled.impl.hh" + +#include + +namespace tmol { +namespace pack { +namespace compiled { + +// AnnealerDispatch for MPS reuses the same CPU SA implementation. +// The SA algorithm is sequential by nature and runs on the host CPU; tensor +// data is accessible via unified memory on Apple Silicon. + +template +void set_quench_order( + TView quench_order, + int const n_rots, + int const pose_rotamer_offset) { + for (int i = 0; i < n_rots; ++i) { + quench_order[i] = i + pose_rotamer_offset; + } + for (int i = 0; i <= n_rots - 2; ++i) { + int j = i + rand() % (n_rots - i); + int jval = quench_order[j]; + quench_order[j] = quench_order[i]; + quench_order[i] = jval; + } +} + +template +auto AnnealerDispatch::forward( + ContextManager&, + int max_n_rotamers_per_pose, + TView pose_n_res, + TView n_rotamers_for_pose, + TView rotamer_offset_for_pose, + TView n_rotamers_for_res, + TView oneb_offsets, + TView res_for_rot, + int32_t chunk_size, + TView chunk_offset_offsets, + TView chunk_offsets, + TView energy1b, + TView energy2b) + -> std::tuple, TPack > { + clock_t start = clock(); + + int const n_poses = pose_n_res.size(0); + int const max_n_res = n_rotamers_for_res.size(1); + int const n_rotamers = res_for_rot.size(0); + int const n_traj = 1; + int const n_outer_iterations = 20; + int const n_inner_iterations_factor = 20; + + auto scores_t = TPack::zeros({n_poses, n_traj}); + auto current_rotamer_assignments_t = + TPack::zeros({n_poses, n_traj, max_n_res}); + auto best_rotamer_assignments_t = + TPack::zeros({n_poses, n_traj, max_n_res}); + auto quench_order_t = TPack::zeros({n_rotamers}); + + auto scores = scores_t.view; + auto current_rotamer_assignments = current_rotamer_assignments_t.view; + auto best_rotamer_assignments = best_rotamer_assignments_t.view; + auto quench_order = quench_order_t.view; + + float const high_temp = 100; + float const low_temp = 0.2; + + for (int pose = 0; pose < n_poses; ++pose) { + int const n_res = pose_n_res[pose]; + int const pose_n_rotamers = n_rotamers_for_pose[pose]; + int const pose_rotamer_offset = rotamer_offset_for_pose[pose]; + int const n_inner_iterations = n_inner_iterations_factor * pose_n_rotamers; + + for (int traj = 0; traj < n_traj; ++traj) { + // Residues with 0 rotamers (padding for shorter poses) get -1. + for (int i = 0; i < max_n_res; ++i) { + int const i_n_rots = n_rotamers_for_res[pose][i]; + if (i_n_rots == 0) { + current_rotamer_assignments[pose][traj][i] = -1; + best_rotamer_assignments[pose][traj][i] = -1; + } else { + int rand_rot = rand() % i_n_rots; + current_rotamer_assignments[pose][traj][i] = rand_rot; + best_rotamer_assignments[pose][traj][i] = rand_rot; + } + } + + float temperature = high_temp; + double best_energy = total_energy_for_assignment( + n_rotamers_for_res[pose], + oneb_offsets[pose], + chunk_size, + chunk_offset_offsets[pose], + chunk_offsets, + energy1b, + energy2b, + current_rotamer_assignments[pose][traj]); + double current_total_energy = best_energy; + int naccepts = 0; + + for (int i = 0; i < n_outer_iterations; ++i) { + bool quench = false; + if (i == n_outer_iterations - 1) { + quench = true; + temperature = 0; + for (int j = 0; j < n_res; ++j) { + current_rotamer_assignments[pose][traj][j] = + best_rotamer_assignments[pose][traj][j]; + } + current_total_energy = total_energy_for_assignment( + n_rotamers_for_res[pose], + oneb_offsets[pose], + chunk_size, + chunk_offset_offsets[pose], + chunk_offsets, + energy1b, + energy2b, + current_rotamer_assignments[pose][traj]); + } + + for (int j = 0; j < n_inner_iterations; ++j) { + int global_ran_rot; + if (quench) { + if (j % pose_n_rotamers == 0) { + set_quench_order(quench_order, pose_n_rotamers, pose_rotamer_offset); + } + global_ran_rot = quench_order[j % pose_n_rotamers]; + } else { + global_ran_rot = rand() % pose_n_rotamers + pose_rotamer_offset; + } + + int const ran_res = res_for_rot[global_ran_rot]; + int const local_prev_rot = + current_rotamer_assignments[pose][traj][ran_res]; + int const ran_res_n_rots = n_rotamers_for_res[pose][ran_res]; + int const ran_res_n_chunks = (ran_res_n_rots - 1) / chunk_size + 1; + int const ran_res_offset = oneb_offsets[pose][ran_res]; + int const local_ran_rot = global_ran_rot - ran_res_offset; + int const ran_rot_chunk = local_ran_rot / chunk_size; + int const prev_rot_chunk = local_prev_rot / chunk_size; + int const ran_rot_in_chunk = + local_ran_rot - chunk_size * ran_rot_chunk; + int const prev_rot_in_chunk = + local_prev_rot - chunk_size * prev_rot_chunk; + int const ran_rot_chunk_size = + std::min(chunk_size, ran_res_n_rots - chunk_size * ran_rot_chunk); + int const prev_rot_chunk_size = std::min( + chunk_size, ran_res_n_rots - chunk_size * prev_rot_chunk); + int const global_prev_rot = local_prev_rot + ran_res_offset; + + double new_e = energy1b[global_ran_rot]; + double prev_e = energy1b[global_prev_rot]; + double deltaE = new_e - prev_e; + + for (int k = 0; k < n_res; ++k) { + if (k == ran_res) continue; + int64_t const k_ran_chunk_offset_offset = + chunk_offset_offsets[pose][k][ran_res]; + if (k_ran_chunk_offset_offset == -1) continue; + int const local_k_rot = + current_rotamer_assignments[pose][traj][k]; + int const k_n_rots = n_rotamers_for_res[pose][k]; + int const kres_n_chunks = (k_n_rots - 1) / chunk_size + 1; + int const krot_chunk = local_k_rot / chunk_size; + int const krot_in_chunk = local_k_rot - krot_chunk * chunk_size; + int64_t const krot_ranrot_chunk_offset = chunk_offsets + [k_ran_chunk_offset_offset + krot_chunk * ran_res_n_chunks + + ran_rot_chunk]; + int64_t const krot_prevrot_chunk_offset = chunk_offsets + [k_ran_chunk_offset_offset + krot_chunk * ran_res_n_chunks + + prev_rot_chunk]; + + double k_new_e = 0; + double k_prev_e = 0; + if (krot_ranrot_chunk_offset >= 0) { + k_new_e = energy2b + [krot_ranrot_chunk_offset + krot_in_chunk * ran_rot_chunk_size + + ran_rot_in_chunk]; + } + if (krot_prevrot_chunk_offset >= 0) { + k_prev_e = energy2b + [krot_prevrot_chunk_offset + + krot_in_chunk * prev_rot_chunk_size + prev_rot_in_chunk]; + } + deltaE += k_new_e - k_prev_e; + new_e += k_new_e; + prev_e += k_prev_e; + } + + float const uniform_random = float(rand()) / RAND_MAX; + if (pass_metropolis(temperature, uniform_random, deltaE, prev_e, quench)) { + current_rotamer_assignments[pose][traj][ran_res] = local_ran_rot; + current_total_energy += deltaE; + ++naccepts; + if (naccepts > 1000) { + naccepts = 0; + current_total_energy = total_energy_for_assignment( + n_rotamers_for_res[pose], + oneb_offsets[pose], + chunk_size, + chunk_offset_offsets[pose], + chunk_offsets, + energy1b, + energy2b, + current_rotamer_assignments[pose][traj]); + } + if (current_total_energy < best_energy) { + for (int k = 0; k < n_res; ++k) { + best_rotamer_assignments[pose][traj][k] = + current_rotamer_assignments[pose][traj][k]; + } + best_energy = current_total_energy; + } + } + } // inner loop + + temperature = + (high_temp - low_temp) * std::exp(-1 * (i + 1)) + low_temp; + } // outer loop + + scores[pose][traj] = total_energy_for_assignment( + n_rotamers_for_res[pose], + oneb_offsets[pose], + chunk_size, + chunk_offset_offsets[pose], + chunk_offsets, + energy1b, + energy2b, + best_rotamer_assignments[pose][traj]); + } // traj loop + } // pose loop + + clock_t stop = clock(); + std::cout << "MPS simulated annealing (CPU path) in " + << ((double)stop - start) / CLOCKS_PER_SEC << " seconds" + << std::endl; + + return {scores_t, best_rotamer_assignments_t}; +} + +template struct AnnealerDispatch; + +template struct InteractionGraphBuilder< + score::common::DeviceOperations, + tmol::Device::MPS, + float, + int64_t>; +template struct InteractionGraphBuilder< + score::common::DeviceOperations, + tmol::Device::MPS, + double, + int64_t>; + +} // namespace compiled +} // namespace pack +} // namespace tmol diff --git a/tmol/pack/compiled/compiled.ops.cpp b/tmol/pack/compiled/compiled.ops.cpp index ccad89ffc..219f433a7 100644 --- a/tmol/pack/compiled/compiled.ops.cpp +++ b/tmol/pack/compiled/compiled.ops.cpp @@ -23,6 +23,11 @@ namespace compiled { ContextManager mgr; using torch::Tensor; +// MPS round-trip: TPack allocates CPU for MPS inputs; move output back. +static inline at::Tensor mps_to_dev(at::Tensor t, c10::Device dev) { + return dev.is_mps() ? t.to(dev) : t; +} + std::vector build_interaction_graph( int64_t const bump_check, int64_t const chunk_size, @@ -55,6 +60,7 @@ std::vector build_interaction_graph( at::Tensor chunk_pair_offset; at::Tensor energy2b; + c10::Device orig_device = sparse_energies.device(); using Int = int64_t; TMOL_DISPATCH_FLOATING_DEVICE( @@ -98,6 +104,28 @@ std::vector build_interaction_graph( chunk_pair_offset = std::get<12>(result).tensor; energy2b = std::get<13>(result).tensor; })); + + // max_n_bump_checked_rotamers_per_pose is always CPU-resident; the rest + // may have been computed via a CPU round-trip for MPS inputs. + n_molten_blocks_per_pose = mps_to_dev(n_molten_blocks_per_pose, orig_device); + n_bc_rots_per_pose = mps_to_dev(n_bc_rots_per_pose, orig_device); + bc_rot_offset_for_pose = mps_to_dev(bc_rot_offset_for_pose, orig_device); + n_bc_rots_for_molten_block = + mps_to_dev(n_bc_rots_for_molten_block, orig_device); + bc_rot_offset_for_molten_block = + mps_to_dev(bc_rot_offset_for_molten_block, orig_device); + molten_block_ind_for_bc_rot = + mps_to_dev(molten_block_ind_for_bc_rot, orig_device); + rotamer_for_nonmolten_block = + mps_to_dev(rotamer_for_nonmolten_block, orig_device); + bc_rot_to_orig_rot = mps_to_dev(bc_rot_to_orig_rot, orig_device); + bg_bg_energies = mps_to_dev(bg_bg_energies, orig_device); + energy1b = mps_to_dev(energy1b, orig_device); + chunk_pair_offset_for_block_pair = + mps_to_dev(chunk_pair_offset_for_block_pair, orig_device); + chunk_pair_offset = mps_to_dev(chunk_pair_offset, orig_device); + energy2b = mps_to_dev(energy2b, orig_device); + std::vector result( {max_n_bump_checked_rotamers_per_pose, n_molten_blocks_per_pose, @@ -133,6 +161,7 @@ std::vector anneal( at::Tensor scores; at::Tensor rotamer_assignments; + c10::Device orig_device = energy1b.device(); TMOL_DISPATCH_FLOATING_DEVICE(energy1b.options(), "pack_anneal", ([&] { constexpr tmol::Device Dev = device_t; @@ -155,6 +184,9 @@ std::vector anneal( std::get<1>(result).tensor; })); + scores = mps_to_dev(scores, orig_device); + rotamer_assignments = mps_to_dev(rotamer_assignments, orig_device); + std::vector result({scores, rotamer_assignments}); return result; } diff --git a/tmol/pack/rotamer/dunbrack/compiled.mps.mm b/tmol/pack/rotamer/dunbrack/compiled.mps.mm new file mode 100644 index 000000000..1d4179f8b --- /dev/null +++ b/tmol/pack/rotamer/dunbrack/compiled.mps.mm @@ -0,0 +1,52 @@ +// compiled.mps.mm — MPS instantiation for Dunbrack rotamer sampler. +// +// ComplexDispatch in complex_dispatch.cpu.impl.hh is a primary template +// (device-agnostic CPU loops). We include it directly for MPS. + +#pragma once + +#include +#include + +#include +#include +#include + +#include +#include +#include + +#include + +#include "dispatch.impl.hh" + +namespace tmol { +namespace pack { +namespace rotamer { +namespace dunbrack { + +template struct DunbrackChiSampler< + score::common::ComplexDispatch, + tmol::Device::MPS, + float, + int32_t>; +template struct DunbrackChiSampler< + score::common::ComplexDispatch, + tmol::Device::MPS, + double, + int32_t>; +template struct DunbrackChiSampler< + score::common::ComplexDispatch, + tmol::Device::MPS, + float, + int64_t>; +template struct DunbrackChiSampler< + score::common::ComplexDispatch, + tmol::Device::MPS, + double, + int64_t>; + +} // namespace dunbrack +} // namespace rotamer +} // namespace pack +} // namespace tmol diff --git a/tmol/pack/rotamer/dunbrack/compiled.ops.cpp b/tmol/pack/rotamer/dunbrack/compiled.ops.cpp index ee973decb..f0493ae37 100644 --- a/tmol/pack/rotamer/dunbrack/compiled.ops.cpp +++ b/tmol/pack/rotamer/dunbrack/compiled.ops.cpp @@ -19,6 +19,11 @@ ContextManager mgr; using torch::Tensor; +// MPS round-trip: TPack allocates CPU for MPS inputs; move output back. +static inline at::Tensor mps_to_dev(at::Tensor t, c10::Device dev) { + return dev.is_mps() ? t.to(dev) : t; +} + template