Line data Source code
1 : !--------------------------------------------------------------------------------------------------!
2 : ! CP2K: A general program to perform molecular dynamics simulations !
3 : ! Copyright 2000-2026 CP2K developers group <https://cp2k.org> !
4 : ! !
5 : ! SPDX-License-Identifier: GPL-2.0-or-later !
6 : !--------------------------------------------------------------------------------------------------!
7 : MODULE torch_api
8 : USE ISO_C_BINDING, ONLY: C_ASSOCIATED, &
9 : C_BOOL, &
10 : C_CHAR, &
11 : C_FLOAT, &
12 : C_DOUBLE, &
13 : C_F_POINTER, &
14 : C_INT, &
15 : C_NULL_CHAR, &
16 : C_NULL_PTR, &
17 : C_PTR, &
18 : C_INT32_T, &
19 : C_INT64_T
20 :
21 : USE kinds, ONLY: sp, int_4, int_8, dp, default_string_length
22 :
23 : #include "./base/base_uses.f90"
24 :
25 : IMPLICIT NONE
26 :
27 : PRIVATE
28 :
29 : TYPE torch_tensor_type
30 : PRIVATE
31 : TYPE(C_PTR) :: c_ptr = C_NULL_PTR
32 : END TYPE torch_tensor_type
33 :
34 : TYPE torch_dict_type
35 : PRIVATE
36 : TYPE(C_PTR) :: c_ptr = C_NULL_PTR
37 : END TYPE torch_dict_type
38 :
39 : TYPE torch_model_type
40 : PRIVATE
41 : TYPE(C_PTR) :: c_ptr = C_NULL_PTR
42 : END TYPE torch_model_type
43 :
44 : #:set max_dim = 3
45 : INTERFACE torch_tensor_from_array
46 : #:for ndims in range(1, max_dim+1)
47 : MODULE PROCEDURE torch_tensor_from_array_int32_${ndims}$d
48 : MODULE PROCEDURE torch_tensor_from_array_float_${ndims}$d
49 : MODULE PROCEDURE torch_tensor_from_array_int64_${ndims}$d
50 : MODULE PROCEDURE torch_tensor_from_array_double_${ndims}$d
51 : #:endfor
52 : END INTERFACE torch_tensor_from_array
53 :
54 : INTERFACE torch_tensor_reset_from_array
55 : #:for ndims in range(1, max_dim+1)
56 : MODULE PROCEDURE torch_tensor_reset_from_array_double_${ndims}$d
57 : #:endfor
58 : END INTERFACE torch_tensor_reset_from_array
59 :
60 : INTERFACE torch_tensor_data_ptr
61 : #:for ndims in range(1, max_dim+1)
62 : MODULE PROCEDURE torch_tensor_data_ptr_int32_${ndims}$d
63 : MODULE PROCEDURE torch_tensor_data_ptr_float_${ndims}$d
64 : MODULE PROCEDURE torch_tensor_data_ptr_int64_${ndims}$d
65 : MODULE PROCEDURE torch_tensor_data_ptr_double_${ndims}$d
66 : #:endfor
67 : END INTERFACE torch_tensor_data_ptr
68 :
69 : INTERFACE torch_model_get_attr
70 : MODULE PROCEDURE torch_model_get_attr_string
71 : MODULE PROCEDURE torch_model_get_attr_double
72 : MODULE PROCEDURE torch_model_get_attr_int64
73 : MODULE PROCEDURE torch_model_get_attr_int32
74 : MODULE PROCEDURE torch_model_get_attr_strlist
75 : END INTERFACE torch_model_get_attr
76 :
77 : PUBLIC :: torch_tensor_type, torch_tensor_expand_dim, torch_tensor_from_array, &
78 : torch_tensor_narrow, torch_tensor_release
79 : PUBLIC :: torch_tensor_reset_from_array
80 : PUBLIC :: torch_tensor_data_ptr, torch_tensor_backward, torch_tensor_backward_scalar
81 : PUBLIC :: torch_tensor_grad, torch_tensor_grad_batch3
82 : PUBLIC :: torch_tensor_to_device_leaf
83 : PUBLIC :: torch_tensor_item_double, torch_tensor_weighted_sum
84 : PUBLIC :: torch_dict_type, torch_dict_clone, torch_dict_create, torch_dict_insert
85 : PUBLIC :: torch_dict_get, torch_dict_release
86 : PUBLIC :: torch_model_type, torch_model_load, torch_model_load_with_metadata, &
87 : torch_model_forward, torch_model_release
88 : PUBLIC :: torch_model_disable_parameter_gradients, torch_model_forward_mol_tensor, &
89 : torch_model_remap_device_constants
90 : PUBLIC :: torch_model_get_attr, torch_model_read_metadata
91 : PUBLIC :: torch_cuda_device_count, torch_cuda_is_available
92 : PUBLIC :: torch_allow_tf32, torch_model_freeze, torch_use_cuda
93 :
94 : CONTAINS
95 :
96 : #:set typenames = ['int32', 'float', 'int64', 'double']
97 : #:set types_f = ['INTEGER(kind=int_4)', 'REAL(sp)', 'INTEGER(kind=int_8)', 'REAL(dp)']
98 : #:set types_c = ['INTEGER(kind=C_INT32_T)', 'REAL(kind=C_FLOAT)', 'INTEGER(kind=C_INT64_T)', 'REAL(kind=C_DOUBLE)']
99 :
100 : #:for ndims in range(1, max_dim+1)
101 : #:for typename, type_f, type_c in zip(typenames, types_f, types_c)
102 :
103 : ! **************************************************************************************************
104 : !> \brief Creates a Torch tensor from an array. The passed array has to outlive the tensor!
105 : !> The source must be an ALLOCATABLE to prevent passing a temporary array.
106 : !> \author Ole Schuett
107 : ! **************************************************************************************************
108 2442 : SUBROUTINE torch_tensor_from_array_${typename}$_${ndims}$d(tensor, source, requires_grad)
109 : TYPE(torch_tensor_type), INTENT(INOUT) :: tensor
110 : #:set arraydims = ", ".join(":" for i in range(ndims))
111 : ${type_f}$, DIMENSION(${arraydims}$), ALLOCATABLE, INTENT(IN) :: source
112 : LOGICAL, OPTIONAL, INTENT(IN) :: requires_grad
113 :
114 : #if defined(__LIBTORCH)
115 : INTEGER(kind=int_8), DIMENSION(${ndims}$) :: sizes_c
116 : LOGICAL :: my_req_grad
117 :
118 : INTERFACE
119 : SUBROUTINE torch_c_tensor_from_array_${typename}$ (tensor, req_grad, ndims, sizes, source) &
120 : BIND(C, name="torch_c_tensor_from_array_${typename}$")
121 : IMPORT :: C_PTR, C_INT, C_INT32_T, C_INT64_T, C_FLOAT, C_DOUBLE, C_BOOL
122 : TYPE(C_PTR) :: tensor
123 : LOGICAL(kind=C_BOOL), VALUE :: req_grad
124 : INTEGER(kind=C_INT), VALUE :: ndims
125 : INTEGER(kind=C_INT64_T), DIMENSION(*) :: sizes
126 : ${type_c}$, DIMENSION(*) :: source
127 : END SUBROUTINE torch_c_tensor_from_array_${typename}$
128 : END INTERFACE
129 :
130 2442 : my_req_grad = .FALSE.
131 2442 : IF (PRESENT(requires_grad)) my_req_grad = requires_grad
132 :
133 : #:for axis in range(ndims)
134 2442 : sizes_c(${axis + 1}$) = SIZE(source, ${ndims - axis}$) ! C arrays are stored row-major.
135 : #:endfor
136 :
137 2442 : CPASSERT(.NOT. C_ASSOCIATED(tensor%c_ptr))
138 : CALL torch_c_tensor_from_array_${typename}$ (tensor=tensor%c_ptr, &
139 : req_grad=LOGICAL(my_req_grad, C_BOOL), &
140 : ndims=${ndims}$, &
141 : sizes=sizes_c, &
142 2442 : source=source)
143 2442 : CPASSERT(C_ASSOCIATED(tensor%c_ptr))
144 : #else
145 : CPABORT("CP2K compiled without the Torch library.")
146 : MARK_USED(tensor)
147 : MARK_USED(source)
148 : MARK_USED(requires_grad)
149 : #endif
150 2442 : END SUBROUTINE torch_tensor_from_array_${typename}$_${ndims}$d
151 :
152 : ! **************************************************************************************************
153 : !> \brief Copies data from a Torch tensor to an array.
154 : !> The returned pointer is only valide during the tensor's lifetime!
155 : !> \author Ole Schuett
156 : ! **************************************************************************************************
157 1730 : SUBROUTINE torch_tensor_data_ptr_${typename}$_${ndims}$d(tensor, data_ptr)
158 : TYPE(torch_tensor_type), INTENT(IN) :: tensor
159 : #:set arraydims = ", ".join(":" for i in range(ndims))
160 : ${type_f}$, DIMENSION(${arraydims}$), POINTER :: data_ptr
161 :
162 : #if defined(__LIBTORCH)
163 : INTEGER(kind=int_8), DIMENSION(${ndims}$) :: sizes_f, sizes_c
164 : TYPE(C_PTR) :: data_ptr_c
165 :
166 : INTERFACE
167 : SUBROUTINE torch_c_tensor_data_ptr_${typename}$ (tensor, ndims, sizes, data_ptr) &
168 : BIND(C, name="torch_c_tensor_data_ptr_${typename}$")
169 : IMPORT :: C_CHAR, C_PTR, C_INT, C_INT32_T, C_INT64_T
170 : TYPE(C_PTR), VALUE :: tensor
171 : INTEGER(kind=C_INT), VALUE :: ndims
172 : INTEGER(kind=C_INT64_T), DIMENSION(*) :: sizes
173 : TYPE(C_PTR) :: data_ptr
174 : END SUBROUTINE torch_c_tensor_data_ptr_${typename}$
175 : END INTERFACE
176 :
177 5258 : sizes_c(:) = -1
178 1730 : data_ptr_c = C_NULL_PTR
179 1730 : CPASSERT(C_ASSOCIATED(tensor%c_ptr))
180 1730 : CPASSERT(.NOT. ASSOCIATED(data_ptr))
181 : CALL torch_c_tensor_data_ptr_${typename}$ (tensor=tensor%c_ptr, &
182 : ndims=${ndims}$, &
183 : sizes=sizes_c, &
184 1730 : data_ptr=data_ptr_c)
185 :
186 : #:for axis in range(ndims)
187 1730 : sizes_f(${axis + 1}$) = sizes_c(${ndims - axis}$) ! C arrays are stored row-major.
188 : #:endfor
189 :
190 5258 : IF (ALL(sizes_f /= 0)) THEN ! Torch returns null pointer for zero-sized tensors.
191 1730 : CPASSERT(C_ASSOCIATED(data_ptr_c))
192 5258 : CALL C_F_POINTER(data_ptr_c, data_ptr, shape=sizes_f)
193 : END IF
194 : #else
195 : CPABORT("CP2K compiled without the Torch library.")
196 : MARK_USED(tensor)
197 : MARK_USED(data_ptr)
198 : #endif
199 1730 : END SUBROUTINE torch_tensor_data_ptr_${typename}$_${ndims}$d
200 :
201 : #:endfor
202 : #:endfor
203 :
204 : #:for ndims in range(1, max_dim+1)
205 :
206 : ! **************************************************************************************************
207 : !> \brief Reuses or creates a device leaf tensor and copies data into it.
208 : !> The source must be an ALLOCATABLE to prevent passing a temporary array.
209 : ! **************************************************************************************************
210 966 : SUBROUTINE torch_tensor_reset_from_array_double_${ndims}$d(tensor, source, requires_grad)
211 : TYPE(torch_tensor_type), INTENT(INOUT) :: tensor
212 : #:set arraydims = ", ".join(":" for i in range(ndims))
213 : REAL(dp), DIMENSION(${arraydims}$), ALLOCATABLE, INTENT(IN) :: source
214 : LOGICAL, OPTIONAL, INTENT(IN) :: requires_grad
215 :
216 : #if defined(__LIBTORCH)
217 : INTEGER(kind=int_8), DIMENSION(${ndims}$) :: sizes_c
218 : LOGICAL :: my_req_grad
219 :
220 : INTERFACE
221 : SUBROUTINE torch_c_tensor_reset_from_array_double(tensor, req_grad, ndims, sizes, source) &
222 : BIND(C, name="torch_c_tensor_reset_from_array_double")
223 : IMPORT :: C_PTR, C_INT, C_INT64_T, C_DOUBLE, C_BOOL
224 : TYPE(C_PTR) :: tensor
225 : LOGICAL(kind=C_BOOL), VALUE :: req_grad
226 : INTEGER(kind=C_INT), VALUE :: ndims
227 : INTEGER(kind=C_INT64_T), DIMENSION(*) :: sizes
228 : REAL(kind=C_DOUBLE), DIMENSION(*) :: source
229 : END SUBROUTINE torch_c_tensor_reset_from_array_double
230 : END INTERFACE
231 :
232 966 : my_req_grad = .FALSE.
233 966 : IF (PRESENT(requires_grad)) my_req_grad = requires_grad
234 :
235 : #:for axis in range(ndims)
236 966 : sizes_c(${axis + 1}$) = SIZE(source, ${ndims - axis}$) ! C arrays are stored row-major.
237 : #:endfor
238 :
239 : CALL torch_c_tensor_reset_from_array_double(tensor=tensor%c_ptr, &
240 : req_grad=LOGICAL(my_req_grad, C_BOOL), &
241 : ndims=${ndims}$, &
242 : sizes=sizes_c, &
243 966 : source=source)
244 966 : CPASSERT(C_ASSOCIATED(tensor%c_ptr))
245 : #else
246 : CPABORT("CP2K compiled without the Torch library.")
247 : MARK_USED(tensor)
248 : MARK_USED(source)
249 : MARK_USED(requires_grad)
250 : #endif
251 966 : END SUBROUTINE torch_tensor_reset_from_array_double_${ndims}$d
252 :
253 : #:endfor
254 :
255 : ! **************************************************************************************************
256 : !> \brief Creates an expanded tensor view along one singleton dimension.
257 : ! **************************************************************************************************
258 372 : SUBROUTINE torch_tensor_expand_dim(tensor, dim, extent, result)
259 : TYPE(torch_tensor_type), INTENT(IN) :: tensor
260 : INTEGER, INTENT(IN) :: dim, extent
261 : TYPE(torch_tensor_type), INTENT(INOUT) :: result
262 :
263 : #if defined(__LIBTORCH)
264 : INTERFACE
265 : SUBROUTINE torch_c_tensor_expand_dim(tensor, dim, extent, result) &
266 : BIND(C, name="torch_c_tensor_expand_dim")
267 : IMPORT :: C_INT64_T, C_PTR
268 : TYPE(C_PTR), VALUE :: tensor
269 : INTEGER(kind=C_INT64_T), VALUE :: dim, extent
270 : TYPE(C_PTR) :: result
271 : END SUBROUTINE torch_c_tensor_expand_dim
272 : END INTERFACE
273 :
274 372 : CPASSERT(C_ASSOCIATED(tensor%c_ptr))
275 372 : CPASSERT(.NOT. C_ASSOCIATED(result%c_ptr))
276 372 : CPASSERT(dim >= 0)
277 372 : CPASSERT(extent >= 0)
278 : CALL torch_c_tensor_expand_dim(tensor=tensor%c_ptr, &
279 : dim=INT(dim, C_INT64_T), &
280 : extent=INT(extent, C_INT64_T), &
281 372 : result=result%c_ptr)
282 372 : CPASSERT(C_ASSOCIATED(result%c_ptr))
283 : #else
284 : CPABORT("CP2K compiled without the Torch library.")
285 : MARK_USED(tensor)
286 : MARK_USED(dim)
287 : MARK_USED(extent)
288 : MARK_USED(result)
289 : #endif
290 372 : END SUBROUTINE torch_tensor_expand_dim
291 :
292 : ! **************************************************************************************************
293 : !> \brief Creates a view of a contiguous tensor slice.
294 : ! **************************************************************************************************
295 32 : SUBROUTINE torch_tensor_narrow(tensor, dim, start_index, length, result)
296 : TYPE(torch_tensor_type), INTENT(IN) :: tensor
297 : INTEGER, INTENT(IN) :: dim, start_index, length
298 : TYPE(torch_tensor_type), INTENT(INOUT) :: result
299 :
300 : #if defined(__LIBTORCH)
301 : INTERFACE
302 : SUBROUTINE torch_c_tensor_narrow(tensor, dim, start_index, length, result) &
303 : BIND(C, name="torch_c_tensor_narrow")
304 : IMPORT :: C_INT64_T, C_PTR
305 : TYPE(C_PTR), VALUE :: tensor
306 : INTEGER(kind=C_INT64_T), VALUE :: dim, start_index, length
307 : TYPE(C_PTR) :: result
308 : END SUBROUTINE torch_c_tensor_narrow
309 : END INTERFACE
310 :
311 32 : CPASSERT(C_ASSOCIATED(tensor%c_ptr))
312 32 : CPASSERT(.NOT. C_ASSOCIATED(result%c_ptr))
313 32 : CPASSERT(dim >= 0)
314 32 : CPASSERT(start_index >= 0)
315 32 : CPASSERT(length >= 0)
316 : CALL torch_c_tensor_narrow(tensor=tensor%c_ptr, &
317 : dim=INT(dim, C_INT64_T), &
318 : start_index=INT(start_index, C_INT64_T), &
319 : length=INT(length, C_INT64_T), &
320 32 : result=result%c_ptr)
321 32 : CPASSERT(C_ASSOCIATED(result%c_ptr))
322 : #else
323 : CPABORT("CP2K compiled without the Torch library.")
324 : MARK_USED(tensor)
325 : MARK_USED(dim)
326 : MARK_USED(start_index)
327 : MARK_USED(length)
328 : MARK_USED(result)
329 : #endif
330 32 : END SUBROUTINE torch_tensor_narrow
331 :
332 : ! **************************************************************************************************
333 : !> \brief Runs autograd on a Torch tensor.
334 : !> \author Ole Schuett
335 : ! **************************************************************************************************
336 6 : SUBROUTINE torch_tensor_backward(tensor, outer_grad)
337 : TYPE(torch_tensor_type), INTENT(IN) :: tensor
338 : TYPE(torch_tensor_type), INTENT(IN) :: outer_grad
339 :
340 : #if defined(__LIBTORCH)
341 : CHARACTER(len=*), PARAMETER :: routineN = 'torch_tensor_backward'
342 : INTEGER :: handle
343 :
344 : INTERFACE
345 : SUBROUTINE torch_c_tensor_backward(tensor, outer_grad) &
346 : BIND(C, name="torch_c_tensor_backward")
347 : IMPORT :: C_CHAR, C_PTR
348 : TYPE(C_PTR), VALUE :: tensor
349 : TYPE(C_PTR), VALUE :: outer_grad
350 : END SUBROUTINE torch_c_tensor_backward
351 : END INTERFACE
352 :
353 6 : CALL timeset(routineN, handle)
354 6 : CPASSERT(C_ASSOCIATED(tensor%c_ptr))
355 6 : CPASSERT(C_ASSOCIATED(outer_grad%c_ptr))
356 6 : CALL torch_c_tensor_backward(tensor=tensor%c_ptr, outer_grad=outer_grad%c_ptr)
357 6 : CALL timestop(handle)
358 : #else
359 : CPABORT("CP2K compiled without the Torch library.")
360 : MARK_USED(tensor)
361 : MARK_USED(outer_grad)
362 : #endif
363 6 : END SUBROUTINE torch_tensor_backward
364 :
365 : ! **************************************************************************************************
366 : !> \brief Runs autograd on a scalar Torch tensor.
367 : ! **************************************************************************************************
368 362 : SUBROUTINE torch_tensor_backward_scalar(tensor)
369 : TYPE(torch_tensor_type), INTENT(IN) :: tensor
370 :
371 : #if defined(__LIBTORCH)
372 : INTERFACE
373 : SUBROUTINE torch_c_tensor_backward_scalar(tensor) &
374 : BIND(C, name="torch_c_tensor_backward_scalar")
375 : IMPORT :: C_PTR
376 : TYPE(C_PTR), VALUE :: tensor
377 : END SUBROUTINE torch_c_tensor_backward_scalar
378 : END INTERFACE
379 :
380 362 : CPASSERT(C_ASSOCIATED(tensor%c_ptr))
381 362 : CALL torch_c_tensor_backward_scalar(tensor=tensor%c_ptr)
382 : #else
383 : CPABORT("CP2K compiled without the Torch library.")
384 : MARK_USED(tensor)
385 : #endif
386 362 : END SUBROUTINE torch_tensor_backward_scalar
387 :
388 : ! **************************************************************************************************
389 : !> \brief Moves a tensor to the active Torch device and makes it an autograd leaf.
390 : ! **************************************************************************************************
391 2180 : SUBROUTINE torch_tensor_to_device_leaf(tensor, requires_grad)
392 : TYPE(torch_tensor_type), INTENT(INOUT) :: tensor
393 : LOGICAL, INTENT(IN) :: requires_grad
394 :
395 : #if defined(__LIBTORCH)
396 : INTERFACE
397 : SUBROUTINE torch_c_tensor_to_device_leaf(tensor, req_grad) &
398 : BIND(C, name="torch_c_tensor_to_device_leaf")
399 : IMPORT :: C_BOOL, C_PTR
400 : TYPE(C_PTR) :: tensor
401 : LOGICAL(kind=C_BOOL), VALUE :: req_grad
402 : END SUBROUTINE torch_c_tensor_to_device_leaf
403 : END INTERFACE
404 :
405 2180 : CPASSERT(C_ASSOCIATED(tensor%c_ptr))
406 : CALL torch_c_tensor_to_device_leaf(tensor=tensor%c_ptr, &
407 2180 : req_grad=LOGICAL(requires_grad, C_BOOL))
408 2180 : CPASSERT(C_ASSOCIATED(tensor%c_ptr))
409 : #else
410 : CPABORT("CP2K compiled without the Torch library.")
411 : MARK_USED(tensor)
412 : MARK_USED(requires_grad)
413 : #endif
414 2180 : END SUBROUTINE torch_tensor_to_device_leaf
415 :
416 : ! **************************************************************************************************
417 : !> \brief Select whether Torch wrappers should use CUDA when available.
418 : ! **************************************************************************************************
419 772 : SUBROUTINE torch_use_cuda(use_cuda)
420 : LOGICAL, INTENT(IN) :: use_cuda
421 :
422 : #if defined(__LIBTORCH)
423 : INTERFACE
424 : SUBROUTINE torch_c_use_cuda(use_cuda) BIND(C, name="torch_c_use_cuda")
425 : IMPORT :: C_BOOL
426 : LOGICAL(kind=C_BOOL), VALUE :: use_cuda
427 : END SUBROUTINE torch_c_use_cuda
428 : END INTERFACE
429 :
430 772 : CALL torch_c_use_cuda(use_cuda=LOGICAL(use_cuda, C_BOOL))
431 : #else
432 : MARK_USED(use_cuda)
433 : #endif
434 772 : END SUBROUTINE torch_use_cuda
435 :
436 : ! **************************************************************************************************
437 : !> \brief Returns the gradient of a Torch tensor which was computed by autograd.
438 : !> \author Ole Schuett
439 : ! **************************************************************************************************
440 552 : SUBROUTINE torch_tensor_grad(tensor, grad)
441 : TYPE(torch_tensor_type), INTENT(IN) :: tensor
442 : TYPE(torch_tensor_type), INTENT(INOUT) :: grad
443 :
444 : #if defined(__LIBTORCH)
445 : INTERFACE
446 : SUBROUTINE torch_c_tensor_grad(tensor, grad) &
447 : BIND(C, name="torch_c_tensor_grad")
448 : IMPORT :: C_PTR
449 : TYPE(C_PTR), VALUE :: tensor
450 : TYPE(C_PTR) :: grad
451 : END SUBROUTINE torch_c_tensor_grad
452 : END INTERFACE
453 :
454 552 : CPASSERT(C_ASSOCIATED(tensor%c_ptr))
455 552 : CPASSERT(.NOT. C_ASSOCIATED(grad%c_ptr))
456 552 : CALL torch_c_tensor_grad(tensor=tensor%c_ptr, grad=grad%c_ptr)
457 552 : CPASSERT(C_ASSOCIATED(grad%c_ptr))
458 : #else
459 : CPABORT("CP2K compiled without the Torch library.")
460 : MARK_USED(tensor)
461 : MARK_USED(grad)
462 : #endif
463 552 : END SUBROUTINE torch_tensor_grad
464 :
465 : ! **************************************************************************************************
466 : !> \brief Copies three autograd gradients to CPU memory.
467 : ! **************************************************************************************************
468 360 : SUBROUTINE torch_tensor_grad_batch3(tensor1, tensor2, tensor3, grad1, grad2, grad3)
469 : TYPE(torch_tensor_type), INTENT(IN) :: tensor1, tensor2, tensor3
470 : TYPE(torch_tensor_type), INTENT(INOUT) :: grad1, grad2, grad3
471 :
472 : #if defined(__LIBTORCH)
473 : INTERFACE
474 : SUBROUTINE torch_c_tensor_grad_batch3(tensor1, tensor2, tensor3, grad1, grad2, grad3) &
475 : BIND(C, name="torch_c_tensor_grad_batch3")
476 : IMPORT :: C_PTR
477 : TYPE(C_PTR), VALUE :: tensor1, tensor2, tensor3
478 : TYPE(C_PTR) :: grad1, grad2, grad3
479 : END SUBROUTINE torch_c_tensor_grad_batch3
480 : END INTERFACE
481 :
482 360 : CPASSERT(C_ASSOCIATED(tensor1%c_ptr))
483 360 : CPASSERT(C_ASSOCIATED(tensor2%c_ptr))
484 360 : CPASSERT(C_ASSOCIATED(tensor3%c_ptr))
485 360 : CPASSERT(.NOT. C_ASSOCIATED(grad1%c_ptr))
486 360 : CPASSERT(.NOT. C_ASSOCIATED(grad2%c_ptr))
487 360 : CPASSERT(.NOT. C_ASSOCIATED(grad3%c_ptr))
488 : CALL torch_c_tensor_grad_batch3(tensor1=tensor1%c_ptr, tensor2=tensor2%c_ptr, &
489 : tensor3=tensor3%c_ptr, grad1=grad1%c_ptr, &
490 360 : grad2=grad2%c_ptr, grad3=grad3%c_ptr)
491 360 : CPASSERT(C_ASSOCIATED(grad1%c_ptr))
492 360 : CPASSERT(C_ASSOCIATED(grad2%c_ptr))
493 360 : CPASSERT(C_ASSOCIATED(grad3%c_ptr))
494 : #else
495 : CPABORT("CP2K compiled without the Torch library.")
496 : MARK_USED(tensor1)
497 : MARK_USED(tensor2)
498 : MARK_USED(tensor3)
499 : MARK_USED(grad1)
500 : MARK_USED(grad2)
501 : MARK_USED(grad3)
502 : #endif
503 360 : END SUBROUTINE torch_tensor_grad_batch3
504 :
505 : ! **************************************************************************************************
506 : !> \brief Returns the weighted sum of two Torch tensors.
507 : ! **************************************************************************************************
508 386 : SUBROUTINE torch_tensor_weighted_sum(values, weights, result)
509 : TYPE(torch_tensor_type), INTENT(IN) :: values, weights
510 : TYPE(torch_tensor_type), INTENT(INOUT) :: result
511 :
512 : #if defined(__LIBTORCH)
513 : INTERFACE
514 : SUBROUTINE torch_c_tensor_weighted_sum(values, weights, result) &
515 : BIND(C, name="torch_c_tensor_weighted_sum")
516 : IMPORT :: C_PTR
517 : TYPE(C_PTR), VALUE :: values
518 : TYPE(C_PTR), VALUE :: weights
519 : TYPE(C_PTR) :: result
520 : END SUBROUTINE torch_c_tensor_weighted_sum
521 : END INTERFACE
522 :
523 386 : CPASSERT(C_ASSOCIATED(values%c_ptr))
524 386 : CPASSERT(C_ASSOCIATED(weights%c_ptr))
525 386 : CPASSERT(.NOT. C_ASSOCIATED(result%c_ptr))
526 386 : CALL torch_c_tensor_weighted_sum(values=values%c_ptr, weights=weights%c_ptr, result=result%c_ptr)
527 386 : CPASSERT(C_ASSOCIATED(result%c_ptr))
528 : #else
529 : CPABORT("CP2K compiled without the Torch library.")
530 : MARK_USED(values)
531 : MARK_USED(weights)
532 : MARK_USED(result)
533 : #endif
534 386 : END SUBROUTINE torch_tensor_weighted_sum
535 :
536 : ! **************************************************************************************************
537 : !> \brief Returns a scalar double value from a Torch tensor.
538 : ! **************************************************************************************************
539 386 : FUNCTION torch_tensor_item_double(tensor) RESULT(value)
540 : TYPE(torch_tensor_type), INTENT(IN) :: tensor
541 : REAL(KIND=dp) :: value
542 :
543 : #if defined(__LIBTORCH)
544 : INTERFACE
545 : FUNCTION torch_c_tensor_item_double(tensor) RESULT(value) &
546 : BIND(C, name="torch_c_tensor_item_double")
547 : IMPORT :: C_DOUBLE, C_PTR
548 : TYPE(C_PTR), VALUE :: tensor
549 : REAL(KIND=C_DOUBLE) :: value
550 : END FUNCTION torch_c_tensor_item_double
551 : END INTERFACE
552 :
553 386 : CPASSERT(C_ASSOCIATED(tensor%c_ptr))
554 386 : value = torch_c_tensor_item_double(tensor=tensor%c_ptr)
555 : #else
556 : value = 0.0_dp
557 : CPABORT("CP2K compiled without the Torch library.")
558 : MARK_USED(tensor)
559 : #endif
560 386 : END FUNCTION torch_tensor_item_double
561 :
562 : ! **************************************************************************************************
563 : !> \brief Releases a Torch tensor and all its ressources.
564 : !> \author Ole Schuett
565 : ! **************************************************************************************************
566 4202 : SUBROUTINE torch_tensor_release(tensor)
567 : TYPE(torch_tensor_type), INTENT(INOUT) :: tensor
568 :
569 : #if defined(__LIBTORCH)
570 : INTERFACE
571 : SUBROUTINE torch_c_tensor_release(tensor) BIND(C, name="torch_c_tensor_release")
572 : IMPORT :: C_PTR
573 : TYPE(C_PTR), VALUE :: tensor
574 : END SUBROUTINE torch_c_tensor_release
575 : END INTERFACE
576 :
577 4202 : CPASSERT(C_ASSOCIATED(tensor%c_ptr))
578 4202 : CALL torch_c_tensor_release(tensor=tensor%c_ptr)
579 4202 : tensor%c_ptr = C_NULL_PTR
580 : #else
581 : CPABORT("CP2K was compiled without Torch library.")
582 : MARK_USED(tensor)
583 : #endif
584 4202 : END SUBROUTINE torch_tensor_release
585 :
586 : ! **************************************************************************************************
587 : !> \brief Creates an empty Torch dictionary.
588 : !> \author Ole Schuett
589 : ! **************************************************************************************************
590 450 : SUBROUTINE torch_dict_create(dict)
591 : TYPE(torch_dict_type), INTENT(INOUT) :: dict
592 :
593 : #if defined(__LIBTORCH)
594 : INTERFACE
595 : SUBROUTINE torch_c_dict_create(dict) BIND(C, name="torch_c_dict_create")
596 : IMPORT :: C_PTR
597 : TYPE(C_PTR) :: dict
598 : END SUBROUTINE torch_c_dict_create
599 : END INTERFACE
600 :
601 450 : CPASSERT(.NOT. C_ASSOCIATED(dict%c_ptr))
602 450 : CALL torch_c_dict_create(dict=dict%c_ptr)
603 450 : CPASSERT(C_ASSOCIATED(dict%c_ptr))
604 : #else
605 : CPABORT("CP2K was compiled without Torch library.")
606 : MARK_USED(dict)
607 : #endif
608 450 : END SUBROUTINE torch_dict_create
609 :
610 : ! **************************************************************************************************
611 : !> \brief Clones a Torch dictionary.
612 : ! **************************************************************************************************
613 132 : SUBROUTINE torch_dict_clone(source, target)
614 : TYPE(torch_dict_type), INTENT(IN) :: source
615 : TYPE(torch_dict_type), INTENT(INOUT) :: target
616 :
617 : #if defined(__LIBTORCH)
618 : INTERFACE
619 : SUBROUTINE torch_c_dict_clone(source, target) BIND(C, name="torch_c_dict_clone")
620 : IMPORT :: C_PTR
621 : TYPE(C_PTR), VALUE :: source
622 : TYPE(C_PTR) :: target
623 : END SUBROUTINE torch_c_dict_clone
624 : END INTERFACE
625 :
626 132 : CPASSERT(C_ASSOCIATED(source%c_ptr))
627 132 : CPASSERT(.NOT. C_ASSOCIATED(target%c_ptr))
628 132 : CALL torch_c_dict_clone(source=source%c_ptr, target=target%c_ptr)
629 132 : CPASSERT(C_ASSOCIATED(target%c_ptr))
630 : #else
631 : CPABORT("CP2K was compiled without Torch library.")
632 : MARK_USED(source)
633 : MARK_USED(target)
634 : #endif
635 132 : END SUBROUTINE torch_dict_clone
636 :
637 : ! **************************************************************************************************
638 : !> \brief Inserts a Torch tensor into a Torch dictionary.
639 : !> \author Ole Schuett
640 : ! **************************************************************************************************
641 2934 : SUBROUTINE torch_dict_insert(dict, key, tensor)
642 : TYPE(torch_dict_type), INTENT(INOUT) :: dict
643 : CHARACTER(len=*), INTENT(IN) :: key
644 : TYPE(torch_tensor_type), INTENT(IN) :: tensor
645 :
646 : #if defined(__LIBTORCH)
647 :
648 : INTERFACE
649 : SUBROUTINE torch_c_dict_insert(dict, key, tensor) &
650 : BIND(C, name="torch_c_dict_insert")
651 : IMPORT :: C_CHAR, C_PTR
652 : TYPE(C_PTR), VALUE :: dict
653 : CHARACTER(kind=C_CHAR), DIMENSION(*) :: key
654 : TYPE(C_PTR), VALUE :: tensor
655 : END SUBROUTINE torch_c_dict_insert
656 : END INTERFACE
657 :
658 2934 : CPASSERT(C_ASSOCIATED(dict%c_ptr))
659 2934 : CPASSERT(C_ASSOCIATED(tensor%c_ptr))
660 2934 : CALL torch_c_dict_insert(dict=dict%c_ptr, key=TRIM(key)//C_NULL_CHAR, tensor=tensor%c_ptr)
661 : #else
662 : CPABORT("CP2K compiled without the Torch library.")
663 : MARK_USED(dict)
664 : MARK_USED(key)
665 : MARK_USED(tensor)
666 : #endif
667 2934 : END SUBROUTINE torch_dict_insert
668 :
669 : ! **************************************************************************************************
670 : !> \brief Retrieves a Torch tensor from a Torch dictionary.
671 : !> \author Ole Schuett
672 : ! **************************************************************************************************
673 76 : SUBROUTINE torch_dict_get(dict, key, tensor)
674 : TYPE(torch_dict_type), INTENT(IN) :: dict
675 : CHARACTER(len=*), INTENT(IN) :: key
676 : TYPE(torch_tensor_type), INTENT(INOUT) :: tensor
677 :
678 : #if defined(__LIBTORCH)
679 :
680 : INTERFACE
681 : SUBROUTINE torch_c_dict_get(dict, key, tensor) &
682 : BIND(C, name="torch_c_dict_get")
683 : IMPORT :: C_CHAR, C_PTR
684 : TYPE(C_PTR), VALUE :: dict
685 : CHARACTER(kind=C_CHAR), DIMENSION(*) :: key
686 : TYPE(C_PTR) :: tensor
687 : END SUBROUTINE torch_c_dict_get
688 : END INTERFACE
689 :
690 76 : CPASSERT(C_ASSOCIATED(dict%c_ptr))
691 76 : CPASSERT(.NOT. C_ASSOCIATED(tensor%c_ptr))
692 76 : CALL torch_c_dict_get(dict=dict%c_ptr, key=TRIM(key)//C_NULL_CHAR, tensor=tensor%c_ptr)
693 76 : CPASSERT(C_ASSOCIATED(tensor%c_ptr))
694 :
695 : #else
696 : CPABORT("CP2K compiled without the Torch library.")
697 : MARK_USED(dict)
698 : MARK_USED(key)
699 : MARK_USED(tensor)
700 : #endif
701 76 : END SUBROUTINE torch_dict_get
702 :
703 : ! **************************************************************************************************
704 : !> \brief Releases a Torch dictionary and all its ressources.
705 : !> \author Ole Schuett
706 : ! **************************************************************************************************
707 344 : SUBROUTINE torch_dict_release(dict)
708 : TYPE(torch_dict_type), INTENT(INOUT) :: dict
709 :
710 : #if defined(__LIBTORCH)
711 : INTERFACE
712 : SUBROUTINE torch_c_dict_release(dict) BIND(C, name="torch_c_dict_release")
713 : IMPORT :: C_PTR
714 : TYPE(C_PTR), VALUE :: dict
715 : END SUBROUTINE torch_c_dict_release
716 : END INTERFACE
717 :
718 344 : CPASSERT(C_ASSOCIATED(dict%c_ptr))
719 344 : CALL torch_c_dict_release(dict=dict%c_ptr)
720 344 : dict%c_ptr = C_NULL_PTR
721 : #else
722 : CPABORT("CP2K was compiled without Torch library.")
723 : MARK_USED(dict)
724 : #endif
725 344 : END SUBROUTINE torch_dict_release
726 :
727 : ! **************************************************************************************************
728 : !> \brief Loads a Torch model from given "*.pth" file. (In Torch lingo models are called modules)
729 : !> \author Ole Schuett
730 : ! **************************************************************************************************
731 16 : SUBROUTINE torch_model_load(model, filename)
732 : TYPE(torch_model_type), INTENT(INOUT) :: model
733 : CHARACTER(len=*), INTENT(IN) :: filename
734 :
735 : #if defined(__LIBTORCH)
736 : CHARACTER(len=*), PARAMETER :: routineN = 'torch_model_load'
737 : INTEGER :: handle
738 :
739 : INTERFACE
740 : SUBROUTINE torch_c_model_load(model, filename) BIND(C, name="torch_c_model_load")
741 : IMPORT :: C_PTR, C_CHAR
742 : TYPE(C_PTR) :: model
743 : CHARACTER(kind=C_CHAR), DIMENSION(*) :: filename
744 : END SUBROUTINE torch_c_model_load
745 : END INTERFACE
746 :
747 16 : CALL timeset(routineN, handle)
748 16 : CPASSERT(.NOT. C_ASSOCIATED(model%c_ptr))
749 16 : CALL torch_c_model_load(model=model%c_ptr, filename=TRIM(filename)//C_NULL_CHAR)
750 16 : CPASSERT(C_ASSOCIATED(model%c_ptr))
751 16 : CALL timestop(handle)
752 : #else
753 : CPABORT("CP2K was compiled without Torch library.")
754 : MARK_USED(model)
755 : MARK_USED(filename)
756 : #endif
757 16 : END SUBROUTINE torch_model_load
758 :
759 : ! **************************************************************************************************
760 : !> \brief Loads a Torch model and reads two metadata entries in the same archive pass.
761 : ! **************************************************************************************************
762 97 : SUBROUTINE torch_model_load_with_metadata(model, filename, key1, value1, key2, value2)
763 : TYPE(torch_model_type), INTENT(INOUT) :: model
764 : CHARACTER(len=*), INTENT(IN) :: filename, key1, key2
765 : CHARACTER(:), ALLOCATABLE, INTENT(OUT) :: value1, value2
766 :
767 : #if defined(__LIBTORCH)
768 : CHARACTER(len=*), PARAMETER :: routineN = 'torch_model_load_with_metadata'
769 : INTEGER :: handle, length1, length2
770 : TYPE(C_PTR) :: content1_c, content2_c
771 :
772 : INTERFACE
773 : SUBROUTINE torch_c_model_load_with_metadata(model, filename, key1, key2, &
774 : content1, length1, content2, length2) &
775 : BIND(C, name="torch_c_model_load_with_metadata")
776 : IMPORT :: C_CHAR, C_INT, C_PTR
777 : TYPE(C_PTR) :: model
778 : CHARACTER(kind=C_CHAR), DIMENSION(*) :: filename, key1, key2
779 : TYPE(C_PTR) :: content1, content2
780 : INTEGER(kind=C_INT) :: length1, length2
781 : END SUBROUTINE torch_c_model_load_with_metadata
782 : END INTERFACE
783 :
784 97 : CALL timeset(routineN, handle)
785 97 : CPASSERT(.NOT. C_ASSOCIATED(model%c_ptr))
786 97 : content1_c = C_NULL_PTR
787 97 : content2_c = C_NULL_PTR
788 97 : length1 = -1
789 97 : length2 = -1
790 : CALL torch_c_model_load_with_metadata(model=model%c_ptr, &
791 : filename=TRIM(filename)//C_NULL_CHAR, &
792 : key1=TRIM(key1)//C_NULL_CHAR, &
793 : key2=TRIM(key2)//C_NULL_CHAR, &
794 : content1=content1_c, length1=length1, &
795 97 : content2=content2_c, length2=length2)
796 97 : CPASSERT(C_ASSOCIATED(model%c_ptr))
797 97 : CALL c_string_to_allocatable(content1_c, length1, value1)
798 97 : CALL c_string_to_allocatable(content2_c, length2, value2)
799 97 : CALL timestop(handle)
800 : #else
801 : CPABORT("CP2K was compiled without Torch library.")
802 : MARK_USED(model)
803 : MARK_USED(filename)
804 : MARK_USED(key1)
805 : MARK_USED(value1)
806 : MARK_USED(key2)
807 : MARK_USED(value2)
808 : #endif
809 97 : END SUBROUTINE torch_model_load_with_metadata
810 :
811 : ! **************************************************************************************************
812 : !> \brief Maps serialized TorchScript device constants to the active Torch device.
813 : ! **************************************************************************************************
814 97 : SUBROUTINE torch_model_remap_device_constants(model)
815 : TYPE(torch_model_type), INTENT(INOUT) :: model
816 :
817 : #if defined(__LIBTORCH)
818 : INTERFACE
819 : SUBROUTINE torch_c_model_remap_device_constants(model) &
820 : BIND(C, name="torch_c_model_remap_device_constants")
821 : IMPORT :: C_PTR
822 : TYPE(C_PTR), VALUE :: model
823 : END SUBROUTINE torch_c_model_remap_device_constants
824 : END INTERFACE
825 :
826 97 : CPASSERT(C_ASSOCIATED(model%c_ptr))
827 97 : CALL torch_c_model_remap_device_constants(model=model%c_ptr)
828 : #else
829 : CPABORT("CP2K was compiled without Torch library.")
830 : MARK_USED(model)
831 : #endif
832 97 : END SUBROUTINE torch_model_remap_device_constants
833 :
834 : ! **************************************************************************************************
835 : !> \brief Disable gradients for inference-only model parameters.
836 : ! **************************************************************************************************
837 97 : SUBROUTINE torch_model_disable_parameter_gradients(model)
838 : TYPE(torch_model_type), INTENT(INOUT) :: model
839 :
840 : #if defined(__LIBTORCH)
841 : INTERFACE
842 : SUBROUTINE torch_c_model_disable_parameter_gradients(model) &
843 : BIND(C, name="torch_c_model_disable_parameter_gradients")
844 : IMPORT :: C_PTR
845 : TYPE(C_PTR), VALUE :: model
846 : END SUBROUTINE torch_c_model_disable_parameter_gradients
847 : END INTERFACE
848 :
849 97 : CPASSERT(C_ASSOCIATED(model%c_ptr))
850 97 : CALL torch_c_model_disable_parameter_gradients(model=model%c_ptr)
851 : #else
852 : CPABORT("CP2K was compiled without Torch library.")
853 : MARK_USED(model)
854 : #endif
855 97 : END SUBROUTINE torch_model_disable_parameter_gradients
856 :
857 : ! **************************************************************************************************
858 : !> \brief Evaluates the given Torch model.
859 : !> \author Ole Schuett
860 : ! **************************************************************************************************
861 62 : SUBROUTINE torch_model_forward(model, inputs, outputs)
862 : TYPE(torch_model_type), INTENT(INOUT) :: model
863 : TYPE(torch_dict_type), INTENT(IN) :: inputs
864 : TYPE(torch_dict_type), INTENT(INOUT) :: outputs
865 :
866 : #if defined(__LIBTORCH)
867 : CHARACTER(len=*), PARAMETER :: routineN = 'torch_model_forward'
868 : INTEGER :: handle
869 :
870 : INTERFACE
871 : SUBROUTINE torch_c_model_forward(model, inputs, outputs) BIND(C, name="torch_c_model_forward")
872 : IMPORT :: C_PTR
873 : TYPE(C_PTR), VALUE :: model
874 : TYPE(C_PTR), VALUE :: inputs
875 : TYPE(C_PTR), VALUE :: outputs
876 : END SUBROUTINE torch_c_model_forward
877 : END INTERFACE
878 :
879 62 : CALL timeset(routineN, handle)
880 62 : CPASSERT(C_ASSOCIATED(model%c_ptr))
881 62 : CPASSERT(C_ASSOCIATED(inputs%c_ptr))
882 62 : CPASSERT(C_ASSOCIATED(outputs%c_ptr))
883 62 : CALL torch_c_model_forward(model=model%c_ptr, inputs=inputs%c_ptr, outputs=outputs%c_ptr)
884 62 : CALL timestop(handle)
885 : #else
886 : CPABORT("CP2K was compiled without Torch library.")
887 : MARK_USED(model)
888 : MARK_USED(inputs)
889 : MARK_USED(outputs)
890 : #endif
891 62 : END SUBROUTINE torch_model_forward
892 :
893 : ! **************************************************************************************************
894 : !> \brief Evaluates a TorchScript model method expecting keyword argument "mol".
895 : ! **************************************************************************************************
896 386 : SUBROUTINE torch_model_forward_mol_tensor(model, method_name, inputs, output)
897 : TYPE(torch_model_type), INTENT(INOUT) :: model
898 : CHARACTER(len=*), INTENT(IN) :: method_name
899 : TYPE(torch_dict_type), INTENT(IN) :: inputs
900 : TYPE(torch_tensor_type), INTENT(INOUT) :: output
901 :
902 : #if defined(__LIBTORCH)
903 : CHARACTER(len=*), PARAMETER :: routineN = 'torch_model_forward_mol_tensor'
904 : INTEGER :: handle
905 :
906 : INTERFACE
907 : SUBROUTINE torch_c_model_forward_mol_tensor(model, method_name, inputs, output) &
908 : BIND(C, name="torch_c_model_forward_mol_tensor")
909 : IMPORT :: C_CHAR, C_PTR
910 : TYPE(C_PTR), VALUE :: model
911 : CHARACTER(kind=C_CHAR), DIMENSION(*) :: method_name
912 : TYPE(C_PTR), VALUE :: inputs
913 : TYPE(C_PTR) :: output
914 : END SUBROUTINE torch_c_model_forward_mol_tensor
915 : END INTERFACE
916 :
917 386 : CALL timeset(routineN, handle)
918 386 : CPASSERT(C_ASSOCIATED(model%c_ptr))
919 386 : CPASSERT(C_ASSOCIATED(inputs%c_ptr))
920 386 : CPASSERT(.NOT. C_ASSOCIATED(output%c_ptr))
921 : CALL torch_c_model_forward_mol_tensor(model=model%c_ptr, &
922 : method_name=TRIM(method_name)//C_NULL_CHAR, &
923 : inputs=inputs%c_ptr, &
924 386 : output=output%c_ptr)
925 386 : CPASSERT(C_ASSOCIATED(output%c_ptr))
926 386 : CALL timestop(handle)
927 : #else
928 : CPABORT("CP2K was compiled without Torch library.")
929 : MARK_USED(model)
930 : MARK_USED(method_name)
931 : MARK_USED(inputs)
932 : MARK_USED(output)
933 : #endif
934 386 : END SUBROUTINE torch_model_forward_mol_tensor
935 :
936 : ! **************************************************************************************************
937 : !> \brief Releases a Torch model and all its ressources.
938 : !> \author Ole Schuett
939 : ! **************************************************************************************************
940 16 : SUBROUTINE torch_model_release(model)
941 : TYPE(torch_model_type), INTENT(INOUT) :: model
942 :
943 : #if defined(__LIBTORCH)
944 : INTERFACE
945 : SUBROUTINE torch_c_model_release(model) BIND(C, name="torch_c_model_release")
946 : IMPORT :: C_PTR
947 : TYPE(C_PTR), VALUE :: model
948 : END SUBROUTINE torch_c_model_release
949 : END INTERFACE
950 :
951 16 : CPASSERT(C_ASSOCIATED(model%c_ptr))
952 16 : CALL torch_c_model_release(model=model%c_ptr)
953 16 : model%c_ptr = C_NULL_PTR
954 : #else
955 : CPABORT("CP2K was compiled without Torch library.")
956 : MARK_USED(model)
957 : #endif
958 16 : END SUBROUTINE torch_model_release
959 :
960 : ! **************************************************************************************************
961 : !> \brief Reads metadata entry from given "*.pth" file. (In Torch lingo they are called extra files)
962 : !> \author Ole Schuett
963 : ! **************************************************************************************************
964 40 : FUNCTION torch_model_read_metadata(filename, key) RESULT(res)
965 : CHARACTER(len=*), INTENT(IN) :: filename, key
966 : CHARACTER(:), ALLOCATABLE :: res
967 :
968 : #if defined(__LIBTORCH)
969 : CHARACTER(len=*), PARAMETER :: routineN = 'torch_model_read_metadata'
970 : INTEGER :: handle
971 :
972 : INTEGER :: length
973 : TYPE(C_PTR) :: content_c
974 :
975 : INTERFACE
976 : SUBROUTINE torch_c_model_read_metadata(filename, key, content, length) &
977 : BIND(C, name="torch_c_model_read_metadata")
978 : IMPORT :: C_CHAR, C_PTR, C_INT
979 : CHARACTER(kind=C_CHAR), DIMENSION(*) :: filename, key
980 : TYPE(C_PTR) :: content
981 : INTEGER(kind=C_INT) :: length
982 : END SUBROUTINE torch_c_model_read_metadata
983 : END INTERFACE
984 :
985 40 : CALL timeset(routineN, handle)
986 40 : content_c = C_NULL_PTR
987 40 : length = -1
988 : CALL torch_c_model_read_metadata(filename=TRIM(filename)//C_NULL_CHAR, &
989 : key=TRIM(key)//C_NULL_CHAR, &
990 : content=content_c, &
991 40 : length=length)
992 40 : CALL c_string_to_allocatable(content_c, length, res)
993 40 : CALL timestop(handle)
994 : #else
995 : res = ""
996 : MARK_USED(filename)
997 : MARK_USED(key)
998 : CPABORT("CP2K was compiled without Torch library.")
999 : #endif
1000 40 : END FUNCTION torch_model_read_metadata
1001 :
1002 : ! **************************************************************************************************
1003 : !> \brief Move a C-allocated null-terminated string into an allocatable Fortran string.
1004 : ! **************************************************************************************************
1005 234 : SUBROUTINE c_string_to_allocatable(content_c, length, res)
1006 : TYPE(C_PTR), INTENT(INOUT) :: content_c
1007 : INTEGER, INTENT(IN) :: length
1008 : CHARACTER(:), ALLOCATABLE, INTENT(OUT) :: res
1009 :
1010 : #if defined(__LIBTORCH)
1011 : CHARACTER(LEN=1, KIND=C_CHAR), DIMENSION(:), &
1012 234 : POINTER :: content_f
1013 : INTEGER :: i
1014 :
1015 : INTERFACE
1016 : SUBROUTINE torch_c_free_string(content) BIND(C, name="torch_c_free_string")
1017 : IMPORT :: C_PTR
1018 : TYPE(C_PTR), VALUE :: content
1019 : END SUBROUTINE torch_c_free_string
1020 : END INTERFACE
1021 :
1022 0 : CPASSERT(C_ASSOCIATED(content_c))
1023 234 : CPASSERT(length >= 0)
1024 :
1025 468 : CALL C_F_POINTER(content_c, content_f, shape=[length + 1])
1026 234 : CPASSERT(content_f(length + 1) == C_NULL_CHAR)
1027 :
1028 234 : ALLOCATE (CHARACTER(LEN=length) :: res)
1029 15914 : DO i = 1, length
1030 15680 : CPASSERT(content_f(i) /= C_NULL_CHAR)
1031 15914 : res(i:i) = content_f(i)
1032 : END DO
1033 :
1034 234 : NULLIFY (content_f)
1035 234 : CALL torch_c_free_string(content_c)
1036 234 : content_c = C_NULL_PTR
1037 :
1038 : #else
1039 : res = ""
1040 : MARK_USED(content_c)
1041 : MARK_USED(length)
1042 : CPABORT("CP2K was compiled without Torch library.")
1043 : #endif
1044 234 : END SUBROUTINE c_string_to_allocatable
1045 :
1046 : ! **************************************************************************************************
1047 : !> \brief Returns true iff the Torch CUDA backend is available.
1048 : !> \author Ole Schuett
1049 : ! **************************************************************************************************
1050 2 : FUNCTION torch_cuda_is_available() RESULT(res)
1051 : LOGICAL :: res
1052 :
1053 : #if defined(__LIBTORCH)
1054 : INTERFACE
1055 : FUNCTION torch_c_cuda_is_available() BIND(C, name="torch_c_cuda_is_available")
1056 : IMPORT :: C_BOOL
1057 : LOGICAL(C_BOOL) :: torch_c_cuda_is_available
1058 : END FUNCTION torch_c_cuda_is_available
1059 : END INTERFACE
1060 :
1061 2 : res = torch_c_cuda_is_available()
1062 : #else
1063 : CPABORT("CP2K was compiled without Torch library.")
1064 : res = .FALSE.
1065 : #endif
1066 2 : END FUNCTION torch_cuda_is_available
1067 :
1068 : ! **************************************************************************************************
1069 : !> \brief Return the number of CUDA devices visible to Torch.
1070 : ! **************************************************************************************************
1071 0 : FUNCTION torch_cuda_device_count() RESULT(count)
1072 : INTEGER :: count
1073 :
1074 : #if defined(__LIBTORCH)
1075 : INTERFACE
1076 : FUNCTION torch_c_cuda_device_count() BIND(C, name="torch_c_cuda_device_count")
1077 : IMPORT :: C_INT
1078 : INTEGER(C_INT) :: torch_c_cuda_device_count
1079 : END FUNCTION torch_c_cuda_device_count
1080 : END INTERFACE
1081 :
1082 0 : count = torch_c_cuda_device_count()
1083 : #else
1084 : CPABORT("CP2K was compiled without Torch library.")
1085 : count = 0
1086 : #endif
1087 0 : END FUNCTION torch_cuda_device_count
1088 :
1089 : ! **************************************************************************************************
1090 : !> \brief Set whether to allow the use of TF32.
1091 : !> Needed due to changes in defaults from pytorch 1.7 to 1.11 to >=1.12
1092 : !> See https://pytorch.org/docs/stable/notes/cuda.html
1093 : !> \author Gabriele Tocci
1094 : ! **************************************************************************************************
1095 6 : SUBROUTINE torch_allow_tf32(allow_tf32)
1096 : LOGICAL, INTENT(IN) :: allow_tf32
1097 :
1098 : #if defined(__LIBTORCH)
1099 : INTERFACE
1100 : SUBROUTINE torch_c_allow_tf32(allow_tf32) BIND(C, name="torch_c_allow_tf32")
1101 : IMPORT :: C_BOOL
1102 : LOGICAL(C_BOOL), VALUE :: allow_tf32
1103 : END SUBROUTINE torch_c_allow_tf32
1104 : END INTERFACE
1105 :
1106 6 : CALL torch_c_allow_tf32(allow_tf32=LOGICAL(allow_tf32, C_BOOL))
1107 : #else
1108 : CPABORT("CP2K was compiled without Torch library.")
1109 : MARK_USED(allow_tf32)
1110 : #endif
1111 6 : END SUBROUTINE torch_allow_tf32
1112 :
1113 : ! **************************************************************************************************
1114 : !> \brief Freeze the given Torch model: applies generic optimization that speed up model.
1115 : !> See https://pytorch.org/docs/stable/generated/torch.jit.freeze.html
1116 : !> \author Gabriele Tocci
1117 : ! **************************************************************************************************
1118 6 : SUBROUTINE torch_model_freeze(model)
1119 : TYPE(torch_model_type), INTENT(INOUT) :: model
1120 :
1121 : #if defined(__LIBTORCH)
1122 : CHARACTER(len=*), PARAMETER :: routineN = 'torch_model_freeze'
1123 : INTEGER :: handle
1124 :
1125 : INTERFACE
1126 : SUBROUTINE torch_c_model_freeze(model) BIND(C, name="torch_c_model_freeze")
1127 : IMPORT :: C_PTR
1128 : TYPE(C_PTR), VALUE :: model
1129 : END SUBROUTINE torch_c_model_freeze
1130 : END INTERFACE
1131 :
1132 6 : CALL timeset(routineN, handle)
1133 6 : CPASSERT(C_ASSOCIATED(model%c_ptr))
1134 6 : CALL torch_c_model_freeze(model=model%c_ptr)
1135 6 : CALL timestop(handle)
1136 : #else
1137 : CPABORT("CP2K was compiled without Torch library.")
1138 : MARK_USED(model)
1139 : #endif
1140 6 : END SUBROUTINE torch_model_freeze
1141 :
1142 : #:set typenames = ['int64', 'double', 'string']
1143 : #:set types_f = ['INTEGER(kind=int_8)', 'REAL(dp)', 'CHARACTER(LEN=default_string_length)']
1144 : #:set types_c = ['INTEGER(kind=C_INT64_T)', 'REAL(kind=C_DOUBLE)', 'CHARACTER(kind=C_CHAR), DIMENSION(*)']
1145 : #:set zeros_f = ['0', '0.0_dp', '""']
1146 :
1147 : #:for typename, type_f, type_c, zero_f in zip(typenames, types_f, types_c, zeros_f)
1148 : ! **************************************************************************************************
1149 : !> \brief Retrieves an attribute from a Torch model. Must be called before torch_model_freeze.
1150 : !> \author Ole Schuett
1151 : ! **************************************************************************************************
1152 64 : SUBROUTINE torch_model_get_attr_${typename}$ (model, key, dest)
1153 : TYPE(torch_model_type), INTENT(IN) :: model
1154 : CHARACTER(len=*), INTENT(IN) :: key
1155 : ${type_f}$, INTENT(OUT) :: dest
1156 :
1157 : #if defined(__LIBTORCH)
1158 :
1159 : INTERFACE
1160 : SUBROUTINE torch_c_model_get_attr_${typename}$ (model, key, dest) &
1161 : BIND(C, name="torch_c_model_get_attr_${typename}$")
1162 : IMPORT :: C_PTR, C_CHAR, C_INT64_T, C_DOUBLE
1163 : TYPE(C_PTR), VALUE :: model
1164 : CHARACTER(kind=C_CHAR), DIMENSION(*) :: key
1165 : ${type_c}$ :: dest
1166 : END SUBROUTINE torch_c_model_get_attr_${typename}$
1167 : END INTERFACE
1168 :
1169 : CALL torch_c_model_get_attr_${typename}$ (model=model%c_ptr, &
1170 : key=TRIM(key)//C_NULL_CHAR, &
1171 64 : dest=dest)
1172 : #else
1173 : dest = ${zero_f}$
1174 : MARK_USED(model)
1175 : MARK_USED(key)
1176 : CPABORT("CP2K compiled without the Torch library.")
1177 : #endif
1178 64 : END SUBROUTINE torch_model_get_attr_${typename}$
1179 : #:endfor
1180 :
1181 : ! **************************************************************************************************
1182 : !> \brief Retrieves an attribute from a Torch model. Must be called before torch_model_freeze.
1183 : !> \author Ole Schuett
1184 : ! **************************************************************************************************
1185 40 : SUBROUTINE torch_model_get_attr_int32(model, key, dest)
1186 : TYPE(torch_model_type), INTENT(IN) :: model
1187 : CHARACTER(len=*), INTENT(IN) :: key
1188 : INTEGER, INTENT(OUT) :: dest
1189 :
1190 : INTEGER(kind=int_8) :: temp
1191 40 : CALL torch_model_get_attr_int64(model, key, temp)
1192 40 : CPASSERT(ABS(temp) < HUGE(dest))
1193 40 : dest = INT(temp)
1194 40 : END SUBROUTINE torch_model_get_attr_int32
1195 :
1196 : ! **************************************************************************************************
1197 : !> \brief Retrieves a list attribute from a Torch model. Must be called before torch_model_freeze.
1198 : !> \author Ole Schuett
1199 : ! **************************************************************************************************
1200 8 : SUBROUTINE torch_model_get_attr_strlist(model, key, dest)
1201 : TYPE(torch_model_type), INTENT(IN) :: model
1202 : CHARACTER(len=*), INTENT(IN) :: key
1203 : CHARACTER(LEN=default_string_length), &
1204 : ALLOCATABLE, DIMENSION(:) :: dest
1205 :
1206 : #if defined(__LIBTORCH)
1207 :
1208 : INTEGER :: num_items, i
1209 :
1210 : INTERFACE
1211 : SUBROUTINE torch_c_model_get_attr_list_size(model, key, size) &
1212 : BIND(C, name="torch_c_model_get_attr_list_size")
1213 : IMPORT :: C_PTR, C_CHAR, C_INT
1214 : TYPE(C_PTR), VALUE :: model
1215 : CHARACTER(kind=C_CHAR), DIMENSION(*) :: key
1216 : INTEGER(kind=C_INT) :: size
1217 : END SUBROUTINE torch_c_model_get_attr_list_size
1218 : END INTERFACE
1219 :
1220 : INTERFACE
1221 : SUBROUTINE torch_c_model_get_attr_strlist(model, key, index, dest) &
1222 : BIND(C, name="torch_c_model_get_attr_strlist")
1223 : IMPORT :: C_PTR, C_CHAR, C_INT
1224 : TYPE(C_PTR), VALUE :: model
1225 : CHARACTER(kind=C_CHAR), DIMENSION(*) :: key
1226 : INTEGER(kind=C_INT), VALUE :: index
1227 : CHARACTER(kind=C_CHAR), DIMENSION(*) :: dest
1228 : END SUBROUTINE torch_c_model_get_attr_strlist
1229 : END INTERFACE
1230 :
1231 : CALL torch_c_model_get_attr_list_size(model=model%c_ptr, &
1232 : key=TRIM(key)//C_NULL_CHAR, &
1233 8 : size=num_items)
1234 24 : ALLOCATE (dest(num_items))
1235 24 : dest(:) = ""
1236 :
1237 24 : DO i = 1, num_items
1238 : CALL torch_c_model_get_attr_strlist(model=model%c_ptr, &
1239 : key=TRIM(key)//C_NULL_CHAR, &
1240 : index=i - 1, &
1241 24 : dest=dest(i))
1242 :
1243 : END DO
1244 : #else
1245 : CPABORT("CP2K compiled without the Torch library.")
1246 : MARK_USED(model)
1247 : MARK_USED(key)
1248 : MARK_USED(dest)
1249 : #endif
1250 :
1251 8 : END SUBROUTINE torch_model_get_attr_strlist
1252 :
1253 0 : END MODULE torch_api
|
