comp-chem/cp2k
1
0
Fork 0
mirror of https://github.com/cp2k/cp2k.git synced 2026-07-28 06:05:29 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add support for Torch which is needed by NequIP

Browse source
This commit is contained in:
Ole Schütt 2022-11-29 13:49:14 +01:00 • committed by Ole Schütt
parent 72b5337cad
commit fcb0f647ee
12 changed files with 856 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

13
INSTALL.md
View file

@ -320,7 +320,12 @@ SIRIUS is a domain specific library for electronic structure calculations.
please see <https://brehm-research.de/bqb> for more information as well as
the `bqbtool` to inspect BQB files.
### 2t. ROCM/HIP (Support for AMD GPU)
### 2t. Torch (Machine Learning Framework needed for NequIP)
- The C++ API of PyTorch can be downloaded from https://pytorch.org/get-started/locally/.
- Add `-D__LIBTORCH` to the DFLAGS to enable support for libtorch.
### 2u. ROCM/HIP (Support for AMD GPU)
The code for the HIP based grid backend was developed and tested on Mi100 but
should work out of the box on Nvidia hardware as well.
@ -349,7 +354,7 @@ should work out of the box on Nvidia hardware as well.
- Use `-D__OFFLOAD_PROFILING` to turn on the AMD ROC TX and Tracer libray.
It requires to link `-lroctx64 -lroctracer64`.
### 2u. OpenCL Devices
### 2v. OpenCL Devices
OpenCL devices are currently supported for DBCSR and can cover GPUs and other devices.
Kernels can be automatically tuned like for the CUDA/HIP backend in DBCSR.
@ -394,7 +399,7 @@ Note: the OpenCL backend uses some functionality from LIBXSMM (dependency).
- Refer to <https://cp2k.github.io/dbcsr/> for, e.g., environment variables,
or how to tune kernels (auto tuned parameters).
### 2v. matrix-matrix multiplication offloading on GPU using SPLA
### 2w. matrix-matrix multiplication offloading on GPU using SPLA
The SPLA library is a hard dependency of SIRIUS but can also be used as a
standalone library. It provides a generic interface to the blas gemm family with
@ -408,7 +413,7 @@ internal criteria to decide if it is worth to do the operation on GPU or not.
Calls to `offload_dgemm` also accept pointers on GPU or a combination of them.
<!---
### 2w. LibMaxwell (External Maxwell Solver)
### 2x. LibMaxwell (External Maxwell Solver)
- LibMaxwell is a library to solve the time-dependent Maxwell equations
and use the resulting electric field in MD runs or real-time propagation.

8
Makefile
View file

@ -96,6 +96,9 @@ ifneq ($(filter nvcc hipcc, $(notdir $(OFFLOAD_CC))),)
OBJ_SRC_FILES += $(shell cd $(SRCDIR); find . -name "*.cu")
endif
# Begrudgingly tolerated because PyTorch has no C API.
OBJ_SRC_FILES += $(SRCDIR)/torch_c_api.cpp
# Included files used by Fypp preprocessor
INCLUDED_SRC_FILES = $(notdir $(shell find $(SRCDIR) -name "*.fypp"))
@ -481,6 +484,7 @@ vpath %.h $(ALL_SRC_DIRS)
vpath %.f90 $(ALL_SRC_DIRS)
vpath %.cu $(ALL_SRC_DIRS)
vpath %.c $(ALL_SRC_DIRS)
vpath %.cpp $(ALL_SRC_DIRS)
#
# Add additional dependency of cp2k_info.F to git-HEAD.
@ -530,6 +534,10 @@ FYPPFLAGS ?= -n
%.o: %.C
@echo "Error: C++ is not supported: $<"; false
# Begrudgingly tolerated because PyTorch has no C API.
torch_c_api.o: torch_c_api.cpp
$(CXX) -c $(CXXFLAGS) $<
ifneq ($(LIBDIR),)
$(LIBDIR)/%:
ifneq ($(LD_SHARED),)

BIN
data/NequIP/water.pth Normal file
View file

Binary file not shown.

3
src/common/cp_files.F
View file

@ -30,7 +30,8 @@ MODULE cp_files
open_file, &
get_unit_number, &
file_exists, &
get_data_dir
get_data_dir, &
discover_file
CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'cp_files'

3
src/cp2k_info.F
View file

@ -215,6 +215,9 @@ CONTAINS
#if defined(__LIBMAXWELL)
flags = TRIM(flags)//" libmaxwell"
#endif
#if defined(__LIBTORCH)
flags = TRIM(flags)//" libtorch"
#endif
#if defined(__OFFLOAD_CUDA)
flags = TRIM(flags)//" offload_cuda"
#endif

233
src/nequip_unittest.F Normal file
View file

