pet-kokkos 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- pet_kokkos-0.1.0/.github/workflows/ci.yml +50 -0
- pet_kokkos-0.1.0/.github/workflows/wheels.yml +146 -0
- pet_kokkos-0.1.0/.gitignore +29 -0
- pet_kokkos-0.1.0/CMakeLists.txt +253 -0
- pet_kokkos-0.1.0/CMakePresets.json +70 -0
- pet_kokkos-0.1.0/LICENSE +29 -0
- pet_kokkos-0.1.0/PKG-INFO +407 -0
- pet_kokkos-0.1.0/README.md +362 -0
- pet_kokkos-0.1.0/apps/CMakeLists.txt +14 -0
- pet_kokkos-0.1.0/apps/batch_probe.cpp +150 -0
- pet_kokkos-0.1.0/apps/pet_eval.cpp +457 -0
- pet_kokkos-0.1.0/cmake/FindVesin.cmake +48 -0
- pet_kokkos-0.1.0/include/pet/calculator.hpp +242 -0
- pet_kokkos-0.1.0/include/pet/checkpoint.hpp +28 -0
- pet_kokkos-0.1.0/include/pet/config.hpp +68 -0
- pet_kokkos-0.1.0/include/pet/cutoff.hpp +46 -0
- pet_kokkos-0.1.0/include/pet/device_cell_list.hpp +63 -0
- pet_kokkos-0.1.0/include/pet/device_geometry.hpp +238 -0
- pet_kokkos-0.1.0/include/pet/device_neighbors.hpp +69 -0
- pet_kokkos-0.1.0/include/pet/exchange.hpp +43 -0
- pet_kokkos-0.1.0/include/pet/gemm.hpp +352 -0
- pet_kokkos-0.1.0/include/pet/graph.hpp +229 -0
- pet_kokkos-0.1.0/include/pet/io.hpp +32 -0
- pet_kokkos-0.1.0/include/pet/kokkos.hpp +265 -0
- pet_kokkos-0.1.0/include/pet/model.hpp +343 -0
- pet_kokkos-0.1.0/include/pet/neighbors.hpp +116 -0
- pet_kokkos-0.1.0/include/pet/ozaki.hpp +86 -0
- pet_kokkos-0.1.0/include/pet/safetensors.hpp +37 -0
- pet_kokkos-0.1.0/include/pet/vesin_neighbors.hpp +20 -0
- pet_kokkos-0.1.0/lammps/KOKKOS/pair_pet_kokkos.cpp +497 -0
- pet_kokkos-0.1.0/lammps/KOKKOS/pair_pet_kokkos.h +81 -0
- pet_kokkos-0.1.0/lammps/ML-PET/README +12 -0
- pet_kokkos-0.1.0/lammps/ML-PET/pair_pet.cpp +747 -0
- pet_kokkos-0.1.0/lammps/ML-PET/pair_pet.h +154 -0
- pet_kokkos-0.1.0/lammps/README.md +123 -0
- pet_kokkos-0.1.0/lammps/cmake/ML-PET.cmake +25 -0
- pet_kokkos-0.1.0/lammps/examples/README +9 -0
- pet_kokkos-0.1.0/lammps/examples/compare.py +36 -0
- pet_kokkos-0.1.0/lammps/examples/diamond_1728.xyz +1730 -0
- pet_kokkos-0.1.0/lammps/examples/diamond_216.xyz +218 -0
- pet_kokkos-0.1.0/lammps/examples/diamond_64.xyz +66 -0
- pet_kokkos-0.1.0/lammps/examples/in.pet +40 -0
- pet_kokkos-0.1.0/lammps/examples/water_648.xyz +650 -0
- pet_kokkos-0.1.0/lammps/examples/xyz2data.py +37 -0
- pet_kokkos-0.1.0/lammps/patch_lammps.sh +34 -0
- pet_kokkos-0.1.0/metatomic/CMakeLists.txt +44 -0
- pet_kokkos-0.1.0/metatomic/examples/ase_calculator.py +153 -0
- pet_kokkos-0.1.0/metatomic/kokkos_array.hpp +276 -0
- pet_kokkos-0.1.0/metatomic/pet_metatomic.cpp +793 -0
- pet_kokkos-0.1.0/metatomic/test_plugin.py +581 -0
- pet_kokkos-0.1.0/pyproject.toml +123 -0
- pet_kokkos-0.1.0/python/CMakeLists.txt +32 -0
- pet_kokkos-0.1.0/python/core.cpp +206 -0
- pet_kokkos-0.1.0/python/pet_kokkos/__init__.py +138 -0
- pet_kokkos-0.1.0/python/pet_kokkos/_backend.py +32 -0
- pet_kokkos-0.1.0/python/pet_kokkos/ase.py +51 -0
- pet_kokkos-0.1.0/python/pet_kokkos/models.py +67 -0
- pet_kokkos-0.1.0/python/tests/test_pet_kokkos.py +111 -0
- pet_kokkos-0.1.0/python/wheel_meta/pet_variant.py +44 -0
- pet_kokkos-0.1.0/python/wheels/cuda.toml +30 -0
- pet_kokkos-0.1.0/python/wheels/rocm.toml +30 -0
- pet_kokkos-0.1.0/src/CMakeLists.txt +142 -0
- pet_kokkos-0.1.0/src/attention.cpp +1289 -0
- pet_kokkos-0.1.0/src/calculator.cpp +609 -0
- pet_kokkos-0.1.0/src/checkpoint.cpp +127 -0
- pet_kokkos-0.1.0/src/device_cell_list.cpp +295 -0
- pet_kokkos-0.1.0/src/device_geometry.cpp +780 -0
- pet_kokkos-0.1.0/src/device_neighbors.cpp +378 -0
- pet_kokkos-0.1.0/src/feedforward.cpp +323 -0
- pet_kokkos-0.1.0/src/forces.cpp +395 -0
- pet_kokkos-0.1.0/src/fused_gemm.cpp +524 -0
- pet_kokkos-0.1.0/src/gemm_lowp.cpp +366 -0
- pet_kokkos-0.1.0/src/io/extxyz.cpp +183 -0
- pet_kokkos-0.1.0/src/model.cpp +769 -0
- pet_kokkos-0.1.0/src/neighbors.cpp +312 -0
- pet_kokkos-0.1.0/src/ops.cpp +311 -0
- pet_kokkos-0.1.0/src/ops.hpp +218 -0
- pet_kokkos-0.1.0/src/ozaki.cpp +326 -0
- pet_kokkos-0.1.0/src/residual.cpp +237 -0
- pet_kokkos-0.1.0/src/vesin_neighbors.cpp +120 -0
- pet_kokkos-0.1.0/tests/CMakeLists.txt +56 -0
- pet_kokkos-0.1.0/tests/golden/pet-mad-xs-v1.6_crystal.json +170 -0
- pet_kokkos-0.1.0/tests/golden/pet-mad-xs-v1.6_disordered.json +170 -0
- pet_kokkos-0.1.0/tests/golden/pet-mad-xs-v1.6_molecule.json +93 -0
- pet_kokkos-0.1.0/tests/golden/pet-mad-xs_crystal.json +182 -0
- pet_kokkos-0.1.0/tests/golden/pet-mad-xs_disordered.json +182 -0
- pet_kokkos-0.1.0/tests/golden/pet-mad-xs_molecule.json +97 -0
- pet_kokkos-0.1.0/tests/golden/pet-mols-s-v1.1_crystal.json +1 -0
- pet_kokkos-0.1.0/tests/golden/pet-mols-s-v1.1_disordered.json +1 -0
- pet_kokkos-0.1.0/tests/golden/pet-mols-s-v1.1_molecule.json +1 -0
- pet_kokkos-0.1.0/tests/test_cell_list.cpp +212 -0
- pet_kokkos-0.1.0/tests/test_determinism.cpp +491 -0
- pet_kokkos-0.1.0/tests/test_device_vs_host.cpp +292 -0
- pet_kokkos-0.1.0/tests/test_edges.cpp +779 -0
- pet_kokkos-0.1.0/tests/test_fd.cpp +161 -0
- pet_kokkos-0.1.0/tests/test_gemm_modes.cpp +122 -0
- pet_kokkos-0.1.0/tests/test_golden.cpp +213 -0
- pet_kokkos-0.1.0/tests/test_io.cpp +141 -0
- pet_kokkos-0.1.0/tests/test_main.cpp +22 -0
- pet_kokkos-0.1.0/tests/test_neighbors.cpp +359 -0
- pet_kokkos-0.1.0/tests/test_ozaki.cpp +209 -0
- pet_kokkos-0.1.0/tests/test_support.hpp +208 -0
- pet_kokkos-0.1.0/tools/bench_vs_upet.py +112 -0
- pet_kokkos-0.1.0/tools/convert_pet.py +347 -0
- pet_kokkos-0.1.0/tools/make_fp64_golden.py +65 -0
- pet_kokkos-0.1.0/tools/make_golden.py +251 -0
- pet_kokkos-0.1.0/tools/test_all_models.py +405 -0
- pet_kokkos-0.1.0/tools/wheel_index.py +67 -0
|
@@ -0,0 +1,50 @@
|
|
|
1
|
+
name: CI
|
|
2
|
+
|
|
3
|
+
on:
|
|
4
|
+
push:
|
|
5
|
+
branches: [main]
|
|
6
|
+
pull_request:
|
|
7
|
+
|
|
8
|
+
jobs:
|
|
9
|
+
test:
|
|
10
|
+
runs-on: ubuntu-24.04
|
|
11
|
+
env:
|
|
12
|
+
CPM_SOURCE_CACHE: ${{ github.workspace }}/.cpm
|
|
13
|
+
steps:
|
|
14
|
+
- uses: actions/checkout@v4
|
|
15
|
+
|
|
16
|
+
- uses: astral-sh/setup-uv@v6
|
|
17
|
+
|
|
18
|
+
# The test models: published checkpoints, converted here. Weights are never
|
|
19
|
+
# committed; the cache keeps them between runs.
|
|
20
|
+
- name: Cache models
|
|
21
|
+
id: models
|
|
22
|
+
uses: actions/cache@v4
|
|
23
|
+
with:
|
|
24
|
+
path: models
|
|
25
|
+
key: models-${{ hashFiles('tools/convert_pet.py') }}
|
|
26
|
+
|
|
27
|
+
- name: Make models
|
|
28
|
+
if: steps.models.outputs.cache-hit != 'true'
|
|
29
|
+
run: |
|
|
30
|
+
uv run tools/convert_pet.py --model pet-mad-xs --version 1.5.0 --out models/pet-mad-xs
|
|
31
|
+
uv run tools/convert_pet.py --model pet-mad-xs --version 1.6.0 --out models/pet-mad-xs-v1.6
|
|
32
|
+
uv run tools/convert_pet.py --model pet-mols-s --version 1.1.0 --out models/pet-mols-s-v1.1
|
|
33
|
+
uv run tools/convert_pet.py --model pet-omol-s --version 1.0.0 --out models/pet-omol-s
|
|
34
|
+
|
|
35
|
+
- name: Cache dependencies
|
|
36
|
+
uses: actions/cache@v4
|
|
37
|
+
with:
|
|
38
|
+
path: .cpm
|
|
39
|
+
key: cpm-${{ hashFiles('CMakeLists.txt', 'tests/CMakeLists.txt') }}
|
|
40
|
+
|
|
41
|
+
- name: Configure
|
|
42
|
+
run: cmake --preset serial -D PET_ARCH_NATIVE=OFF
|
|
43
|
+
|
|
44
|
+
- name: Build
|
|
45
|
+
run: cmake --build build-serial -j 4
|
|
46
|
+
|
|
47
|
+
- name: Test
|
|
48
|
+
env:
|
|
49
|
+
PET_TEST_QUICK: "1" # skips the batch-invariance test: minutes on the serial backend
|
|
50
|
+
run: ctest --preset serial -j 4
|
|
@@ -0,0 +1,146 @@
|
|
|
1
|
+
name: Wheels
|
|
2
|
+
|
|
3
|
+
# pet-kokkos's wheels: on a tag, by hand, or on a pull request that touches
|
|
4
|
+
# the packaging.
|
|
5
|
+
# - CPU (pyproject.toml's [tool.cibuildwheel]): Linux x86_64 and aarch64
|
|
6
|
+
# (OpenMP, OpenBLAS), macOS arm64 (threads, Accelerate), and the sdist: to
|
|
7
|
+
# PyPI on a tag, by trusted publishing (the pypi environment).
|
|
8
|
+
# - GPU variants (python/wheels/): pet-kokkos 0.1.0+cu12, +cu13 (x86_64,
|
|
9
|
+
# aarch64), +rocm6, +rocm7 (x86_64). PyPI takes no local versions: on a tag
|
|
10
|
+
# they go to the GitHub release, and tools/wheel_index.py publishes an index
|
|
11
|
+
# for each on GitHub Pages (Settings > Pages > Source: GitHub Actions).
|
|
12
|
+
# Built without a GPU; tested on real hardware by hand.
|
|
13
|
+
|
|
14
|
+
on:
|
|
15
|
+
workflow_dispatch:
|
|
16
|
+
push:
|
|
17
|
+
tags: ["v*"]
|
|
18
|
+
pull_request:
|
|
19
|
+
paths: ["pyproject.toml", "python/**", "tools/wheel_index.py", ".github/workflows/wheels.yml"]
|
|
20
|
+
|
|
21
|
+
permissions:
|
|
22
|
+
contents: read
|
|
23
|
+
|
|
24
|
+
jobs:
|
|
25
|
+
cpu:
|
|
26
|
+
name: cpu ${{ matrix.os }}
|
|
27
|
+
runs-on: ${{ matrix.os }}
|
|
28
|
+
strategy:
|
|
29
|
+
fail-fast: false
|
|
30
|
+
matrix:
|
|
31
|
+
os: [ubuntu-24.04, ubuntu-24.04-arm, macos-14]
|
|
32
|
+
steps:
|
|
33
|
+
- uses: actions/checkout@v4
|
|
34
|
+
- uses: astral-sh/setup-uv@v6
|
|
35
|
+
- uses: pypa/cibuildwheel@v3.2
|
|
36
|
+
with:
|
|
37
|
+
output-dir: wheelhouse
|
|
38
|
+
- uses: actions/upload-artifact@v4
|
|
39
|
+
with:
|
|
40
|
+
name: wheels-cpu-${{ matrix.os }}
|
|
41
|
+
path: wheelhouse/*.whl
|
|
42
|
+
|
|
43
|
+
gpu:
|
|
44
|
+
name: ${{ matrix.variant }} ${{ matrix.os }}
|
|
45
|
+
runs-on: ${{ matrix.os }}
|
|
46
|
+
strategy:
|
|
47
|
+
fail-fast: false
|
|
48
|
+
matrix:
|
|
49
|
+
include:
|
|
50
|
+
- { variant: cu12, os: ubuntu-24.04, cuda: "12-8" }
|
|
51
|
+
- { variant: cu12, os: ubuntu-24.04-arm, cuda: "12-8" }
|
|
52
|
+
- { variant: cu13, os: ubuntu-24.04, cuda: "13-0" }
|
|
53
|
+
- { variant: cu13, os: ubuntu-24.04-arm, cuda: "13-0" }
|
|
54
|
+
- { variant: rocm6, os: ubuntu-24.04, rocm: "6.3.2", archs: "--offload-arch=gfx942" }
|
|
55
|
+
- { variant: rocm7, os: ubuntu-24.04, rocm: "7.0.3", archs: "--offload-arch=gfx942 --offload-arch=gfx950" }
|
|
56
|
+
env:
|
|
57
|
+
PET_KOKKOS_VARIANT: ${{ matrix.variant }}
|
|
58
|
+
PET_CUDA: ${{ matrix.cuda }}
|
|
59
|
+
PET_ROCM: ${{ matrix.rocm }}
|
|
60
|
+
PET_ROCM_ARCHS: ${{ matrix.archs }}
|
|
61
|
+
steps:
|
|
62
|
+
- uses: actions/checkout@v4
|
|
63
|
+
- uses: astral-sh/setup-uv@v6
|
|
64
|
+
# The toolkits are large: make room on the runner first.
|
|
65
|
+
- run: sudo rm -rf /usr/share/dotnet /usr/local/lib/android /opt/ghc /opt/hostedtoolcache/CodeQL
|
|
66
|
+
- uses: pypa/cibuildwheel@v3.2
|
|
67
|
+
with:
|
|
68
|
+
config-file: python/wheels/${{ startsWith(matrix.variant, 'cu') && 'cuda' || 'rocm' }}.toml
|
|
69
|
+
output-dir: wheelhouse
|
|
70
|
+
- uses: actions/upload-artifact@v4
|
|
71
|
+
with:
|
|
72
|
+
name: wheels-${{ matrix.variant }}-${{ matrix.os }}
|
|
73
|
+
path: wheelhouse/*.whl
|
|
74
|
+
|
|
75
|
+
sdist:
|
|
76
|
+
runs-on: ubuntu-24.04
|
|
77
|
+
steps:
|
|
78
|
+
- uses: actions/checkout@v4
|
|
79
|
+
- uses: astral-sh/setup-uv@v6
|
|
80
|
+
- run: uv build --sdist -o dist
|
|
81
|
+
- uses: actions/upload-artifact@v4
|
|
82
|
+
with:
|
|
83
|
+
name: sdist
|
|
84
|
+
path: dist/*.tar.gz
|
|
85
|
+
|
|
86
|
+
# On a tag: every wheel and the sdist to the GitHub release.
|
|
87
|
+
release:
|
|
88
|
+
if: startsWith(github.ref, 'refs/tags/v')
|
|
89
|
+
needs: [cpu, gpu, sdist]
|
|
90
|
+
runs-on: ubuntu-24.04
|
|
91
|
+
permissions:
|
|
92
|
+
contents: write
|
|
93
|
+
steps:
|
|
94
|
+
- uses: actions/download-artifact@v4
|
|
95
|
+
with:
|
|
96
|
+
path: dist
|
|
97
|
+
merge-multiple: true
|
|
98
|
+
- env:
|
|
99
|
+
GH_TOKEN: ${{ github.token }}
|
|
100
|
+
run: |
|
|
101
|
+
gh release view "$GITHUB_REF_NAME" -R "$GITHUB_REPOSITORY" >/dev/null 2>&1 ||
|
|
102
|
+
gh release create "$GITHUB_REF_NAME" -R "$GITHUB_REPOSITORY" --verify-tag --generate-notes
|
|
103
|
+
gh release upload "$GITHUB_REF_NAME" -R "$GITHUB_REPOSITORY" --clobber dist/*
|
|
104
|
+
|
|
105
|
+
# On a tag: the CPU wheels and the sdist to PyPI.
|
|
106
|
+
pypi:
|
|
107
|
+
if: startsWith(github.ref, 'refs/tags/v')
|
|
108
|
+
needs: [cpu, sdist]
|
|
109
|
+
runs-on: ubuntu-24.04
|
|
110
|
+
environment:
|
|
111
|
+
name: pypi
|
|
112
|
+
url: https://pypi.org/p/pet-kokkos
|
|
113
|
+
permissions:
|
|
114
|
+
id-token: write
|
|
115
|
+
steps:
|
|
116
|
+
- uses: actions/download-artifact@v4
|
|
117
|
+
with:
|
|
118
|
+
pattern: wheels-cpu-*
|
|
119
|
+
path: dist
|
|
120
|
+
merge-multiple: true
|
|
121
|
+
- uses: actions/download-artifact@v4
|
|
122
|
+
with:
|
|
123
|
+
name: sdist
|
|
124
|
+
path: dist
|
|
125
|
+
- uses: pypa/gh-action-pypi-publish@release/v1
|
|
126
|
+
|
|
127
|
+
# Then the GPU variants' indexes, over every release, to GitHub Pages.
|
|
128
|
+
index:
|
|
129
|
+
needs: release
|
|
130
|
+
runs-on: ubuntu-24.04
|
|
131
|
+
permissions:
|
|
132
|
+
pages: write
|
|
133
|
+
id-token: write
|
|
134
|
+
environment:
|
|
135
|
+
name: github-pages
|
|
136
|
+
url: ${{ steps.deploy.outputs.page_url }}
|
|
137
|
+
steps:
|
|
138
|
+
- uses: actions/checkout@v4
|
|
139
|
+
- env:
|
|
140
|
+
GH_TOKEN: ${{ github.token }}
|
|
141
|
+
run: python3 tools/wheel_index.py "$GITHUB_REPOSITORY" site
|
|
142
|
+
- uses: actions/upload-pages-artifact@v3
|
|
143
|
+
with:
|
|
144
|
+
path: site
|
|
145
|
+
- id: deploy
|
|
146
|
+
uses: actions/deploy-pages@v4
|
|
@@ -0,0 +1,29 @@
|
|
|
1
|
+
# Build trees
|
|
2
|
+
build/
|
|
3
|
+
build-*/
|
|
4
|
+
_deps/
|
|
5
|
+
compile_commands.json
|
|
6
|
+
CMakeUserPresets.json
|
|
7
|
+
|
|
8
|
+
# Model weights are NEVER committed. Some checkpoints are distributed under
|
|
9
|
+
# terms that do not permit redistribution; none of them belong in git.
|
|
10
|
+
models/
|
|
11
|
+
*.safetensors
|
|
12
|
+
*.ckpt
|
|
13
|
+
*.pt
|
|
14
|
+
|
|
15
|
+
# Python
|
|
16
|
+
.venv/
|
|
17
|
+
__pycache__/
|
|
18
|
+
*.egg-info/
|
|
19
|
+
dist/
|
|
20
|
+
*.so
|
|
21
|
+
*.whl
|
|
22
|
+
|
|
23
|
+
# Editors / OS
|
|
24
|
+
.vscode/
|
|
25
|
+
.idea/
|
|
26
|
+
.DS_Store
|
|
27
|
+
*.swp
|
|
28
|
+
build-models/
|
|
29
|
+
build-py/
|
|
@@ -0,0 +1,253 @@
|
|
|
1
|
+
# pet-kokkos — pure C++/Kokkos evaluation of PET (Point Edge Transformer)
|
|
2
|
+
# machine-learning interatomic potentials: energy, conservative forces and the
|
|
3
|
+
# virial, with no libtorch at run time.
|
|
4
|
+
#
|
|
5
|
+
# Kokkos and nlohmann/json are fetched here unless the consumer already has
|
|
6
|
+
# them.
|
|
7
|
+
cmake_minimum_required(VERSION 3.18)
|
|
8
|
+
|
|
9
|
+
project(pet_kokkos
|
|
10
|
+
VERSION 0.1.0
|
|
11
|
+
DESCRIPTION "Kokkos evaluation of PET machine-learning interatomic potentials"
|
|
12
|
+
LANGUAGES CXX)
|
|
13
|
+
|
|
14
|
+
# C++17 is all the code needs; 20 by default, as Kokkos 5 requires. A consumer
|
|
15
|
+
# on Kokkos 4 (LAMMPS bundles 4.6) may set 17.
|
|
16
|
+
if(NOT CMAKE_CXX_STANDARD)
|
|
17
|
+
set(CMAKE_CXX_STANDARD 20)
|
|
18
|
+
endif()
|
|
19
|
+
set(CMAKE_CXX_STANDARD_REQUIRED ON)
|
|
20
|
+
set(CMAKE_CXX_EXTENSIONS OFF)
|
|
21
|
+
set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
|
|
22
|
+
set(CMAKE_POSITION_INDEPENDENT_CODE ON)
|
|
23
|
+
|
|
24
|
+
# Whether we are the top-level project or a consumer's add_subdirectory/CPM
|
|
25
|
+
# dependency. Tests, apps and installation default OFF for a consumer.
|
|
26
|
+
if(CMAKE_SOURCE_DIR STREQUAL CMAKE_CURRENT_SOURCE_DIR)
|
|
27
|
+
set(PET_IS_TOP_LEVEL ON)
|
|
28
|
+
else()
|
|
29
|
+
set(PET_IS_TOP_LEVEL OFF)
|
|
30
|
+
endif()
|
|
31
|
+
|
|
32
|
+
if(NOT CMAKE_BUILD_TYPE AND NOT CMAKE_CONFIGURATION_TYPES)
|
|
33
|
+
set(CMAKE_BUILD_TYPE Release CACHE STRING "Build type" FORCE)
|
|
34
|
+
endif()
|
|
35
|
+
|
|
36
|
+
# --- options -----------------------------------------------------------------
|
|
37
|
+
|
|
38
|
+
option(PET_BUILD_TESTS "Build the test suite" ${PET_IS_TOP_LEVEL})
|
|
39
|
+
option(PET_BUILD_APPS "Build the command-line tools" ${PET_IS_TOP_LEVEL})
|
|
40
|
+
option(PET_INSTALL "Install the library and headers" ${PET_IS_TOP_LEVEL})
|
|
41
|
+
# The metatomic plugin (metatomic/): needs metatomic and metatensor installed.
|
|
42
|
+
option(PET_BUILD_METATOMIC "Build the metatomic plugin" OFF)
|
|
43
|
+
# The Python package's compiled module (python/): needs nanobind. A wheel build
|
|
44
|
+
# (pyproject.toml) turns it on.
|
|
45
|
+
option(PET_BUILD_PYTHON "Build the Python module pet_kokkos._core" OFF)
|
|
46
|
+
|
|
47
|
+
# Optional vesin backend for the host neighbour search. The built-in search is an
|
|
48
|
+
# O(N^2 x images) brute force, fine for small cells and the dominant cost for
|
|
49
|
+
# large ones; vesin's cell list does the same work in O(N).
|
|
50
|
+
#
|
|
51
|
+
# AUTO uses it when it can be found (the `vesin` Python wheel ships the header
|
|
52
|
+
# and the library, so a venv counts) and silently does without otherwise; ON
|
|
53
|
+
# fails the configure if it is missing; OFF never looks.
|
|
54
|
+
set(PET_WITH_VESIN "AUTO" CACHE STRING "Use vesin for neighbour lists: AUTO | ON | OFF")
|
|
55
|
+
set_property(CACHE PET_WITH_VESIN PROPERTY STRINGS AUTO ON OFF)
|
|
56
|
+
set(PET_WITH_CUTLASS "AUTO" CACHE STRING "Fused GEMM kernels on CUTLASS (CUDA only): AUTO | ON | OFF")
|
|
57
|
+
set_property(CACHE PET_WITH_CUTLASS PROPERTY STRINGS AUTO ON OFF)
|
|
58
|
+
|
|
59
|
+
# Tune for the building machine's CPU. A real win on the host kernels, but it
|
|
60
|
+
# bakes THIS machine's instruction set into the objects: turn it off when the
|
|
61
|
+
# build host and the run host may differ, or when objects are cached and reused
|
|
62
|
+
# elsewhere (ccache hashes the literal "-march=native", not its expansion, so it
|
|
63
|
+
# can reuse an object on a CPU that cannot run it). Never for CUDA: nvcc does
|
|
64
|
+
# not accept the flag.
|
|
65
|
+
option(PET_ARCH_NATIVE "Optimize for the building machine's CPU (-march=native)" ON)
|
|
66
|
+
option(PET_USE_BLAS "Host builds: the GEMMs through a CBLAS (OpenBLAS, MKL, ...) when one is found" ON)
|
|
67
|
+
# GPU builds run Kokkos's host space on OpenMP; off, on one thread (a wheel
|
|
68
|
+
# without a host OpenMP runtime to ship: the ROCm one).
|
|
69
|
+
option(PET_HOST_OPENMP "GPU builds: OpenMP for the host" ON)
|
|
70
|
+
|
|
71
|
+
# threads: Kokkos's own C++ threads, for host compilers without OpenMP (Apple's).
|
|
72
|
+
set(PET_BACKEND "serial" CACHE STRING "Kokkos backend: serial, openmp, threads, cuda, hip")
|
|
73
|
+
set_property(CACHE PET_BACKEND PROPERTY STRINGS serial openmp threads cuda hip)
|
|
74
|
+
# Inside an application that brings its own Kokkos (LAMMPS's KOKKOS package),
|
|
75
|
+
# that Kokkos's device is ours: in-tree it leaves Kokkos_ENABLE_*, installed
|
|
76
|
+
# Kokkos_DEVICES.
|
|
77
|
+
if(TARGET Kokkos::kokkos)
|
|
78
|
+
foreach(dev CUDA HIP OPENMP THREADS)
|
|
79
|
+
if(Kokkos_ENABLE_${dev} OR "${dev}" IN_LIST Kokkos_DEVICES)
|
|
80
|
+
string(TOLOWER ${dev} PET_BACKEND)
|
|
81
|
+
break()
|
|
82
|
+
endif()
|
|
83
|
+
endforeach()
|
|
84
|
+
set(PET_BACKEND ${PET_BACKEND} CACHE STRING "Kokkos backend" FORCE)
|
|
85
|
+
endif()
|
|
86
|
+
|
|
87
|
+
# Precision mode. PUBLIC on the library target so a consumer sees the same
|
|
88
|
+
# Net/Real types the library was compiled with.
|
|
89
|
+
# mixed (default) : Net=float, Real=double -- fp32 network, fp64 geometry,
|
|
90
|
+
# energy, force and virial accumulation. What the goldens
|
|
91
|
+
# are validated against.
|
|
92
|
+
# fp32 : everything single. Fastest on consumer GPUs.
|
|
93
|
+
# fp64 : everything double. Slow on consumer GPUs (1/64 rate on
|
|
94
|
+
# GeForce); a correctness instrument, not a speed mode.
|
|
95
|
+
set(PET_PRECISION "mixed" CACHE STRING "PET precision: mixed | fp32 | fp64")
|
|
96
|
+
set_property(CACHE PET_PRECISION PROPERTY STRINGS mixed fp32 fp64)
|
|
97
|
+
|
|
98
|
+
if(PET_ARCH_NATIVE AND NOT PET_BACKEND STREQUAL "cuda")
|
|
99
|
+
include(CheckCXXCompilerFlag)
|
|
100
|
+
check_cxx_compiler_flag(-march=native PET_HAS_MARCH_NATIVE)
|
|
101
|
+
if(PET_HAS_MARCH_NATIVE)
|
|
102
|
+
add_compile_options(-march=native)
|
|
103
|
+
endif()
|
|
104
|
+
endif()
|
|
105
|
+
|
|
106
|
+
# --- CPM ---------------------------------------------------------------------
|
|
107
|
+
|
|
108
|
+
set(CPM_DOWNLOAD_VERSION 0.40.8)
|
|
109
|
+
if(NOT EXISTS "${CMAKE_BINARY_DIR}/cmake/CPM.cmake")
|
|
110
|
+
file(DOWNLOAD
|
|
111
|
+
"https://github.com/cpm-cmake/CPM.cmake/releases/download/v${CPM_DOWNLOAD_VERSION}/CPM.cmake"
|
|
112
|
+
"${CMAKE_BINARY_DIR}/cmake/CPM.cmake")
|
|
113
|
+
endif()
|
|
114
|
+
include(${CMAKE_BINARY_DIR}/cmake/CPM.cmake)
|
|
115
|
+
|
|
116
|
+
# --- Kokkos ------------------------------------------------------------------
|
|
117
|
+
#
|
|
118
|
+
# CPM's OPTIONS are not promoted to cache variables, so anything that later does
|
|
119
|
+
# find_package(Kokkos) or reads Kokkos_ENABLE_* cannot see them. Set every flag
|
|
120
|
+
# as a FORCE'd cache var up front so Kokkos and any downstream consumer agree.
|
|
121
|
+
#
|
|
122
|
+
# Only fetch when the consumer has not already provided Kokkos: a Kokkos
|
|
123
|
+
# application owns one already, and two Kokkos configurations cannot share a
|
|
124
|
+
# build tree.
|
|
125
|
+
if(NOT TARGET Kokkos::kokkos)
|
|
126
|
+
set(Kokkos_ENABLE_SERIAL ON CACHE BOOL "" FORCE)
|
|
127
|
+
set(Kokkos_ENABLE_OPENMP OFF CACHE BOOL "" FORCE)
|
|
128
|
+
set(Kokkos_ENABLE_CUDA OFF CACHE BOOL "" FORCE)
|
|
129
|
+
set(Kokkos_ENABLE_HIP OFF CACHE BOOL "" FORCE)
|
|
130
|
+
set(Kokkos_ENABLE_THREADS OFF CACHE BOOL "" FORCE)
|
|
131
|
+
|
|
132
|
+
if(PET_BACKEND STREQUAL "openmp")
|
|
133
|
+
set(Kokkos_ENABLE_OPENMP ON CACHE BOOL "" FORCE)
|
|
134
|
+
elseif(PET_BACKEND STREQUAL "threads")
|
|
135
|
+
set(Kokkos_ENABLE_THREADS ON CACHE BOOL "" FORCE)
|
|
136
|
+
elseif(PET_BACKEND STREQUAL "cuda")
|
|
137
|
+
set(Kokkos_ENABLE_CUDA ON CACHE BOOL "" FORCE)
|
|
138
|
+
set(Kokkos_ENABLE_CUDA_LAMBDA ON CACHE BOOL "" FORCE)
|
|
139
|
+
set(Kokkos_ENABLE_CUDA_CONSTEXPR ON CACHE BOOL "" FORCE)
|
|
140
|
+
set(Kokkos_ENABLE_OPENMP ${PET_HOST_OPENMP} CACHE BOOL "" FORCE)
|
|
141
|
+
elseif(PET_BACKEND STREQUAL "hip")
|
|
142
|
+
set(Kokkos_ENABLE_HIP ON CACHE BOOL "" FORCE)
|
|
143
|
+
set(Kokkos_ENABLE_OPENMP ${PET_HOST_OPENMP} CACHE BOOL "" FORCE)
|
|
144
|
+
elseif(NOT PET_BACKEND STREQUAL "serial")
|
|
145
|
+
message(FATAL_ERROR "PET_BACKEND must be one of: serial, openmp, threads, cuda, hip")
|
|
146
|
+
endif()
|
|
147
|
+
|
|
148
|
+
CPMAddPackage(
|
|
149
|
+
NAME Kokkos
|
|
150
|
+
GITHUB_REPOSITORY kokkos/kokkos
|
|
151
|
+
GIT_TAG 5.2.2
|
|
152
|
+
)
|
|
153
|
+
else()
|
|
154
|
+
message(STATUS "pet-kokkos: using the consumer's Kokkos::kokkos (${PET_BACKEND})")
|
|
155
|
+
endif()
|
|
156
|
+
|
|
157
|
+
# --- nlohmann/json -----------------------------------------------------------
|
|
158
|
+
|
|
159
|
+
if(NOT TARGET nlohmann_json::nlohmann_json)
|
|
160
|
+
CPMAddPackage(
|
|
161
|
+
NAME nlohmann_json
|
|
162
|
+
VERSION 3.11.3
|
|
163
|
+
GITHUB_REPOSITORY nlohmann/json
|
|
164
|
+
OPTIONS "JSON_BuildTests OFF"
|
|
165
|
+
)
|
|
166
|
+
endif()
|
|
167
|
+
|
|
168
|
+
# --- CUTLASS (optional, CUDA) -----------------------------------------------
|
|
169
|
+
# Header-only: its fp32 GEMM main loop carries the fused kernels in
|
|
170
|
+
# src/fused_gemm.cpp. Without it those operations run unfused.
|
|
171
|
+
|
|
172
|
+
set(PET_HAVE_CUTLASS OFF)
|
|
173
|
+
if(PET_BACKEND STREQUAL "cuda" AND NOT PET_WITH_CUTLASS STREQUAL "OFF")
|
|
174
|
+
CPMAddPackage(
|
|
175
|
+
NAME cutlass
|
|
176
|
+
GITHUB_REPOSITORY NVIDIA/cutlass
|
|
177
|
+
GIT_TAG v4.8.0
|
|
178
|
+
DOWNLOAD_ONLY YES
|
|
179
|
+
)
|
|
180
|
+
if(cutlass_ADDED)
|
|
181
|
+
set(PET_HAVE_CUTLASS ON)
|
|
182
|
+
message(STATUS "pet-kokkos: fused GEMMs on CUTLASS from ${cutlass_SOURCE_DIR}")
|
|
183
|
+
elseif(PET_WITH_CUTLASS STREQUAL "ON")
|
|
184
|
+
message(FATAL_ERROR "pet-kokkos: PET_WITH_CUTLASS=ON but CUTLASS could not be fetched")
|
|
185
|
+
endif()
|
|
186
|
+
endif()
|
|
187
|
+
|
|
188
|
+
# --- vesin (optional) --------------------------------------------------------
|
|
189
|
+
|
|
190
|
+
set(PET_HAVE_VESIN OFF)
|
|
191
|
+
if(NOT PET_WITH_VESIN STREQUAL "OFF")
|
|
192
|
+
list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake")
|
|
193
|
+
if(PET_WITH_VESIN STREQUAL "ON")
|
|
194
|
+
find_package(Vesin REQUIRED)
|
|
195
|
+
else()
|
|
196
|
+
find_package(Vesin QUIET)
|
|
197
|
+
endif()
|
|
198
|
+
if(Vesin_FOUND)
|
|
199
|
+
set(PET_HAVE_VESIN ON)
|
|
200
|
+
message(STATUS "pet-kokkos: vesin neighbour lists from ${Vesin_LIBRARY}")
|
|
201
|
+
elseif(PET_WITH_VESIN STREQUAL "AUTO")
|
|
202
|
+
message(STATUS "pet-kokkos: vesin not found; using the built-in neighbour search")
|
|
203
|
+
endif()
|
|
204
|
+
endif()
|
|
205
|
+
|
|
206
|
+
# --- warnings ----------------------------------------------------------------
|
|
207
|
+
#
|
|
208
|
+
# Applied to our own targets only, never to Kokkos. Nothing here is an error:
|
|
209
|
+
# -Werror in a project whose kernels are compiled by three different vendor
|
|
210
|
+
# toolchains breaks a build over a warning one of them invented.
|
|
211
|
+
#
|
|
212
|
+
# sign-compare and unused-parameter are off because Kokkos' own idioms trip them
|
|
213
|
+
# in our code.
|
|
214
|
+
add_library(pet_warnings INTERFACE)
|
|
215
|
+
if(MSVC)
|
|
216
|
+
target_compile_options(pet_warnings INTERFACE /W3)
|
|
217
|
+
else()
|
|
218
|
+
target_compile_options(pet_warnings INTERFACE
|
|
219
|
+
-Wall -Wextra -Wno-unused-parameter -Wno-sign-compare)
|
|
220
|
+
endif()
|
|
221
|
+
|
|
222
|
+
# --- the library -------------------------------------------------------------
|
|
223
|
+
|
|
224
|
+
add_subdirectory(src)
|
|
225
|
+
|
|
226
|
+
if(PET_BUILD_APPS)
|
|
227
|
+
add_subdirectory(apps)
|
|
228
|
+
endif()
|
|
229
|
+
|
|
230
|
+
if(PET_BUILD_METATOMIC)
|
|
231
|
+
add_subdirectory(metatomic)
|
|
232
|
+
endif()
|
|
233
|
+
|
|
234
|
+
if(PET_BUILD_PYTHON)
|
|
235
|
+
add_subdirectory(python)
|
|
236
|
+
endif()
|
|
237
|
+
|
|
238
|
+
if(PET_BUILD_TESTS)
|
|
239
|
+
enable_testing()
|
|
240
|
+
add_subdirectory(tests)
|
|
241
|
+
endif()
|
|
242
|
+
|
|
243
|
+
# --- summary -----------------------------------------------------------------
|
|
244
|
+
|
|
245
|
+
message(STATUS "pet-kokkos ${PROJECT_VERSION}")
|
|
246
|
+
message(STATUS " backend : ${PET_BACKEND}")
|
|
247
|
+
message(STATUS " precision : ${PET_PRECISION}")
|
|
248
|
+
message(STATUS " build type : ${CMAKE_BUILD_TYPE}")
|
|
249
|
+
message(STATUS " tests : ${PET_BUILD_TESTS} apps: ${PET_BUILD_APPS}")
|
|
250
|
+
message(STATUS " vesin : ${PET_HAVE_VESIN}")
|
|
251
|
+
message(STATUS " cutlass : ${PET_HAVE_CUTLASS}")
|
|
252
|
+
message(STATUS " metatomic : ${PET_BUILD_METATOMIC}")
|
|
253
|
+
message(STATUS " python : ${PET_BUILD_PYTHON}")
|
|
@@ -0,0 +1,70 @@
|
|
|
1
|
+
{
|
|
2
|
+
"version": 3,
|
|
3
|
+
"cmakeMinimumRequired": { "major": 3, "minor": 21, "patch": 0 },
|
|
4
|
+
"configurePresets": [
|
|
5
|
+
{
|
|
6
|
+
"name": "base",
|
|
7
|
+
"hidden": true,
|
|
8
|
+
"binaryDir": "${sourceDir}/build-${presetName}",
|
|
9
|
+
"cacheVariables": {
|
|
10
|
+
"CMAKE_BUILD_TYPE": "Release",
|
|
11
|
+
"CMAKE_EXPORT_COMPILE_COMMANDS": "ON"
|
|
12
|
+
}
|
|
13
|
+
},
|
|
14
|
+
{
|
|
15
|
+
"name": "serial",
|
|
16
|
+
"inherits": "base",
|
|
17
|
+
"displayName": "Serial (host, single-threaded)",
|
|
18
|
+
"cacheVariables": { "PET_BACKEND": "serial" }
|
|
19
|
+
},
|
|
20
|
+
{
|
|
21
|
+
"name": "openmp",
|
|
22
|
+
"inherits": "base",
|
|
23
|
+
"displayName": "OpenMP (host, threaded)",
|
|
24
|
+
"cacheVariables": { "PET_BACKEND": "openmp" }
|
|
25
|
+
},
|
|
26
|
+
{
|
|
27
|
+
"name": "cuda",
|
|
28
|
+
"inherits": "base",
|
|
29
|
+
"displayName": "CUDA (NVIDIA)",
|
|
30
|
+
"cacheVariables": {
|
|
31
|
+
"PET_BACKEND": "cuda"
|
|
32
|
+
}
|
|
33
|
+
},
|
|
34
|
+
{
|
|
35
|
+
"name": "hip",
|
|
36
|
+
"inherits": "base",
|
|
37
|
+
"displayName": "HIP (AMD)",
|
|
38
|
+
"cacheVariables": {
|
|
39
|
+
"PET_BACKEND": "hip",
|
|
40
|
+
"CMAKE_CXX_COMPILER": "hipcc"
|
|
41
|
+
}
|
|
42
|
+
},
|
|
43
|
+
{
|
|
44
|
+
"name": "fp64",
|
|
45
|
+
"inherits": "serial",
|
|
46
|
+
"displayName": "Serial, all-double (a correctness instrument, not a speed mode)",
|
|
47
|
+
"cacheVariables": { "PET_PRECISION": "fp64" }
|
|
48
|
+
},
|
|
49
|
+
{
|
|
50
|
+
"name": "debug",
|
|
51
|
+
"inherits": "serial",
|
|
52
|
+
"displayName": "Serial, debug",
|
|
53
|
+
"cacheVariables": { "CMAKE_BUILD_TYPE": "Debug" }
|
|
54
|
+
}
|
|
55
|
+
],
|
|
56
|
+
"buildPresets": [
|
|
57
|
+
{ "name": "serial", "configurePreset": "serial" },
|
|
58
|
+
{ "name": "openmp", "configurePreset": "openmp" },
|
|
59
|
+
{ "name": "cuda", "configurePreset": "cuda" },
|
|
60
|
+
{ "name": "hip", "configurePreset": "hip" },
|
|
61
|
+
{ "name": "fp64", "configurePreset": "fp64" },
|
|
62
|
+
{ "name": "debug", "configurePreset": "debug" }
|
|
63
|
+
],
|
|
64
|
+
"testPresets": [
|
|
65
|
+
{ "name": "serial", "configurePreset": "serial", "output": { "outputOnFailure": true } },
|
|
66
|
+
{ "name": "openmp", "configurePreset": "openmp", "output": { "outputOnFailure": true } },
|
|
67
|
+
{ "name": "cuda", "configurePreset": "cuda", "output": { "outputOnFailure": true } },
|
|
68
|
+
{ "name": "fp64", "configurePreset": "fp64", "output": { "outputOnFailure": true } }
|
|
69
|
+
]
|
|
70
|
+
}
|
pet_kokkos-0.1.0/LICENSE
ADDED
|
@@ -0,0 +1,29 @@
|
|
|
1
|
+
BSD 3-Clause License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026, Peter Spackman
|
|
4
|
+
All rights reserved.
|
|
5
|
+
|
|
6
|
+
Redistribution and use in source and binary forms, with or without
|
|
7
|
+
modification, are permitted provided that the following conditions are met:
|
|
8
|
+
|
|
9
|
+
1. Redistributions of source code must retain the above copyright notice, this
|
|
10
|
+
list of conditions and the following disclaimer.
|
|
11
|
+
|
|
12
|
+
2. Redistributions in binary form must reproduce the above copyright notice,
|
|
13
|
+
this list of conditions and the following disclaimer in the documentation
|
|
14
|
+
and/or other materials provided with the distribution.
|
|
15
|
+
|
|
16
|
+
3. Neither the name of the copyright holder nor the names of its
|
|
17
|
+
contributors may be used to endorse or promote products derived from
|
|
18
|
+
this software without specific prior written permission.
|
|
19
|
+
|
|
20
|
+
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
|
21
|
+
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
|
22
|
+
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
|
23
|
+
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
|
24
|
+
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
|
25
|
+
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
|
26
|
+
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
|
27
|
+
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
|
28
|
+
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
|
29
|
+
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|