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