@ -0,0 +1,233 @@
!--------------------------------------------------------------------------------------------------!
! CP2K: A general program to perform molecular dynamics simulations !
! Copyright 2000-2022 CP2K developers group <https://cp2k.org> !
! !
! SPDX-License-Identifier: GPL-2.0-or-later !
!--------------------------------------------------------------------------------------------------!
PROGRAM nequip_unittest
USE cp_files, ONLY: discover_file
USE kinds, ONLY: default_path_length,&
dp,&
int_8,&
sp
USE mathlib, ONLY: inv_3x3
USE torch_api, ONLY: &
torch_dict_create, torch_dict_get, torch_dict_insert, torch_dict_release, torch_dict_type, &
torch_model_eval, torch_model_load, torch_model_read_metadata, torch_model_release, &
torch_model_type
#include "./base/base_uses.f90"
IMPLICIT NONE
CHARACTER(LEN=default_path_length) :: filename, cutoff_str, nequip_version
REAL(dp) :: cutoff
! Inputs.
INTEGER, PARAMETER :: natoms = 96
INTEGER :: nedges
REAL(sp), DIMENSION(3, natoms) :: pos
REAL(dp), DIMENSION(3, 3):: cell, hinv
INTEGER(kind=int_8), DIMENSION(natoms):: atom_types
INTEGER(kind=int_8), DIMENSION(:, :), ALLOCATABLE:: edge_index
REAL(sp), DIMENSION(:, :), ALLOCATABLE:: edge_cell_shift
! Torch objects.
TYPE(torch_model_type) :: model
TYPE(torch_dict_type) :: inputs, outputs
! Outputs.
REAL(sp), DIMENSION(:, :), POINTER :: total_energy, atomic_energy, forces
NULLIFY (total_energy, atomic_energy, forces)
! A box with 32 water molecules.
pos = RESHAPE([ &
42.886169670000001, -0.055681660000000001, 38.329161120000002, &
34.202588720000001, -0.6185484, 37.365568080000003, &
30.080392589999999, -2.0124176500000002, 36.480796079999998, &
28.705791179999999, -2.6880392799999999, 36.602098329999997, &
36.24794267, -0.51634849000000005, 34.492359610000001, &
37.696472460000003, -0.041087279999999997, 35.014073510000003, &
27.760669979999999, 7.4854206000000003, 33.927691950000003, &
28.816099959999999, 6.49857774, 34.21636084, &
37.157637200000003, 9.0188280800000005, 31.926581209999998, &
38.606381650000003, 9.5820079600000003, 32.343597279999997, &
34.303195930000001, 2.2195014400000002, 45.988045190000001, &
33.244413940000001, 1.30253325, 46.469842720000003, &
38.728617479999997, -5.0541897699999998, 26.074396839999999, &
38.348392150000002, -6.2832846900000003, 26.986725310000001, &
32.864252090000001, 3.20606327, 30.897116, &
31.290408809999999, 3.0871834699999998, 30.62739775, &
33.751986969999997, -3.13832624, 39.672760769999996, &
34.664297990000001, -3.66438591, 38.646602719999997, &
42.717321439999999, 5.1246883700000003, 32.588340119999998, &
41.56274552, 5.5893544000000004, 33.417490280000003, &
32.428380089999997, 9.1182520999999994, 30.54776786, &
32.6432407, 10.77068308, 30.48427787, &
31.484867080000001, 4.6777144699999997, 37.395719499999998, &
32.317188299999998, -6.2287496200000003, 36.467186439999999, &
26.662134099999999, 3.1708123800000001, 35.682014649999999, &
26.52713675, 1.60390403, 35.488348299999998, &
32.023823659999998, 16.918208029999999, 31.688356989999999, &
31.400657970000001, 7.0315610800000004, 30.239455469999999, &
33.52642531, -3.5594808100000002, 34.26368308, &
34.640485550000001, -3.26538336, 35.497148240000001, &
40.056437510000002, -0.30543864999999998, 29.831207429999999, &
39.478446419999997, -1.09483146, 38.310114040000002, &
39.704076149999999, 1.9584631400000001, 33.390237599999999, &
38.333857080000001, 2.69671781, 42.92619457, &
40.182045549999998, -7.2199289499999999, 27.658039049999999, &
39.320443179999998, -8.4564252700000004, 28.13196589, &
36.38769637, 8.8117085900000003, 38.354536240000002, &
36.320563759999999, 9.0063075799999996, 36.752600139999998, &
29.999158309999999, -5.5637817500000004, 33.929505059999997, &
30.772854550000002, -5.0385870199999996, 35.199806719999998, &
40.059251779999997, 6.3305279499999996, 28.257946189999998, &
40.239836089999997, 5.1745923999999999, 29.29629568, &
26.33209111, 2.4393638599999998, 33.565386850000003, &
26.960697119999999, 1.2711078899999999, 32.592388440000001, &
34.837269769999999, -0.47227084000000003, 30.382436200000001, &
35.396881370000003, -1.92684834, 30.308183759999999, &
32.121760719999997, -0.73334299000000003, 36.510438299999997, &
32.218084349999998, 7.8454304099999996, 35.667196779999998, &
36.37809987, -4.3048878799999999, 36.4539793, &
35.811927560000001, -3.00139282, 27.034893759999999, &
29.645249140000001, 1.0652123600000001, 35.714365399999998, &
30.379465499999998, -0.066814650000000003, 34.988246869999998, &
34.21493366, -1.65591205, 33.887643709999999, &
34.784243549999999, -1.0252141100000001, 32.503483260000003, &
40.464995450000004, 1.14678254, 31.307350320000001, &
41.326246990000001, 0.65508034999999998, 32.455588290000001, &
29.02108599, 3.5038194900000001, 39.90877029, &
29.494542689999999, 3.7276637300000002, 41.376613800000001, &
34.135966400000001, -6.7533422300000003, 32.356841029999998, &
34.954657009999998, -5.7704242399999996, 31.45710669, &
33.253235689999997, 1.5268048299999999, 44.056217189999998, &
33.793166939999999, 0.50146325999999997, 43.059759010000001, &
36.820540960000002, 2.6214681999999998, 40.683400659999997, &
37.555270669999999, 1.56498329, 39.76489351, &
43.209908720000001, -0.062845650000000003, 47.25931559, &
29.394058350000002, -2.3133019500000001, 37.140788360000002, &
36.741570840000001, -0.083871000000000001, 35.259178339999998, &
27.94247769, 6.7622961500000001, 34.564838440000003, &
37.681265600000003, 9.4216399800000001, 32.647864349999999, &
33.317129080000001, 2.0951401700000001, 45.872226509999997, &
37.995135589999997, 4.3611431200000004, 26.55718199, &
32.182467090000003, 2.6611503399999998, 30.457724819999999, &
34.653801209999997, -3.4374573100000001, 39.588924570000003, &
42.292983370000002, 5.9471069500000002, 32.846099549999998, &
32.960469009999997, 9.9050313200000009, 30.158730670000001, &
31.42818866, -5.8338304000000001, 36.673874359999999, &
26.05637308, 2.4973869199999998, 35.348687040000002, &
32.033492709999997, 17.325228939999999, 30.811601379999999, &
33.825218239999998, -2.9520949600000002, 35.022046070000002, &
39.456998159999998, -0.30727594000000003, 38.934782900000002, &
29.48467089, 2.8692561099999998, 43.006186839999998, &
39.286418410000003, -7.62061031, 27.627114779999999, &
35.879750280000003, 8.6515870800000005, 37.522173430000002, &
30.358254389999999, -4.7607656800000004, 34.335564570000003, &
40.709895690000003, 5.8331250199999998, 28.755837589999999, &
26.717908300000001, 2.2415138300000002, 32.657729799999998, &
35.658925619999998, -0.99689035000000004, 30.574953059999999, &
31.585160200000001, -1.31218042, 35.901110940000002, &
35.548938669999998, -3.9056138900000001, 26.821449000000001, &
29.565661609999999, 0.46817945999999999, 34.967071199999999, &
34.76151282, -0.95696802000000003, 33.489136729999998, &
40.485340669999999, 0.40236209000000001, 31.94254162, &
29.672828979999998, 4.0134825300000001, 40.450578, &
34.127228619999997, -5.8796882899999998, 31.892542299999999, &
33.116888420000002, 1.2338084899999999, 43.112771199999997, &
37.199699350000003, 2.5049007099999998, 39.791712660000002], shape=[3, natoms])
cell(1, :) = [9.85_dp, 0.0_dp, 0.0_dp]
cell(2, :) = [0.0_dp, 9.85_dp, 0.0_dp]
cell(3, :) = [0.0_dp, 0.0_dp, 9.85_dp]
hinv = inv_3x3(cell)
atom_types(:64) = 0 ! Hydrogen
atom_types(65:) = 1 ! Oxygen
#if defined(__LIBTORCH)
filename = discover_file('NequIP/water.pth')
WRITE (*, *) "Loading NequIP model from: "//TRIM(filename)
CALL torch_model_load(model, filename)
cutoff_str = torch_model_read_metadata(filename, "r_max")
nequip_version = torch_model_read_metadata(filename, "nequip_version")
READ (cutoff_str, *) cutoff
WRITE (*, *) "Version: ", TRIM(nequip_version)
WRITE (*, *) "Cutoff: ", cutoff
CALL neighbor_search(nedges)
ALLOCATE (edge_index(nedges, 2))
ALLOCATE (edge_cell_shift(3, nedges))
CALL neighbor_search(nedges, edge_index, edge_cell_shift)
WRITE (*, *) "Found", nedges, "neighbor edges between", natoms, "atoms."
CALL torch_dict_create(inputs)
CALL torch_dict_insert(inputs, "pos", pos)
CALL torch_dict_insert(inputs, "edge_index", edge_index)
CALL torch_dict_insert(inputs, "edge_cell_shift", edge_cell_shift)
CALL torch_dict_insert(inputs, "cell", REAL(cell, kind=sp))
CALL torch_dict_insert(inputs, "atom_types", atom_types)
CALL torch_dict_create(outputs)
CALL torch_model_eval(model, inputs, outputs)
CALL torch_dict_get(outputs, "total_energy", total_energy)
CALL torch_dict_get(outputs, "atomic_energy", atomic_energy)
CALL torch_dict_get(outputs, "forces", forces)
WRITE (*, *) "Total Energy: ", total_energy(1, 1)
CPASSERT(ABS(-14985.61621 - total_energy(1, 1)) < 1e-3)
CALL torch_dict_release(inputs)
CALL torch_dict_release(outputs)
CALL torch_model_release(model)
DEALLOCATE (edge_index, edge_cell_shift, total_energy, atomic_energy, forces)
WRITE (*, *) "NequIP unittest was successfully :-)"
#else
WRITE (*, *) "CP2K compiled without the Torch library - skipping NequIP unittest."
#endif
CONTAINS
! **************************************************************************************************
!> \brief Naive neighbor search - beware it scales O(N**2).
!> \param nedges ...
!> \param edge_index ...
!> \param edge_cell_shift ...
! **************************************************************************************************
SUBROUTINE neighbor_search(nedges, edge_index, edge_cell_shift)
INTEGER, INTENT(OUT) :: nedges
INTEGER(kind=int_8), DIMENSION(:, :), &
INTENT(OUT), OPTIONAL :: edge_index
REAL(sp), DIMENSION(:, :), INTENT(OUT), OPTIONAL :: edge_cell_shift
INTEGER:: iatom, jatom
REAL(dp), DIMENSION(3) :: s1, s2, s12, cell_shift, dx
nedges = 0
DO iatom = 1, natoms
DO jatom = 1, natoms
IF (iatom == jatom) CYCLE
s1 = MATMUL(hinv, pos(:, iatom))
s2 = MATMUL(hinv, pos(:, jatom))
s12 = s1 - s2
cell_shift = ANINT(s12)
dx = MATMUL(cell, s12 - cell_shift)
IF (DOT_PRODUCT(dx, dx) <= cutoff**2) THEN
nedges = nedges + 1
IF (PRESENT(edge_index)) THEN
edge_index(nedges, :) = [iatom - 1, jatom - 1]
END IF
IF (PRESENT(edge_cell_shift)) THEN
edge_cell_shift(:, nedges) = REAL(cell_shift, kind=sp)
END IF
END IF
END DO
END DO
END SUBROUTINE neighbor_search
END PROGRAM nequip_unittest

