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