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 :
8 : #if defined(__LIBTORCH)
9 :
10 : #include <ATen/Parallel.h>
11 : #include <c10/core/DeviceGuard.h>
12 : #include <torch/csrc/api/include/torch/cuda.h>
13 : #include <torch/script.h>
14 :
15 : #include "offload/offload_library.h"
16 :
17 : #include <cassert>
18 :
19 : #include <cfenv>
20 : #include <climits>
21 : #include <cstdlib>
22 : #include <cstring>
23 : #include <string>
24 : #include <unordered_map>
25 : #include <vector>
26 :
27 : typedef torch::Tensor torch_c_tensor_t;
28 : typedef c10::Dict<std::string, torch::Tensor> torch_c_dict_t;
29 : typedef torch::jit::Module torch_c_model_t;
30 :
31 : class TorchFloatingPointMaskGuard {
32 : public:
33 392 : TorchFloatingPointMaskGuard() : active_(std::feholdexcept(&env_) == 0) {}
34 392 : ~TorchFloatingPointMaskGuard() {
35 392 : if (active_) {
36 392 : std::feclearexcept(FE_ALL_EXCEPT);
37 392 : std::fesetenv(&env_);
38 : }
39 392 : }
40 :
41 : private:
42 : std::fenv_t env_;
43 : bool active_;
44 : };
45 :
46 : /*******************************************************************************
47 : * \brief Internal helper for selecting the CUDA device when available.
48 : * \author Ole Schuett
49 : ******************************************************************************/
50 : static bool use_cuda_if_available = true;
51 :
52 200 : static bool get_positive_int_env(const char *name, int &value) {
53 200 : const char *raw = std::getenv(name);
54 200 : if (raw == nullptr || raw[0] == '\0') {
55 : return false;
56 : }
57 0 : char *end = nullptr;
58 0 : const long parsed = std::strtol(raw, &end, 10);
59 0 : if (end == raw || *end != '\0' || parsed <= 0 || parsed > INT_MAX) {
60 : return false;
61 : }
62 0 : value = static_cast<int>(parsed);
63 0 : return true;
64 : }
65 :
66 11668 : static void initialize_torch_threads_from_env() {
67 11668 : static bool initialized = false;
68 11668 : if (initialized) {
69 11568 : return;
70 : }
71 100 : initialized = true;
72 :
73 100 : int num_threads = 0;
74 100 : if (get_positive_int_env("CP2K_TORCH_NUM_THREADS", num_threads)) {
75 0 : at::set_num_threads(num_threads);
76 : }
77 100 : if (get_positive_int_env("CP2K_TORCH_NUM_INTEROP_THREADS", num_threads)) {
78 0 : at::set_num_interop_threads(num_threads);
79 : }
80 : }
81 :
82 9092 : static torch::Device get_device() {
83 9092 : initialize_torch_threads_from_env();
84 9092 : if (!use_cuda_if_available || !torch::cuda::is_available()) {
85 9092 : return torch::kCPU;
86 : }
87 0 : const auto device_count = torch::cuda::device_count();
88 0 : if (device_count <= 0) {
89 0 : return torch::kCPU;
90 : }
91 0 : const int chosen_device = offload_get_chosen_device();
92 0 : const int device = (chosen_device >= 0) ? chosen_device : 0;
93 0 : assert(device < device_count);
94 0 : return torch::Device(torch::kCUDA, device);
95 : }
96 :
97 9092 : static torch::Device get_device_with_guard(c10::OptionalDeviceGuard &guard) {
98 9092 : const auto device = get_device();
99 9092 : if (device.is_cuda()) {
100 0 : guard.reset_device(device);
101 : }
102 9092 : return device;
103 : }
104 :
105 104 : static void set_jit_fusion_strategy() {
106 : // JIT Fusion strategy optimization, hardcode dynamic 10, see also
107 : // https://github.com/mir-group/pair_nequip_allegro.git
108 104 : torch::jit::FusionStrategy strategy = {
109 104 : {torch::jit::FusionBehavior::DYNAMIC, 10}};
110 208 : torch::jit::setFusionStrategy(strategy);
111 104 : }
112 :
113 216 : static void copy_string_to_c_buffer(const std::string &source, char **content,
114 : int *length) {
115 216 : *length = source.length();
116 216 : *content = (char *)malloc(source.length() + 1); // +1 for null terminator
117 216 : strcpy(*content, source.c_str());
118 216 : }
119 :
120 104 : static bool can_load_directly_to_device(const torch::Device &device) {
121 104 : return !device.is_cuda() || device.index() == 0 ||
122 0 : torch::cuda::device_count() == 1;
123 : }
124 :
125 104 : static torch::jit::Module load_module_for_device(const char *filename,
126 : const torch::Device &device) {
127 104 : if (can_load_directly_to_device(device)) {
128 208 : return torch::jit::load(filename, device);
129 : }
130 0 : auto model = torch::jit::load(filename, torch::kCPU);
131 0 : model.to(device);
132 0 : return model;
133 0 : }
134 :
135 : /*******************************************************************************
136 : * \brief Internal helper for creating a Torch tensor from an array.
137 : * \author Ole Schuett
138 : ******************************************************************************/
139 2576 : static torch_c_tensor_t *tensor_from_array(const torch::Dtype dtype,
140 : const bool req_grad, const int ndims,
141 : const int64_t sizes[],
142 : void *source) {
143 2576 : initialize_torch_threads_from_env();
144 2576 : const auto opts = torch::TensorOptions().dtype(dtype).requires_grad(req_grad);
145 2576 : const auto sizes_ref = c10::IntArrayRef(sizes, ndims);
146 2576 : return new torch_c_tensor_t(torch::from_blob(source, sizes_ref, opts));
147 : }
148 :
149 870 : static bool tensor_matches(const torch_c_tensor_t *tensor,
150 : const torch::Dtype dtype,
151 : const torch::Device &device, const int ndims,
152 : const int64_t sizes[]) {
153 948 : if (tensor == nullptr || !tensor->defined() ||
154 1818 : tensor->scalar_type() != dtype || tensor->device() != device ||
155 474 : tensor->ndimension() != ndims) {
156 396 : return false;
157 : }
158 1580 : for (int i = 0; i < ndims; i++) {
159 1106 : if (tensor->size(i) != sizes[i]) {
160 : return false;
161 : }
162 : }
163 474 : return tensor->is_contiguous();
164 : }
165 :
166 870 : static void reset_tensor_from_array_double(torch_c_tensor_t **tensor,
167 : const bool req_grad, const int ndims,
168 : const int64_t sizes[],
169 : double source[]) {
170 870 : c10::OptionalDeviceGuard guard;
171 870 : const auto device = get_device_with_guard(guard);
172 870 : const auto sizes_ref = c10::IntArrayRef(sizes, ndims);
173 870 : if (!tensor_matches(*tensor, torch::kFloat64, device, ndims, sizes)) {
174 396 : delete (*tensor);
175 396 : const auto opts =
176 396 : torch::TensorOptions().dtype(torch::kFloat64).device(device);
177 792 : *tensor = new torch_c_tensor_t(torch::empty(sizes_ref, opts).detach());
178 : }
179 870 : const auto source_tensor = torch::from_blob(
180 870 : source, sizes_ref, torch::TensorOptions().dtype(torch::kFloat64));
181 870 : {
182 870 : torch::NoGradGuard no_grad;
183 870 : (*tensor)->copy_(source_tensor);
184 870 : (*tensor)->mutable_grad() = torch::Tensor();
185 0 : }
186 1740 : (*tensor)->set_requires_grad(req_grad);
187 870 : }
188 :
189 : /*******************************************************************************
190 : * \brief Internal helper for getting the data_ptr and sizes of a Torch tensor.
191 : * \author Ole Schuett
192 : ******************************************************************************/
193 1562 : static void *get_data_ptr(const torch_c_tensor_t *tensor,
194 : const torch::Dtype dtype, const int ndims,
195 : int64_t sizes[]) {
196 1562 : assert(tensor->scalar_type() == dtype);
197 1562 : assert(tensor->ndimension() == ndims);
198 4746 : for (int i = 0; i < ndims; i++) {
199 3184 : sizes[i] = tensor->size(i);
200 : }
201 :
202 1562 : assert(tensor->is_contiguous());
203 1562 : return tensor->data_ptr();
204 : };
205 :
206 : #ifdef __cplusplus
207 : extern "C" {
208 : #endif
209 :
210 : /*******************************************************************************
211 : * \brief Creates a Torch tensor from an array of int32s.
212 : * The passed array has to outlive the tensor!
213 : * \author Ole Schuett
214 : ******************************************************************************/
215 0 : void torch_c_tensor_from_array_int32(torch_c_tensor_t **tensor,
216 : const bool req_grad, const int ndims,
217 : const int64_t sizes[], int32_t source[]) {
218 0 : *tensor = tensor_from_array(torch::kInt32, req_grad, ndims, sizes, source);
219 0 : }
220 :
221 : /*******************************************************************************
222 : * \brief Creates a Torch tensor from an array of floats.
223 : * The passed array has to outlive the tensor!
224 : * \author Ole Schuett
225 : ******************************************************************************/
226 66 : void torch_c_tensor_from_array_float(torch_c_tensor_t **tensor,
227 : const bool req_grad, const int ndims,
228 : const int64_t sizes[], float source[]) {
229 66 : *tensor = tensor_from_array(torch::kFloat32, req_grad, ndims, sizes, source);
230 66 : }
231 :
232 : /*******************************************************************************
233 : * \brief Creates a Torch tensor from an array of int64s.
234 : * The passed array has to outlive the tensor!
235 : * \author Ole Schuett
236 : ******************************************************************************/
237 1014 : void torch_c_tensor_from_array_int64(torch_c_tensor_t **tensor,
238 : const bool req_grad, const int ndims,
239 : const int64_t sizes[], int64_t source[]) {
240 1014 : *tensor = tensor_from_array(torch::kInt64, req_grad, ndims, sizes, source);
241 1014 : }
242 :
243 : /*******************************************************************************
244 : * \brief Creates a Torch tensor from an array of doubles.
245 : * The passed array has to outlive the tensor!
246 : * \author Ole Schuett
247 : ******************************************************************************/
248 1496 : void torch_c_tensor_from_array_double(torch_c_tensor_t **tensor,
249 : const bool req_grad, const int ndims,
250 : const int64_t sizes[], double source[]) {
251 1496 : *tensor = tensor_from_array(torch::kFloat64, req_grad, ndims, sizes, source);
252 1496 : }
253 :
254 : /*******************************************************************************
255 : * \brief Releases a string returned from the Torch C API.
256 : ******************************************************************************/
257 216 : void torch_c_free_string(char *content) { free(content); }
258 :
259 : /*******************************************************************************
260 : * \brief Reuses or creates a device tensor and copies double data into it.
261 : ******************************************************************************/
262 870 : void torch_c_tensor_reset_from_array_double(torch_c_tensor_t **tensor,
263 : const bool req_grad,
264 : const int ndims,
265 : const int64_t sizes[],
266 : double source[]) {
267 870 : reset_tensor_from_array_double(tensor, req_grad, ndims, sizes, source);
268 870 : }
269 :
270 : /*******************************************************************************
271 : * \brief Creates an expanded tensor view along one singleton dimension.
272 : ******************************************************************************/
273 30 : void torch_c_tensor_expand_dim(const torch_c_tensor_t *tensor,
274 : const int64_t dim, const int64_t size,
275 : torch_c_tensor_t **result) {
276 30 : c10::OptionalDeviceGuard guard;
277 30 : get_device_with_guard(guard);
278 30 : assert(*result == NULL);
279 30 : assert(dim >= 0);
280 30 : assert(dim < tensor->dim());
281 30 : std::vector<int64_t> sizes(tensor->sizes().begin(), tensor->sizes().end());
282 30 : assert(sizes[dim] == 1);
283 30 : sizes[dim] = size;
284 30 : *result = new torch_c_tensor_t(tensor->expand(sizes));
285 30 : }
286 :
287 : /*******************************************************************************
288 : * \brief Creates a tensor view narrowed along one dimension.
289 : ******************************************************************************/
290 32 : void torch_c_tensor_narrow(const torch_c_tensor_t *tensor, const int64_t dim,
291 : const int64_t start_index, const int64_t length,
292 : torch_c_tensor_t **result) {
293 32 : c10::OptionalDeviceGuard guard;
294 32 : const auto device = get_device_with_guard(guard);
295 32 : assert(*result == NULL);
296 32 : assert(dim >= 0);
297 32 : assert(start_index >= 0);
298 32 : assert(length >= 0);
299 32 : assert(dim < tensor->ndimension());
300 32 : assert(start_index + length <= tensor->size(dim));
301 64 : *result =
302 32 : new torch_c_tensor_t(tensor->narrow(dim, start_index, length).to(device));
303 32 : }
304 :
305 : /*******************************************************************************
306 : * \brief Returns the data_ptr and sizes of a Torch tensor of int32s.
307 : * The returned pointer is only valide during the tensor's live time!
308 : * \author Ole Schuett
309 : ******************************************************************************/
310 0 : void torch_c_tensor_data_ptr_int32(const torch_c_tensor_t *tensor,
311 : const int ndims, int64_t sizes[],
312 : int32_t **data_ptr) {
313 0 : *data_ptr = (int32_t *)get_data_ptr(tensor, torch::kInt32, ndims, sizes);
314 0 : }
315 :
316 : /*******************************************************************************
317 : * \brief Returns the data_ptr and sizes of a Torch tensor of floats.
318 : * The returned pointer is only valide during the tensor's lifetime!
319 : * \author Ole Schuett
320 : ******************************************************************************/
321 66 : void torch_c_tensor_data_ptr_float(const torch_c_tensor_t *tensor,
322 : const int ndims, int64_t sizes[],
323 : float **data_ptr) {
324 66 : *data_ptr = (float *)get_data_ptr(tensor, torch::kFloat32, ndims, sizes);
325 66 : }
326 :
327 : /*******************************************************************************
328 : * \brief Returns the data_ptr and sizes of a Torch tensor of int64s.
329 : * The returned pointer is only valide during the tensor's live time!
330 : * \author Ole Schuett
331 : ******************************************************************************/
332 0 : void torch_c_tensor_data_ptr_int64(const torch_c_tensor_t *tensor,
333 : const int ndims, int64_t sizes[],
334 : int64_t **data_ptr) {
335 0 : *data_ptr = (int64_t *)get_data_ptr(tensor, torch::kInt64, ndims, sizes);
336 0 : }
337 :
338 : /*******************************************************************************
339 : * \brief Returns the data_ptr and sizes of a Torch tensor of doubles.
340 : * The returned pointer is only valide during the tensor's live time!
341 : * \author Ole Schuett
342 : ******************************************************************************/
343 1496 : void torch_c_tensor_data_ptr_double(const torch_c_tensor_t *tensor,
344 : const int ndims, int64_t sizes[],
345 : double **data_ptr) {
346 1496 : *data_ptr = (double *)get_data_ptr(tensor, torch::kFloat64, ndims, sizes);
347 1496 : }
348 :
349 : /*******************************************************************************
350 : * \brief Runs autograd on a Torch tensor.
351 : * \author Ole Schuett
352 : ******************************************************************************/
353 6 : void torch_c_tensor_backward(const torch_c_tensor_t *tensor,
354 : const torch_c_tensor_t *outer_grad) {
355 6 : TorchFloatingPointMaskGuard fpe_guard;
356 6 : c10::OptionalDeviceGuard guard;
357 6 : get_device_with_guard(guard);
358 6 : tensor->backward(*outer_grad);
359 6 : }
360 :
361 : /*******************************************************************************
362 : * \brief Runs autograd on a scalar Torch tensor.
363 : ******************************************************************************/
364 324 : void torch_c_tensor_backward_scalar(const torch_c_tensor_t *tensor) {
365 324 : TorchFloatingPointMaskGuard fpe_guard;
366 324 : c10::OptionalDeviceGuard guard;
367 324 : get_device_with_guard(guard);
368 648 : tensor->backward();
369 324 : }
370 :
371 : /*******************************************************************************
372 : * \brief Moves a tensor to the active device and makes it an autograd leaf.
373 : ******************************************************************************/
374 2314 : void torch_c_tensor_to_device_leaf(torch_c_tensor_t **tensor,
375 : const bool req_grad) {
376 2314 : c10::OptionalDeviceGuard guard;
377 2314 : const auto device = get_device_with_guard(guard);
378 4628 : auto moved = (*tensor)->to(device).detach();
379 2314 : moved.set_requires_grad(req_grad);
380 4628 : delete (*tensor);
381 4628 : *tensor = new torch_c_tensor_t(moved);
382 2314 : }
383 :
384 : /*******************************************************************************
385 : * \brief Select whether Torch wrappers should use CUDA when available.
386 : ******************************************************************************/
387 644 : void torch_c_use_cuda(const bool use_cuda) { use_cuda_if_available = use_cuda; }
388 :
389 : /*******************************************************************************
390 : * \brief Returns the gradient of a Torch tensor which was computed by autograd.
391 : * \author Ole Schuett
392 : ******************************************************************************/
393 1466 : void torch_c_tensor_grad(const torch_c_tensor_t *tensor,
394 : torch_c_tensor_t **grad) {
395 1466 : c10::OptionalDeviceGuard guard;
396 1466 : get_device_with_guard(guard);
397 1466 : const torch::Tensor maybe_grad = tensor->grad();
398 1466 : assert(maybe_grad.defined());
399 1466 : *grad = new torch_c_tensor_t(maybe_grad.cpu().contiguous());
400 1466 : }
401 :
402 : /*******************************************************************************
403 : * \brief Releases a Torch tensor and all its ressources.
404 : * \author Ole Schuett
405 : ******************************************************************************/
406 7396 : void torch_c_tensor_release(torch_c_tensor_t *tensor) { delete (tensor); }
407 :
408 : /*******************************************************************************
409 : * \brief Creates an empty Torch dictionary.
410 : * \author Ole Schuett
411 : ******************************************************************************/
412 476 : void torch_c_dict_create(torch_c_dict_t **dict_out) {
413 476 : assert(*dict_out == NULL);
414 476 : *dict_out = new c10::Dict<std::string, torch::Tensor>();
415 476 : }
416 :
417 : /*******************************************************************************
418 : * \brief Clones a Torch dictionary.
419 : ******************************************************************************/
420 128 : void torch_c_dict_clone(const torch_c_dict_t *dict, torch_c_dict_t **dict_out) {
421 128 : assert(*dict_out == NULL);
422 128 : torch_c_dict_t *clone = new c10::Dict<std::string, torch::Tensor>();
423 768 : for (const auto &entry : *dict) {
424 640 : clone->insert(entry.key(), entry.value());
425 : }
426 128 : *dict_out = clone;
427 128 : }
428 :
429 : /*******************************************************************************
430 : * \brief Inserts a Torch tensor into a Torch dictionary.
431 : * \author Ole Schuett
432 : ******************************************************************************/
433 2912 : void torch_c_dict_insert(const torch_c_dict_t *dict, const char *key,
434 : const torch_c_tensor_t *tensor) {
435 2912 : c10::OptionalDeviceGuard guard;
436 2912 : const auto device = get_device_with_guard(guard);
437 5824 : dict->insert(key, tensor->to(device));
438 2912 : }
439 :
440 : /*******************************************************************************
441 : * \brief Retrieves a Torch tensor from a Torch dictionary.
442 : * \author Ole Schuett
443 : ******************************************************************************/
444 76 : void torch_c_dict_get(const torch_c_dict_t *dict, const char *key,
445 : torch_c_tensor_t **tensor) {
446 76 : assert(dict->contains(key));
447 76 : *tensor = new torch_c_tensor_t(dict->at(key).cpu().contiguous());
448 76 : }
449 :
450 : /*******************************************************************************
451 : * \brief Releases a Torch dictionary and all its ressources.
452 : * \author Ole Schuett
453 : ******************************************************************************/
454 680 : void torch_c_dict_release(torch_c_dict_t *dict) { delete (dict); }
455 :
456 : /*******************************************************************************
457 : * \brief Loads a Torch model from given "*.pth" file.
458 : * In Torch lingo models are called modules.
459 : * \author Ole Schuett
460 : ******************************************************************************/
461 104 : void torch_c_model_load(torch_c_model_t **model_out, const char *filename) {
462 104 : assert(*model_out == NULL);
463 104 : c10::OptionalDeviceGuard guard;
464 104 : const auto device = get_device_with_guard(guard);
465 104 : set_jit_fusion_strategy();
466 104 : torch::jit::Module *model = new torch::jit::Module();
467 104 : *model = load_module_for_device(filename, device);
468 104 : model->eval(); // Set to evaluation mode to disable gradients, drop-out, etc.
469 104 : *model_out = model;
470 104 : }
471 :
472 : /*******************************************************************************
473 : * \brief Evaluates the given Torch model.
474 : * \author Ole Schuett
475 : ******************************************************************************/
476 62 : void torch_c_model_forward(torch_c_model_t *model, const torch_c_dict_t *inputs,
477 : torch_c_dict_t *outputs) {
478 :
479 62 : TorchFloatingPointMaskGuard fpe_guard;
480 62 : c10::OptionalDeviceGuard guard;
481 62 : get_device_with_guard(guard);
482 310 : auto untyped_output = model->forward({*inputs}).toGenericDict();
483 62 : outputs->clear();
484 232 : for (const auto &entry : untyped_output) {
485 170 : outputs->insert(entry.key().toStringView(), entry.value().toTensor());
486 : }
487 186 : }
488 :
489 : /*******************************************************************************
490 : * \brief Evaluates a TorchScript model method expecting argument "mol".
491 : ******************************************************************************/
492 324 : void torch_c_model_forward_mol_tensor(torch_c_model_t *model,
493 : const char *method_name,
494 : const torch_c_dict_t *inputs,
495 : torch_c_tensor_t **output) {
496 :
497 324 : c10::OptionalDeviceGuard guard;
498 324 : get_device_with_guard(guard);
499 324 : assert(*output == NULL);
500 324 : *output = new torch_c_tensor_t(
501 1620 : model->get_method(method_name)({*inputs}).toTensor());
502 972 : }
503 :
504 : /*******************************************************************************
505 : * \brief Returns the weighted sum of two Torch tensors.
506 : ******************************************************************************/
507 324 : void torch_c_tensor_weighted_sum(const torch_c_tensor_t *values,
508 : const torch_c_tensor_t *weights,
509 : torch_c_tensor_t **result) {
510 324 : c10::OptionalDeviceGuard guard;
511 324 : get_device_with_guard(guard);
512 324 : const auto weights_on_device = weights->to(values->device());
513 648 : *result = new torch_c_tensor_t((*values * weights_on_device).sum());
514 324 : }
515 :
516 : /*******************************************************************************
517 : * \brief Returns a scalar double value from a Torch tensor.
518 : ******************************************************************************/
519 324 : double torch_c_tensor_item_double(const torch_c_tensor_t *tensor) {
520 324 : c10::OptionalDeviceGuard guard;
521 324 : get_device_with_guard(guard);
522 324 : return tensor->item<double>();
523 324 : }
524 :
525 : /*******************************************************************************
526 : * \brief Releases a Torch model and all its ressources.
527 : * \author Ole Schuett
528 : ******************************************************************************/
529 16 : void torch_c_model_release(torch_c_model_t *model) { delete (model); }
530 :
531 : /*******************************************************************************
532 : * \brief Reads metadata entry from given "*.pth" file.
533 : * In Torch lingo they are called extra files.
534 : * The returned char array has to be deallocated by caller!
535 : * \author Ole Schuett
536 : ******************************************************************************/
537 216 : void torch_c_model_read_metadata(const char *filename, const char *key,
538 : char **content, int *length) {
539 :
540 648 : std::unordered_map<std::string, std::string> extra_files = {{key, ""}};
541 216 : torch::jit::load(filename, torch::kCPU, extra_files);
542 432 : const std::string &content_str = extra_files[key];
543 216 : copy_string_to_c_buffer(content_str, content, length);
544 432 : }
545 :
546 : /*******************************************************************************
547 : * \brief Returns true iff the Torch CUDA backend is available.
548 : * \author Ole Schuett
549 : ******************************************************************************/
550 2 : bool torch_c_cuda_is_available() { return torch::cuda::is_available(); }
551 :
552 : /*******************************************************************************
553 : * \brief Return the number of CUDA devices visible to Torch.
554 : ******************************************************************************/
555 0 : int torch_c_cuda_device_count() {
556 0 : return torch::cuda::is_available() ? torch::cuda::device_count() : 0;
557 : }
558 :
559 : /*******************************************************************************
560 : * \brief Set whether to allow TF32.
561 : * Needed due to changes in defaults from pytorch 1.7 to 1.11 to >=1.12
562 : * See https://pytorch.org/docs/stable/notes/cuda.html
563 : * \author Gabriele Tocci
564 : ******************************************************************************/
565 6 : void torch_c_allow_tf32(const bool allow_tf32) {
566 :
567 6 : at::globalContext().setAllowTF32CuBLAS(allow_tf32);
568 6 : at::globalContext().setAllowTF32CuDNN(allow_tf32);
569 6 : }
570 :
571 : /******************************************************************************
572 : * \brief Freeze the Torch model: generic optimization that speeds up model.
573 : * See https://pytorch.org/docs/stable/generated/torch.jit.freeze.html
574 : * \author Gabriele Tocci
575 : ******************************************************************************/
576 6 : void torch_c_model_freeze(torch_c_model_t *model) {
577 :
578 6 : *model = torch::jit::freeze(*model);
579 6 : }
580 :
581 : /*******************************************************************************
582 : * \brief Retrieves an int64 attribute. Must be called before model freeze.
583 : * \author Ole Schuett
584 : ******************************************************************************/
585 40 : void torch_c_model_get_attr_int64(const torch_c_model_t *model, const char *key,
586 : int64_t *dest) {
587 40 : *dest = model->attr(key).toInt();
588 40 : }
589 :
590 : /*******************************************************************************
591 : * \brief Retrieves a double attribute. Must be called before model freeze.
592 : * \author Ole Schuett
593 : ******************************************************************************/
594 8 : void torch_c_model_get_attr_double(const torch_c_model_t *model,
595 : const char *key, double *dest) {
596 8 : *dest = model->attr(key).toDouble();
597 8 : }
598 :
599 : /*******************************************************************************
600 : * \brief Retrieves a string attribute. Must be called before model freeze.
601 : * \author Ole Schuett
602 : ******************************************************************************/
603 16 : void torch_c_model_get_attr_string(const torch_c_model_t *model,
604 : const char *key, char *dest) {
605 16 : const std::string &str = model->attr(key).toStringRef();
606 16 : assert(str.size() < 80); // default_string_length
607 144 : for (int i = 0; i < str.size(); i++) {
608 128 : dest[i] = str[i];
609 : }
610 16 : }
611 :
612 : /*******************************************************************************
613 : * \brief Retrieves a list attribute's size. Must be called before model freeze.
614 : * \author Ole Schuett
615 : ******************************************************************************/
616 8 : void torch_c_model_get_attr_list_size(const torch_c_model_t *model,
617 : const char *key, int *size) {
618 8 : *size = model->attr(key).toList().size();
619 8 : }
620 :
621 : /*******************************************************************************
622 : * \brief Retrieves a single item from a string list attribute.
623 : * \author Ole Schuett
624 : ******************************************************************************/
625 16 : void torch_c_model_get_attr_strlist(const torch_c_model_t *model,
626 : const char *key, const int index,
627 : char *dest) {
628 32 : const auto list = model->attr(key).toList();
629 16 : const std::string &str = list[index].toStringRef();
630 16 : assert(str.size() < 80); // default_string_length
631 32 : for (int i = 0; i < str.size(); i++) {
632 16 : dest[i] = str[i];
633 : }
634 16 : }
635 :
636 : #ifdef __cplusplus
637 : }
638 : #endif
639 :
640 : #endif // defined(__LIBTORCH)
641 :
642 : // EOF
|