329
src/torch_api.F Normal file
View file

@ -0,0 +1,329 @@
!--------------------------------------------------------------------------------------------------!
! CP2K: A general program to perform molecular dynamics simulations !
! Copyright 2000-2022 CP2K developers group <https://cp2k.org> !
! !
! SPDX-License-Identifier: GPL-2.0-or-later !
!--------------------------------------------------------------------------------------------------!
MODULE torch_api
USE ISO_C_BINDING, ONLY: C_ASSOCIATED, &
C_CHAR, &
C_FLOAT, &
C_F_POINTER, &
C_INT, &
C_NULL_CHAR, &
C_NULL_PTR, &
C_PTR, &
C_INT64_T
USE kinds, ONLY: sp, int_8
#include "./base/base_uses.f90"
IMPLICIT NONE
PRIVATE
TYPE torch_dict_type
PRIVATE
TYPE(C_PTR) :: c_ptr = C_NULL_PTR
END TYPE torch_dict_type
TYPE torch_model_type
PRIVATE
TYPE(C_PTR) :: c_ptr = C_NULL_PTR
END TYPE torch_model_type
#:set max_dim = 3
INTERFACE torch_dict_insert
#:for ndims in range(1, max_dim+1)
MODULE PROCEDURE torch_dict_insert_float_${ndims}$d
MODULE PROCEDURE torch_dict_insert_int64_${ndims}$d
#:endfor
END INTERFACE torch_dict_insert
INTERFACE torch_dict_get
#:for ndims in range(1, max_dim+1)
MODULE PROCEDURE torch_dict_get_float_${ndims}$d
MODULE PROCEDURE torch_dict_get_int64_${ndims}$d
#:endfor
END INTERFACE torch_dict_get
PUBLIC :: torch_dict_type, torch_dict_create, torch_dict_release
PUBLIC :: torch_dict_insert, torch_dict_get
PUBLIC :: torch_model_type, torch_model_load, torch_model_eval, torch_model_release
PUBLIC :: torch_model_read_metadata
CONTAINS
#:set typenames = ['float', 'int64']
#:set types_f = ['REAL(sp)','INTEGER(kind=int_8)']
#:set types_c = ['REAL(kind=C_FLOAT)','INTEGER(kind=C_INT64_T)']
#:for ndims in range(1, max_dim+1)
#:for typename, type_f, type_c in zip(typenames, types_f, types_c)
! **************************************************************************************************
!> \brief Inserts array into Torch dictionary. The passed array has to outlive the dictionary!
!> \author Ole Schuett
! **************************************************************************************************
SUBROUTINE torch_dict_insert_${typename}$_${ndims}$d(dict, key, source)
TYPE(torch_dict_type), INTENT(INOUT) :: dict
CHARACTER(len=*), INTENT(IN) :: key
#:set arraydims = ", ".join(":" for i in range(ndims))
${type_f}$, CONTIGUOUS, DIMENSION(${arraydims}$), INTENT(IN) :: source
#if defined(__LIBTORCH)
INTEGER(kind=int_8), DIMENSION(${ndims}$) :: sizes_c
INTERFACE
SUBROUTINE torch_c_dict_insert_${typename}$ (dict, key, ndims, sizes, source) &
BIND(C, name="torch_c_dict_insert_${typename}$")
IMPORT :: C_CHAR, C_PTR, C_INT, C_INT64_T, C_FLOAT
TYPE(C_PTR), VALUE :: dict
CHARACTER(kind=C_CHAR), DIMENSION(*) :: key
INTEGER(kind=C_INT), VALUE :: ndims
INTEGER(kind=C_INT64_T), DIMENSION(*) :: sizes
${type_c}$, DIMENSION(*) :: source
END SUBROUTINE torch_c_dict_insert_${typename}$
END INTERFACE
#:for axis in range(ndims)
sizes_c(${axis + 1}$) = SIZE(source, ${ndims - axis}$) ! C arrays are stored row-major.
#:endfor
CPASSERT(C_ASSOCIATED(dict%c_ptr))
CALL torch_c_dict_insert_${typename}$ (dict=dict%c_ptr, &
key=TRIM(key)//C_NULL_CHAR, &
ndims=${ndims}$, &
sizes=sizes_c, &
source=source)
#else
CPABORT("CP2K compiled without the Torch library.")
#endif
END SUBROUTINE torch_dict_insert_${typename}$_${ndims}$d
! **************************************************************************************************
!> \brief Retrieves array from Torch dictionary. The returned array has to deallocated by caller!
!> \author Ole Schuett
! **************************************************************************************************
SUBROUTINE torch_dict_get_${typename}$_${ndims}$d(dict, key, dest)
TYPE(torch_dict_type), INTENT(IN) :: dict
CHARACTER(len=*), INTENT(IN) :: key
#:set arraydims = ", ".join(":" for i in range(ndims))
${type_f}$, DIMENSION(${arraydims}$), POINTER :: dest
#if defined(__LIBTORCH)
INTEGER(kind=int_8), DIMENSION(${ndims}$) :: sizes_f, sizes_c
TYPE(C_PTR) :: dest_c
INTERFACE
SUBROUTINE torch_c_dict_get_${typename}$ (dict, key, ndims, sizes, dest) &
BIND(C, name="torch_c_dict_get_${typename}$")
IMPORT :: C_CHAR, C_PTR, C_INT, C_INT64_T
TYPE(C_PTR), VALUE :: dict
CHARACTER(kind=C_CHAR), DIMENSION(*) :: key
INTEGER(kind=C_INT), VALUE :: ndims
INTEGER(kind=C_INT64_T), DIMENSION(*) :: sizes
TYPE(C_PTR) :: dest
END SUBROUTINE torch_c_dict_get_${typename}$
END INTERFACE
sizes_c(:) = -1
dest_c = C_NULL_PTR
CPASSERT(C_ASSOCIATED(dict%c_ptr))
CPASSERT(.NOT. ASSOCIATED(dest))
CALL torch_c_dict_get_${typename}$ (dict=dict%c_ptr, &
key=TRIM(key)//C_NULL_CHAR, &
ndims=${ndims}$, &
sizes=sizes_c, &
dest=dest_c)
CPASSERT(ALL(sizes_c >= 0))
CPASSERT(C_ASSOCIATED(dest_c))
#:for axis in range(ndims)
sizes_f(${axis + 1}$) = sizes_c(${ndims - axis}$) ! C arrays are stored row-major.
#:endfor
CALL C_F_POINTER(dest_c, dest, shape=sizes_f)
#else
CPABORT("CP2K compiled without the Torch library.")
#endif
END SUBROUTINE torch_dict_get_${typename}$_${ndims}$d
#:endfor
#:endfor
! **************************************************************************************************
!> \brief Creates an empty Torch dictionary.
!> \author Ole Schuett
! **************************************************************************************************
SUBROUTINE torch_dict_create(dict)
TYPE(torch_dict_type), INTENT(INOUT) :: dict
#if defined(__LIBTORCH)
INTERFACE
SUBROUTINE torch_c_dict_create(dict) BIND(C, name="torch_c_dict_create")
IMPORT :: C_PTR
TYPE(C_PTR) :: dict
END SUBROUTINE torch_c_dict_create
END INTERFACE
CPASSERT(.NOT. C_ASSOCIATED(dict%c_ptr))
CALL torch_c_dict_create(dict=dict%c_ptr)
CPASSERT(C_ASSOCIATED(dict%c_ptr))
#else
CPABORT("CP2K was compiled without Torch library.")
#endif
END SUBROUTINE torch_dict_create
! **************************************************************************************************
!> \brief Releases a Torch dictionary and all its ressources.
!> \author Ole Schuett
! **************************************************************************************************
SUBROUTINE torch_dict_release(dict)
TYPE(torch_dict_type), INTENT(INOUT) :: dict
#if defined(__LIBTORCH)
INTERFACE
SUBROUTINE torch_c_dict_release(dict) BIND(C, name="torch_c_dict_release")
IMPORT :: C_PTR
TYPE(C_PTR), VALUE :: dict
END SUBROUTINE torch_c_dict_release
END INTERFACE
CPASSERT(C_ASSOCIATED(dict%c_ptr))
CALL torch_c_dict_release(dict=dict%c_ptr)
dict%c_ptr = C_NULL_PTR
#else
CPABORT("CP2K was compiled without Torch library.")
#endif
END SUBROUTINE torch_dict_release
! **************************************************************************************************
!> \brief Loads a Torch model from given "*.pth" file. (In Torch lingo models are called modules)
!> \author Ole Schuett
! **************************************************************************************************
SUBROUTINE torch_model_load(model, filename)
TYPE(torch_model_type), INTENT(INOUT) :: model
CHARACTER(len=*), INTENT(IN) :: filename
#if defined(__LIBTORCH)
INTERFACE
SUBROUTINE torch_c_model_load(model, filename) BIND(C, name="torch_c_model_load")
IMPORT :: C_PTR, C_CHAR
TYPE(C_PTR) :: model
CHARACTER(kind=C_CHAR), DIMENSION(*) :: filename
END SUBROUTINE torch_c_model_load
END INTERFACE
CPASSERT(.NOT. C_ASSOCIATED(model%c_ptr))
CALL torch_c_model_load(model=model%c_ptr, filename=TRIM(filename)//C_NULL_CHAR)
CPASSERT(C_ASSOCIATED(model%c_ptr))
#else
CPABORT("CP2K was compiled without Torch library.")
#endif
END SUBROUTINE torch_model_load
! **************************************************************************************************
!> \brief Evaluates the given Torch model. (In Torch lingo this operation is called forward())
!> \author Ole Schuett
! **************************************************************************************************
SUBROUTINE torch_model_eval(model, inputs, outputs)
TYPE(torch_model_type), INTENT(INOUT) :: model
TYPE(torch_dict_type), INTENT(IN) :: inputs
TYPE(torch_dict_type), INTENT(INOUT) :: outputs
#if defined(__LIBTORCH)
INTERFACE
SUBROUTINE torch_c_model_eval(model, inputs, outputs) BIND(C, name="torch_c_model_eval")
IMPORT :: C_PTR
TYPE(C_PTR), VALUE :: model
TYPE(C_PTR), VALUE :: inputs
TYPE(C_PTR), VALUE :: outputs
END SUBROUTINE torch_c_model_eval
END INTERFACE
CPASSERT(C_ASSOCIATED(model%c_ptr))
CPASSERT(C_ASSOCIATED(inputs%c_ptr))
CPASSERT(C_ASSOCIATED(outputs%c_ptr))
CALL torch_c_model_eval(model=model%c_ptr, &
inputs=inputs%c_ptr, &
outputs=outputs%c_ptr)
#else
CPABORT("CP2K was compiled without Torch library.")
#endif
END SUBROUTINE torch_model_eval
! **************************************************************************************************
!> \brief Releases a Torch model and all its ressources.
!> \author Ole Schuett
! **************************************************************************************************
SUBROUTINE torch_model_release(model)
TYPE(torch_model_type), INTENT(INOUT) :: model
#if defined(__LIBTORCH)
INTERFACE
SUBROUTINE torch_c_model_release(model) BIND(C, name="torch_c_model_release")
IMPORT :: C_PTR
TYPE(C_PTR), VALUE :: model
END SUBROUTINE torch_c_model_release
END INTERFACE
CPASSERT(C_ASSOCIATED(model%c_ptr))
CALL torch_c_model_release(model=model%c_ptr)
model%c_ptr = C_NULL_PTR
#else
CPABORT("CP2K was compiled without Torch library.")
#endif
END SUBROUTINE torch_model_release
! **************************************************************************************************
!> \brief Reads metadata entry from given "*.pth" file. (In Torch lingo they are called extra files)
!> \author Ole Schuett
! **************************************************************************************************
FUNCTION torch_model_read_metadata(filename, key) RESULT(res)
CHARACTER(len=*), INTENT(IN) :: filename, key
CHARACTER(:), ALLOCATABLE :: res
#if defined(__LIBTORCH)
CHARACTER(LEN=1, KIND=C_CHAR), DIMENSION(:), &
POINTER :: content_f
INTEGER :: i
INTEGER :: length
TYPE(C_PTR) :: content_c
INTERFACE
SUBROUTINE torch_c_model_read_metadata(filename, key, content, length) &
BIND(C, name="torch_c_model_read_metadata")
IMPORT :: C_CHAR, C_PTR, C_INT
CHARACTER(kind=C_CHAR), DIMENSION(*) :: filename, key
TYPE(C_PTR) :: content
INTEGER(kind=C_INT) :: length
END SUBROUTINE torch_c_model_read_metadata
END INTERFACE
content_c = C_NULL_PTR
length = -1
CALL torch_c_model_read_metadata(filename=TRIM(filename)//C_NULL_CHAR, &
key=TRIM(key)//C_NULL_CHAR, &
content=content_c, &
length=length)
CPASSERT(C_ASSOCIATED(content_c))
CPASSERT(length >= 0)
CALL C_F_POINTER(content_c, content_f, shape=(/length + 1/))
CPASSERT(content_f(length + 1) == C_NULL_CHAR)
ALLOCATE (CHARACTER(LEN=length) :: res)
DO i = 1, length
CPASSERT(content_f(i) /= C_NULL_CHAR)
res(i:i) = content_f(i)
END DO
DEALLOCATE (content_f) ! Was allocated on the C side.
#else
CPABORT("CP2K was compiled without Torch library.")
#endif
END FUNCTION torch_model_read_metadata
END MODULE torch_api

169
src/torch_c_api.cpp Normal file
View file

@ -0,0 +1,169 @@
/*----------------------------------------------------------------------------*/
/* CP2K: A general program to perform molecular dynamics simulations */
/* Copyright 2000-2022 CP2K developers group <https://cp2k.org> */
/* */
/* SPDX-License-Identifier: GPL-2.0-or-later */
/*----------------------------------------------------------------------------*/
#if defined(__LIBTORCH)
#include <torch/script.h>
typedef c10::Dict<std::string, torch::Tensor> torch_c_dict_t;
typedef torch::jit::Module torch_c_model_t;
/*******************************************************************************
* \brief Internal helper for retrieving arrays from Torch dictionary.
* \author Ole Schuett
******************************************************************************/
template <typename T>
static void torch_c_dict_get(const torch_c_dict_t *dict, const char *key,
const int ndims, int64_t sizes[], T **dest) {
assert(dict->contains(key));
const torch::Tensor tensor = dict->at(key);
assert(tensor.ndimension() == ndims);
int64_t size_flat = 1;
for (int i = 0; i < ndims; i++) {
sizes[i] = tensor.size(i);
size_flat *= sizes[i];
}
*dest = (T *)malloc(size_flat * sizeof(T));
const torch::Tensor tensor_flat = tensor.flatten();
const auto accessor = tensor_flat.accessor<T, 1>();
for (int i = 0; i < size_flat; i++) {
(*dest)[i] = accessor[i];
}
};
#ifdef __cplusplus
extern "C" {
#endif
/*******************************************************************************
* \brief Inserts array of floats into Torch dictionary.
* The passed array has to outlive the dictionary!
* \author Ole Schuett
******************************************************************************/
void torch_c_dict_insert_float(torch_c_dict_t *dict, const char *key,
const int ndims, const int64_t sizes[],
float source[]) {
const auto options = torch::TensorOptions().dtype(torch::kFloat32);
const auto sizes_ref = c10::IntArrayRef(sizes, ndims);
const torch::Tensor tensor = torch::from_blob(source, sizes_ref, options);
dict->insert(key, tensor);
}
/*******************************************************************************
* \brief Inserts array of int64s into Torch dictionary.
* The passed array has to outlive the dictionary!
* \author Ole Schuett
******************************************************************************/
void torch_c_dict_insert_int64(torch_c_dict_t *dict, const char *key,
const int ndims, const int64_t sizes[],
int64_t source[]) {
const auto options = torch::TensorOptions().dtype(torch::kInt64);
const auto sizes_ref = c10::IntArrayRef(sizes, ndims);
const torch::Tensor tensor = torch::from_blob(source, sizes_ref, options);
dict->insert(key, tensor);
}
/*******************************************************************************
* \brief Retrieves array of floats from Torch dictionary.
* The returned array has to be deallocated by caller!
* \author Ole Schuett
******************************************************************************/
void torch_c_dict_get_float(const torch_c_dict_t *dict, const char *key,
const int ndims, int64_t sizes[], float **dest) {
torch_c_dict_get<float>(dict, key, ndims, sizes, dest);
}
/*******************************************************************************
* \brief Retrieves array of int64s from Torch dictionary.
* The returned array has to be deallocated by caller!
* \author Ole Schuett
******************************************************************************/
void torch_c_dict_get_int64(const torch_c_dict_t *dict, const char *key,
const int ndims, int64_t sizes[], int64_t **dest) {
torch_c_dict_get<int64_t>(dict, key, ndims, sizes, dest);
}
/*******************************************************************************
* \brief Creates an empty Torch dictionary.
* \author Ole Schuett
******************************************************************************/
void torch_c_dict_create(torch_c_dict_t **dict_out) {
assert(*dict_out == NULL);
*dict_out = new c10::Dict<std::string, torch::Tensor>();
}
/*******************************************************************************
* \brief Releases a Torch dictionary and all its ressources.
* \author Ole Schuett
******************************************************************************/
void torch_c_dict_release(torch_c_dict_t *dict) { delete (dict); }
/*******************************************************************************
* \brief Loads a Torch model from given "*.pth" file.
* In Torch lingo models are called modules.
* \author Ole Schuett
******************************************************************************/
void torch_c_model_load(torch_c_model_t **model_out, const char *filename) {
torch::jit::Module *model = new torch::jit::Module();
*model = torch::jit::load(filename);
model->eval();
assert(*model_out == NULL);
*model_out = model;
}
/*******************************************************************************
* \brief Evaluates the given Torch model.
* In Torch lingo this operation is called forward().
* \author Ole Schuett
******************************************************************************/
void torch_c_model_eval(torch_c_model_t *model, const torch_c_dict_t *inputs,
torch_c_dict_t *outputs) {
auto untyped_output = model->forward({*inputs}).toGenericDict();
outputs->clear();
for (const auto &entry : untyped_output) {
outputs->insert(entry.key().toStringView(), entry.value().toTensor());
}
}
/*******************************************************************************
* \brief Releases a Torch model and all its ressources.
* \author Ole Schuett
******************************************************************************/
void torch_c_model_release(torch_c_model_t *model) { delete (model); }
/*******************************************************************************
* \brief Reads metadata entry from given "*.pth" file.
* In Torch lingo they are called extra files.
* The returned char array has to be deallocated by caller!
* \author Ole Schuett
******************************************************************************/
void torch_c_model_read_metadata(const char *filename, const char *key,
char **content, int *length) {
std::unordered_map<std::string, std::string> extra_files = {{key, ""}};
torch::jit::load(filename, torch::kCPU, extra_files);
const std::string &content_str = extra_files[key];
*length = content_str.length();
*content = (char *)malloc(content_str.length() * sizeof(char));
strcpy(*content, content_str.c_str());
}
#ifdef __cplusplus
}
#endif
#endif // defined(__LIBTORCH)
// EOF

6
tools/precommit/precommit.py
View file

@ -193,7 +193,11 @@ def process_file(fn: str, allow_modifications: bool) -> None:
run_remote_tool("clangformat", fn)
if re.match(r".*\.(cc|cpp|cxx|hcc|hpp|hxx)$", fn):
raise Exception(f"C++ is not supported.")
if fn.endswith("/torch_c_api.cpp"):
# Begrudgingly tolerated because PyTorch has no C API.
run_remote_tool("clangformat", fn)
else:
raise Exception(f"C++ is not supported.")
if re.match(r"(.*/PACKAGE)|(.*\.py)$", fn):
ast.parse(orig_content, filename=fn)

2
tools/toolchain/scripts/generate_arch_files.sh
View file

@ -118,7 +118,7 @@ LDFLAGS="IF_STATIC(${STATIC_FLAGS}|) \$(FCFLAGS) ${CP_LDFLAGS}"
# add standard libs
LIBS="${CP_LIBS} -lstdc++"
CXXFLAGS+=" --std=c++11 \$(DFLAGS) -Wno-deprecated-declarations"
CXXFLAGS+=" --std=c++14 \$(DFLAGS) -Wno-deprecated-declarations"
# CUDA handling
if [ "${ENABLE_CUDA}" = __TRUE__ ] && [ "${GPUVER}" != no ]; then
CUDA_LIBS="-lcudart -lnvrtc -lcuda -lcufft -lcublas -lrt IF_DEBUG(-lnvToolsExt|)"

96
tools/toolchain/scripts/stage7/install_libtorch.sh Executable file
View file

@ -0,0 +1,96 @@
#!/bin/bash -e
# TODO: Review and if possible fix shellcheck errors.
# shellcheck disable=all
[ "${BASH_SOURCE[0]}" ] && SCRIPT_NAME="${BASH_SOURCE[0]}" || SCRIPT_NAME=$0
SCRIPT_DIR="$(cd "$(dirname "$SCRIPT_NAME")/.." && pwd -P)"
# From https://pytorch.org/get-started/locally/
libtorch_ver="1.12.1"
libtorch_sha256="82c7be80860f2aa7963f8700004a40af8205e1d721298f2e09b700e766a9d283"
# shellcheck source=/dev/null
source "${SCRIPT_DIR}"/common_vars.sh
source "${SCRIPT_DIR}"/tool_kit.sh
source "${SCRIPT_DIR}"/signal_trap.sh
source "${INSTALLDIR}"/toolchain.conf
source "${INSTALLDIR}"/toolchain.env
[ -f "${BUILDDIR}/setup_libtorch" ] && rm "${BUILDDIR}/setup_libtorch"
! [ -d "${BUILDDIR}" ] && mkdir -p "${BUILDDIR}"
cd "${BUILDDIR}"
case "${with_libtorch:=__INSTALL__}" in
__INSTALL__)
echo "==================== Installing libtorch ===================="
pkg_install_dir="${INSTALLDIR}/libtorch-${libtorch_ver}"
install_lock_file="${pkg_install_dir}/install_successful"
archive_file="libtorch-cxx11-abi-shared-with-deps-${libtorch_ver}+cpu.zip"
if verify_checksums "${install_lock_file}"; then
echo "libtorch-${libtorch_ver} is already installed, skipping it."
else
if [ -f ${archive_file} ]; then
echo "${archive_file} is found"
else
download_pkg_from_cp2k_org "${libtorch_sha256}" "${archive_file}"
fi
echo "Installing from scratch into ${pkg_install_dir}"
[ -d libtorch ] && rm -rf libtorch
[ -d ${pkg_install_dir} ] && rm -rf ${pkg_install_dir}
unzip -q ${archive_file}
mv libtorch ${pkg_install_dir}
write_checksums "${install_lock_file}" "${SCRIPT_DIR}/stage7/$(basename "${SCRIPT_NAME}")"
fi
LIBTORCH_CXXFLAGS="-I${pkg_install_dir}/include"
LIBTORCH_LDFLAGS="-L'${pkg_install_dir}/lib' -Wl,-rpath='${pkg_install_dir}/lib'"
;;
__SYSTEM__)
echo "==================== Finding libtorch from system paths ===================="
check_lib -ltorch "libtorch"
add_include_from_paths LIBTORCH_CXXFLAGS "libtorch.h" $INCLUDE_PATHS
add_lib_from_paths LIBTORCH_LDFLAGS "libtorch.*" "$LIB_PATHS"
;;
__DONTUSE__) ;;
*)
echo "==================== Linking libtorch to user paths ===================="
pkg_install_dir="${with_libtorch}"
# use the lib64 directory if present (multi-abi distros may link lib/ to lib32/ instead)
LIBTORCH_LIBDIR="${pkg_install_dir}/lib"
[ -d "${pkg_install_dir}/lib64" ] && LIBTORCH_LIBDIR="${pkg_install_dir}/lib64"
check_dir "${LIBTORCH_LIBDIR}"
LIBTORCH_CXXFLAGS="-I${pkg_install_dir}/include"
LIBTORCH_LDFLAGS="-L'${LIBTORCH_LIBDIR}' -Wl,-rpath='${LIBTORCH_LIBDIR}'"
;;
esac
if [ "$with_libtorch" != "__DONTUSE__" ]; then
if [ "$with_libtorch" != "__SYSTEM__" ]; then
cat << EOF > "${BUILDDIR}/setup_libtorch"
prepend_path LD_LIBRARY_PATH "${pkg_install_dir}/lib"
prepend_path LD_RUN_PATH "${pkg_install_dir}/lib"
prepend_path LIBRARY_PATH "${pkg_install_dir}/lib"
EOF
fi
cat << EOF >> "${BUILDDIR}/setup_libtorch"
export CP_DFLAGS="\${CP_DFLAGS} -D__LIBTORCH"
export CXXFLAGS="\${CXXFLAGS} ${LIBTORCH_CXXFLAGS}"
export CP_LDFLAGS="\${CP_LDFLAGS} ${LIBTORCH_LDFLAGS}"
export CP_LIBS="\${CP_LIBS} -lc10 -ltorch_cpu -ltorch"
EOF
cat "${BUILDDIR}/setup_libtorch" >> "${SETUPFILE}"
fi
load "${BUILDDIR}/setup_libtorch"
write_toolchain_env "${INSTALLDIR}"
cd "${ROOTDIR}"
report_timing "libtorch"

1
tools/toolchain/scripts/stage7/install_stage7.sh
View file

@ -7,5 +7,6 @@
./scripts/stage7/install_libvdwxc.sh
./scripts/stage7/install_spglib.sh
./scripts/stage7/install_libvori.sh
./scripts/stage7/install_libtorch.sh
#EOF