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.
Files changed (108) hide show
  1. pet_kokkos-0.1.0/.github/workflows/ci.yml +50 -0
  2. pet_kokkos-0.1.0/.github/workflows/wheels.yml +146 -0
  3. pet_kokkos-0.1.0/.gitignore +29 -0
  4. pet_kokkos-0.1.0/CMakeLists.txt +253 -0
  5. pet_kokkos-0.1.0/CMakePresets.json +70 -0
  6. pet_kokkos-0.1.0/LICENSE +29 -0
  7. pet_kokkos-0.1.0/PKG-INFO +407 -0
  8. pet_kokkos-0.1.0/README.md +362 -0
  9. pet_kokkos-0.1.0/apps/CMakeLists.txt +14 -0
  10. pet_kokkos-0.1.0/apps/batch_probe.cpp +150 -0
  11. pet_kokkos-0.1.0/apps/pet_eval.cpp +457 -0
  12. pet_kokkos-0.1.0/cmake/FindVesin.cmake +48 -0
  13. pet_kokkos-0.1.0/include/pet/calculator.hpp +242 -0
  14. pet_kokkos-0.1.0/include/pet/checkpoint.hpp +28 -0
  15. pet_kokkos-0.1.0/include/pet/config.hpp +68 -0
  16. pet_kokkos-0.1.0/include/pet/cutoff.hpp +46 -0
  17. pet_kokkos-0.1.0/include/pet/device_cell_list.hpp +63 -0
  18. pet_kokkos-0.1.0/include/pet/device_geometry.hpp +238 -0
  19. pet_kokkos-0.1.0/include/pet/device_neighbors.hpp +69 -0
  20. pet_kokkos-0.1.0/include/pet/exchange.hpp +43 -0
  21. pet_kokkos-0.1.0/include/pet/gemm.hpp +352 -0
  22. pet_kokkos-0.1.0/include/pet/graph.hpp +229 -0
  23. pet_kokkos-0.1.0/include/pet/io.hpp +32 -0
  24. pet_kokkos-0.1.0/include/pet/kokkos.hpp +265 -0
  25. pet_kokkos-0.1.0/include/pet/model.hpp +343 -0
  26. pet_kokkos-0.1.0/include/pet/neighbors.hpp +116 -0
  27. pet_kokkos-0.1.0/include/pet/ozaki.hpp +86 -0
  28. pet_kokkos-0.1.0/include/pet/safetensors.hpp +37 -0
  29. pet_kokkos-0.1.0/include/pet/vesin_neighbors.hpp +20 -0
  30. pet_kokkos-0.1.0/lammps/KOKKOS/pair_pet_kokkos.cpp +497 -0
  31. pet_kokkos-0.1.0/lammps/KOKKOS/pair_pet_kokkos.h +81 -0
  32. pet_kokkos-0.1.0/lammps/ML-PET/README +12 -0
  33. pet_kokkos-0.1.0/lammps/ML-PET/pair_pet.cpp +747 -0
  34. pet_kokkos-0.1.0/lammps/ML-PET/pair_pet.h +154 -0
  35. pet_kokkos-0.1.0/lammps/README.md +123 -0
  36. pet_kokkos-0.1.0/lammps/cmake/ML-PET.cmake +25 -0
  37. pet_kokkos-0.1.0/lammps/examples/README +9 -0
  38. pet_kokkos-0.1.0/lammps/examples/compare.py +36 -0
  39. pet_kokkos-0.1.0/lammps/examples/diamond_1728.xyz +1730 -0
  40. pet_kokkos-0.1.0/lammps/examples/diamond_216.xyz +218 -0
  41. pet_kokkos-0.1.0/lammps/examples/diamond_64.xyz +66 -0
  42. pet_kokkos-0.1.0/lammps/examples/in.pet +40 -0
  43. pet_kokkos-0.1.0/lammps/examples/water_648.xyz +650 -0
  44. pet_kokkos-0.1.0/lammps/examples/xyz2data.py +37 -0
  45. pet_kokkos-0.1.0/lammps/patch_lammps.sh +34 -0
  46. pet_kokkos-0.1.0/metatomic/CMakeLists.txt +44 -0
  47. pet_kokkos-0.1.0/metatomic/examples/ase_calculator.py +153 -0
  48. pet_kokkos-0.1.0/metatomic/kokkos_array.hpp +276 -0
  49. pet_kokkos-0.1.0/metatomic/pet_metatomic.cpp +793 -0
  50. pet_kokkos-0.1.0/metatomic/test_plugin.py +581 -0
  51. pet_kokkos-0.1.0/pyproject.toml +123 -0
  52. pet_kokkos-0.1.0/python/CMakeLists.txt +32 -0
  53. pet_kokkos-0.1.0/python/core.cpp +206 -0
  54. pet_kokkos-0.1.0/python/pet_kokkos/__init__.py +138 -0
  55. pet_kokkos-0.1.0/python/pet_kokkos/_backend.py +32 -0
  56. pet_kokkos-0.1.0/python/pet_kokkos/ase.py +51 -0
  57. pet_kokkos-0.1.0/python/pet_kokkos/models.py +67 -0
  58. pet_kokkos-0.1.0/python/tests/test_pet_kokkos.py +111 -0
  59. pet_kokkos-0.1.0/python/wheel_meta/pet_variant.py +44 -0
  60. pet_kokkos-0.1.0/python/wheels/cuda.toml +30 -0
  61. pet_kokkos-0.1.0/python/wheels/rocm.toml +30 -0
  62. pet_kokkos-0.1.0/src/CMakeLists.txt +142 -0
  63. pet_kokkos-0.1.0/src/attention.cpp +1289 -0
  64. pet_kokkos-0.1.0/src/calculator.cpp +609 -0
  65. pet_kokkos-0.1.0/src/checkpoint.cpp +127 -0
  66. pet_kokkos-0.1.0/src/device_cell_list.cpp +295 -0
  67. pet_kokkos-0.1.0/src/device_geometry.cpp +780 -0
  68. pet_kokkos-0.1.0/src/device_neighbors.cpp +378 -0
  69. pet_kokkos-0.1.0/src/feedforward.cpp +323 -0
  70. pet_kokkos-0.1.0/src/forces.cpp +395 -0
  71. pet_kokkos-0.1.0/src/fused_gemm.cpp +524 -0
  72. pet_kokkos-0.1.0/src/gemm_lowp.cpp +366 -0
  73. pet_kokkos-0.1.0/src/io/extxyz.cpp +183 -0
  74. pet_kokkos-0.1.0/src/model.cpp +769 -0
  75. pet_kokkos-0.1.0/src/neighbors.cpp +312 -0
  76. pet_kokkos-0.1.0/src/ops.cpp +311 -0
  77. pet_kokkos-0.1.0/src/ops.hpp +218 -0
  78. pet_kokkos-0.1.0/src/ozaki.cpp +326 -0
  79. pet_kokkos-0.1.0/src/residual.cpp +237 -0
  80. pet_kokkos-0.1.0/src/vesin_neighbors.cpp +120 -0
  81. pet_kokkos-0.1.0/tests/CMakeLists.txt +56 -0
  82. pet_kokkos-0.1.0/tests/golden/pet-mad-xs-v1.6_crystal.json +170 -0
  83. pet_kokkos-0.1.0/tests/golden/pet-mad-xs-v1.6_disordered.json +170 -0
  84. pet_kokkos-0.1.0/tests/golden/pet-mad-xs-v1.6_molecule.json +93 -0
  85. pet_kokkos-0.1.0/tests/golden/pet-mad-xs_crystal.json +182 -0
  86. pet_kokkos-0.1.0/tests/golden/pet-mad-xs_disordered.json +182 -0
  87. pet_kokkos-0.1.0/tests/golden/pet-mad-xs_molecule.json +97 -0
  88. pet_kokkos-0.1.0/tests/golden/pet-mols-s-v1.1_crystal.json +1 -0
  89. pet_kokkos-0.1.0/tests/golden/pet-mols-s-v1.1_disordered.json +1 -0
  90. pet_kokkos-0.1.0/tests/golden/pet-mols-s-v1.1_molecule.json +1 -0
  91. pet_kokkos-0.1.0/tests/test_cell_list.cpp +212 -0
  92. pet_kokkos-0.1.0/tests/test_determinism.cpp +491 -0
  93. pet_kokkos-0.1.0/tests/test_device_vs_host.cpp +292 -0
  94. pet_kokkos-0.1.0/tests/test_edges.cpp +779 -0
  95. pet_kokkos-0.1.0/tests/test_fd.cpp +161 -0
  96. pet_kokkos-0.1.0/tests/test_gemm_modes.cpp +122 -0
  97. pet_kokkos-0.1.0/tests/test_golden.cpp +213 -0
  98. pet_kokkos-0.1.0/tests/test_io.cpp +141 -0
  99. pet_kokkos-0.1.0/tests/test_main.cpp +22 -0
  100. pet_kokkos-0.1.0/tests/test_neighbors.cpp +359 -0
  101. pet_kokkos-0.1.0/tests/test_ozaki.cpp +209 -0
  102. pet_kokkos-0.1.0/tests/test_support.hpp +208 -0
  103. pet_kokkos-0.1.0/tools/bench_vs_upet.py +112 -0
  104. pet_kokkos-0.1.0/tools/convert_pet.py +347 -0
  105. pet_kokkos-0.1.0/tools/make_fp64_golden.py +65 -0
  106. pet_kokkos-0.1.0/tools/make_golden.py +251 -0
  107. pet_kokkos-0.1.0/tools/test_all_models.py +405 -0
  108. 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
+ }
@@ -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.