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 : #if defined(__LIBTORCH_CUDA)
12 : #include <ATen/cuda/CUDAContextLight.h>
13 : #include <c10/cuda/CUDAAllocatorConfig.h>
14 : #endif
15 : #include <c10/core/DeviceGuard.h>
16 : #include <torch/csrc/api/include/torch/cuda.h>
17 : #include <torch/script.h>
18 :
19 : #include "offload/offload_library.h"
20 :
21 : #if defined(__OPENBLAS)
22 : #include <cblas.h>
23 : #endif
24 :
25 : #include <cassert>
26 :
27 : #include <cfenv>
28 : #include <climits>
29 : #include <cstdlib>
30 : #include <cstring>
31 : #include <string>
32 : #include <unordered_map>
33 : #include <vector>
34 :
35 : #if defined(__OPENBLAS)
36 : // PyTorch's oneMKL batch ABI is not compatible with OpenBLAS's same-named
37 : // entry points. Expand grouped GEMMs into the portable CBLAS interface.
38 21159 : extern "C" void cblas_sgemm_batch(
39 : const enum CBLAS_ORDER order, const enum CBLAS_TRANSPOSE *trans_a,
40 : const enum CBLAS_TRANSPOSE *trans_b, const int *m, const int *n,
41 : const int *k, const float *alpha, const float **a, const int *lda,
42 : const float **b, const int *ldb, const float *beta, float **c,
43 : const int *ldc, const int group_count, const int *group_size) {
44 21159 : int offset = 0;
45 42318 : for (int group = 0; group < group_count; ++group) {
46 390381 : for (int operation = 0; operation < group_size[group]; ++operation) {
47 369222 : const int index = offset + operation;
48 369222 : cblas_sgemm(order, trans_a[group], trans_b[group], m[group], n[group],
49 369222 : k[group], alpha[group], a[index], lda[group], b[index],
50 369222 : ldb[group], beta[group], c[index], ldc[group]);
51 : }
52 21159 : offset += group_size[group];
53 : }
54 21159 : }
55 :
56 945 : extern "C" void cblas_dgemm_batch(
57 : const enum CBLAS_ORDER order, const enum CBLAS_TRANSPOSE *trans_a,
58 : const enum CBLAS_TRANSPOSE *trans_b, const int *m, const int *n,
59 : const int *k, const double *alpha, const double **a, const int *lda,
60 : const double **b, const int *ldb, const double *beta, double **c,
61 : const int *ldc, const int group_count, const int *group_size) {
62 945 : int offset = 0;
63 1890 : for (int group = 0; group < group_count; ++group) {
64 4125 : for (int operation = 0; operation < group_size[group]; ++operation) {
65 3180 : const int index = offset + operation;
66 3180 : cblas_dgemm(order, trans_a[group], trans_b[group], m[group], n[group],
67 3180 : k[group], alpha[group], a[index], lda[group], b[index],
68 3180 : ldb[group], beta[group], c[index], ldc[group]);
69 : }
70 945 : offset += group_size[group];
71 : }
72 945 : }
73 : #endif
74 :
75 : typedef torch::Tensor torch_c_tensor_t;
76 : typedef c10::Dict<std::string, torch::Tensor> torch_c_dict_t;
77 : typedef torch::jit::Module torch_c_model_t;
78 :
79 : class TorchFloatingPointMaskGuard {
80 : public:
81 430 : TorchFloatingPointMaskGuard() : active_(std::feholdexcept(&env_) == 0) {}
82 430 : ~TorchFloatingPointMaskGuard() {
83 430 : if (active_) {
84 430 : std::feclearexcept(FE_ALL_EXCEPT);
85 430 : std::fesetenv(&env_);
86 : }
87 430 : }
88 :
89 : private:
90 : std::fenv_t env_;
91 : bool active_;
92 : };
93 :
94 : /*******************************************************************************
95 : * \brief Internal helper for selecting the CUDA device when available.
96 : * \author Ole Schuett
97 : ******************************************************************************/
98 : static bool use_cuda_if_available = true;
99 :
100 0 : static void enable_expandable_cuda_segments_if_unconfigured() {
101 : #if defined(__LIBTORCH_CUDA)
102 : static const bool initialized = []() {
103 : const char *legacy_config = std::getenv("PYTORCH_CUDA_ALLOC_CONF");
104 : const char *config = std::getenv("PYTORCH_ALLOC_CONF");
105 : if ((legacy_config == nullptr || legacy_config[0] == '\0') &&
106 : (config == nullptr || config[0] == '\0')) {
107 : c10::cuda::CUDACachingAllocator::setAllocatorSettings(
108 : "expandable_segments:True");
109 : }
110 : return true;
111 : }();
112 : (void)initialized;
113 : #endif
114 0 : }
115 :
116 218 : static bool get_positive_int_env(const char *name, int &value) {
117 218 : const char *raw = std::getenv(name);
118 218 : if (raw == nullptr || raw[0] == '\0') {
119 : return false;
120 : }
121 0 : char *end = nullptr;
122 0 : const long parsed = std::strtol(raw, &end, 10);
123 0 : if (end == raw || *end != '\0' || parsed <= 0 || parsed > INT_MAX) {
124 : return false;
125 : }
126 0 : value = static_cast<int>(parsed);
127 0 : return true;
128 : }
129 :
130 11636 : static void initialize_torch_threads_from_env() {
131 11636 : static bool initialized = false;
132 11636 : if (initialized) {
133 11527 : return;
134 : }
135 109 : initialized = true;
136 :
137 109 : int num_threads = 0;
138 109 : if (get_positive_int_env("CP2K_TORCH_NUM_THREADS", num_threads)) {
139 0 : at::set_num_threads(num_threads);
140 : }
141 109 : if (get_positive_int_env("CP2K_TORCH_NUM_INTEROP_THREADS", num_threads)) {
142 0 : at::set_num_interop_threads(num_threads);
143 : }
144 : }
145 :
146 9194 : static torch::Device get_device() {
147 9194 : initialize_torch_threads_from_env();
148 9194 : if (!use_cuda_if_available || !torch::cuda::is_available()) {
149 9194 : return torch::kCPU;
150 : }
151 0 : enable_expandable_cuda_segments_if_unconfigured();
152 0 : const auto device_count = torch::cuda::device_count();
153 0 : if (device_count <= 0) {
154 0 : return torch::kCPU;
155 : }
156 0 : const int chosen_device = offload_get_chosen_device();
157 0 : const int device = (chosen_device >= 0) ? chosen_device : 0;
158 0 : assert(device < device_count);
159 0 : return torch::Device(torch::kCUDA, device);
160 : }
161 :
162 9194 : static torch::Device get_device_with_guard(c10::OptionalDeviceGuard &guard) {
163 9194 : const auto device = get_device();
164 9194 : if (device.is_cuda()) {
165 0 : guard.reset_device(device);
166 : }
167 9194 : return device;
168 : }
169 :
170 360 : static bool use_batched_gradient_readback(const torch::Tensor &tensor) {
171 : #if defined(__LIBTORCH_CUDA)
172 : return tensor.is_cuda() &&
173 : at::cuda::getDeviceProperties(tensor.device().index())->integrated;
174 : #else
175 360 : (void)tensor;
176 360 : return false;
177 : #endif
178 : }
179 :
180 113 : static void set_jit_fusion_strategy() {
181 : // JIT Fusion strategy optimization, hardcode dynamic 10, see also
182 : // https://github.com/mir-group/pair_nequip_allegro.git
183 113 : torch::jit::FusionStrategy strategy = {
184 113 : {torch::jit::FusionBehavior::DYNAMIC, 10}};
185 226 : torch::jit::setFusionStrategy(strategy);
186 113 : }
187 :
188 234 : static void copy_string_to_c_buffer(const std::string &source, char **content,
189 : int *length) {
190 234 : *length = source.length();
191 234 : *content = (char *)malloc(source.length() + 1); // +1 for null terminator
192 234 : strcpy(*content, source.c_str());
193 234 : }
194 :
195 113 : static bool can_load_directly_to_device(const torch::Device &device) {
196 113 : return !device.is_cuda() || device.index() == 0 ||
197 0 : torch::cuda::device_count() == 1;
198 : }
199 :
200 113 : static torch::jit::Module load_module_for_device(
201 : const char *filename, const torch::Device &device,
202 : std::unordered_map<std::string, std::string> *extra_files = nullptr) {
203 113 : if (can_load_directly_to_device(device)) {
204 113 : if (extra_files != nullptr) {
205 194 : return torch::jit::load(filename, device, *extra_files);
206 : }
207 32 : return torch::jit::load(filename, device);
208 : }
209 0 : auto model = (extra_files != nullptr)
210 0 : ? torch::jit::load(filename, torch::kCPU, *extra_files)
211 0 : : torch::jit::load(filename, torch::kCPU);
212 0 : model.to(device);
213 0 : return model;
214 0 : }
215 :
216 0 : static void remap_device_constants(torch::jit::Block *block,
217 : const torch::Device &device) {
218 0 : for (torch::jit::Node *node : block->nodes()) {
219 0 : if (node->kind() == torch::jit::prim::Constant &&
220 0 : node->outputs().size() == 1 &&
221 0 : node->output()->type()->kind() == c10::TypeKind::DeviceObjType &&
222 0 : node->hasAttribute(torch::jit::attr::value) &&
223 0 : node->kindOf(torch::jit::attr::value) == torch::jit::AttributeKind::s) {
224 0 : node->s_(torch::jit::attr::value, device.str());
225 : }
226 0 : for (torch::jit::Block *nested_block : node->blocks()) {
227 0 : remap_device_constants(nested_block, device);
228 : }
229 : }
230 0 : }
231 :
232 0 : static void remap_model_device_constants(torch::jit::Module &model,
233 : const torch::Device &device) {
234 0 : for (const auto &method : model.get_methods()) {
235 0 : remap_device_constants(method.graph()->block(), device);
236 0 : }
237 0 : for (auto child : model.children()) {
238 0 : remap_model_device_constants(child, device);
239 0 : }
240 0 : }
241 :
242 : /*******************************************************************************
243 : * \brief Internal helper for creating a Torch tensor from an array.
244 : * \author Ole Schuett
245 : ******************************************************************************/
246 2442 : static torch_c_tensor_t *tensor_from_array(const torch::Dtype dtype,
247 : const bool req_grad, const int ndims,
248 : const int64_t sizes[],
249 : void *source) {
250 2442 : initialize_torch_threads_from_env();
251 2442 : const auto opts = torch::TensorOptions().dtype(dtype).requires_grad(req_grad);
252 2442 : const auto sizes_ref = c10::IntArrayRef(sizes, ndims);
253 2442 : return new torch_c_tensor_t(torch::from_blob(source, sizes_ref, opts));
254 : }
255 :
256 966 : static bool tensor_matches(const torch_c_tensor_t *tensor,
257 : const torch::Dtype dtype,
258 : const torch::Device &device, const int ndims,
259 : const int64_t sizes[]) {
260 756 : if (tensor == nullptr || !tensor->defined() ||
261 1722 : tensor->scalar_type() != dtype || tensor->device() != device ||
262 378 : tensor->ndimension() != ndims) {
263 588 : return false;
264 : }
265 1260 : for (int i = 0; i < ndims; i++) {
266 882 : if (tensor->size(i) != sizes[i]) {
267 : return false;
268 : }
269 : }
270 378 : return tensor->is_contiguous();
271 : }
272 :
273 966 : static void reset_tensor_from_array_double(torch_c_tensor_t **tensor,
274 : const bool req_grad, const int ndims,
275 : const int64_t sizes[],
276 : double source[]) {
277 966 : c10::OptionalDeviceGuard guard;
278 966 : const auto device = get_device_with_guard(guard);
279 966 : const auto sizes_ref = c10::IntArrayRef(sizes, ndims);
280 966 : if (!tensor_matches(*tensor, torch::kFloat64, device, ndims, sizes)) {
281 588 : delete (*tensor);
282 588 : const auto opts =
283 588 : torch::TensorOptions().dtype(torch::kFloat64).device(device);
284 1176 : *tensor = new torch_c_tensor_t(torch::empty(sizes_ref, opts).detach());
285 : }
286 966 : const auto source_tensor = torch::from_blob(
287 966 : source, sizes_ref, torch::TensorOptions().dtype(torch::kFloat64));
288 966 : {
289 966 : torch::NoGradGuard no_grad;
290 966 : (*tensor)->copy_(source_tensor);
291 966 : (*tensor)->mutable_grad() = torch::Tensor();
292 0 : }
293 1932 : (*tensor)->set_requires_grad(req_grad);
294 966 : }
295 :
296 : /*******************************************************************************
297 : * \brief Internal helper for getting the data_ptr and sizes of a Torch tensor.
298 : * \author Ole Schuett
299 : ******************************************************************************/
300 1730 : static void *get_data_ptr(const torch_c_tensor_t *tensor,
301 : const torch::Dtype dtype, const int ndims,
302 : int64_t sizes[]) {
303 1730 : assert(tensor->scalar_type() == dtype);
304 1730 : assert(tensor->ndimension() == ndims);
305 5258 : for (int i = 0; i < ndims; i++) {
306 3528 : sizes[i] = tensor->size(i);
307 : }
308 :
309 1730 : assert(tensor->is_contiguous());
310 1730 : return tensor->data_ptr();
311 : };
312 :
313 : #ifdef __cplusplus
314 : extern "C" {
315 : #endif
316 :
317 : /*******************************************************************************
318 : * \brief Creates a Torch tensor from an array of int32s.
319 : * The passed array has to outlive the tensor!
320 : * \author Ole Schuett
321 : ******************************************************************************/
322 0 : void torch_c_tensor_from_array_int32(torch_c_tensor_t **tensor,
323 : const bool req_grad, const int ndims,
324 : const int64_t sizes[], int32_t source[]) {
325 0 : *tensor = tensor_from_array(torch::kInt32, req_grad, ndims, sizes, source);
326 0 : }
327 :
328 : /*******************************************************************************
329 : * \brief Creates a Torch tensor from an array of floats.
330 : * The passed array has to outlive the tensor!
331 : * \author Ole Schuett
332 : ******************************************************************************/
333 66 : void torch_c_tensor_from_array_float(torch_c_tensor_t **tensor,
334 : const bool req_grad, const int ndims,
335 : const int64_t sizes[], float source[]) {
336 66 : *tensor = tensor_from_array(torch::kFloat32, req_grad, ndims, sizes, source);
337 66 : }
338 :
339 : /*******************************************************************************
340 : * \brief Creates a Torch tensor from an array of int64s.
341 : * The passed array has to outlive the tensor!
342 : * \author Ole Schuett
343 : ******************************************************************************/
344 894 : void torch_c_tensor_from_array_int64(torch_c_tensor_t **tensor,
345 : const bool req_grad, const int ndims,
346 : const int64_t sizes[], int64_t source[]) {
347 894 : *tensor = tensor_from_array(torch::kInt64, req_grad, ndims, sizes, source);
348 894 : }
349 :
350 : /*******************************************************************************
351 : * \brief Creates a Torch tensor from an array of doubles.
352 : * The passed array has to outlive the tensor!
353 : * \author Ole Schuett
354 : ******************************************************************************/
355 1482 : void torch_c_tensor_from_array_double(torch_c_tensor_t **tensor,
356 : const bool req_grad, const int ndims,
357 : const int64_t sizes[], double source[]) {
358 1482 : *tensor = tensor_from_array(torch::kFloat64, req_grad, ndims, sizes, source);
359 1482 : }
360 :
361 : /*******************************************************************************
362 : * \brief Releases a string returned from the Torch C API.
363 : ******************************************************************************/
364 234 : void torch_c_free_string(char *content) { free(content); }
365 :
366 : /*******************************************************************************
367 : * \brief Reuses or creates a device tensor and copies double data into it.
368 : ******************************************************************************/
369 966 : void torch_c_tensor_reset_from_array_double(torch_c_tensor_t **tensor,
370 : const bool req_grad,
371 : const int ndims,
372 : const int64_t sizes[],
373 : double source[]) {
374 966 : reset_tensor_from_array_double(tensor, req_grad, ndims, sizes, source);
375 966 : }
376 :
377 : /*******************************************************************************
378 : * \brief Creates an expanded tensor view along one singleton dimension.
379 : ******************************************************************************/
380 372 : void torch_c_tensor_expand_dim(const torch_c_tensor_t *tensor,
381 : const int64_t dim, const int64_t size,
382 : torch_c_tensor_t **result) {
383 372 : c10::OptionalDeviceGuard guard;
384 372 : get_device_with_guard(guard);
385 372 : assert(*result == NULL);
386 372 : assert(dim >= 0);
387 372 : assert(dim < tensor->dim());
388 372 : std::vector<int64_t> sizes(tensor->sizes().begin(), tensor->sizes().end());
389 372 : assert(sizes[dim] == 1);
390 372 : sizes[dim] = size;
391 372 : *result = new torch_c_tensor_t(tensor->expand(sizes));
392 372 : }
393 :
394 : /*******************************************************************************
395 : * \brief Creates a tensor view narrowed along one dimension.
396 : ******************************************************************************/
397 32 : void torch_c_tensor_narrow(const torch_c_tensor_t *tensor, const int64_t dim,
398 : const int64_t start_index, const int64_t length,
399 : torch_c_tensor_t **result) {
400 32 : c10::OptionalDeviceGuard guard;
401 32 : const auto device = get_device_with_guard(guard);
402 32 : assert(*result == NULL);
403 32 : assert(dim >= 0);
404 32 : assert(start_index >= 0);
405 32 : assert(length >= 0);
406 32 : assert(dim < tensor->ndimension());
407 32 : assert(start_index + length <= tensor->size(dim));
408 64 : *result =
409 32 : new torch_c_tensor_t(tensor->narrow(dim, start_index, length).to(device));
410 32 : }
411 :
412 : /*******************************************************************************
413 : * \brief Returns the data_ptr and sizes of a Torch tensor of int32s.
414 : * The returned pointer is only valide during the tensor's live time!
415 : * \author Ole Schuett
416 : ******************************************************************************/
417 0 : void torch_c_tensor_data_ptr_int32(const torch_c_tensor_t *tensor,
418 : const int ndims, int64_t sizes[],
419 : int32_t **data_ptr) {
420 0 : *data_ptr = (int32_t *)get_data_ptr(tensor, torch::kInt32, ndims, sizes);
421 0 : }
422 :
423 : /*******************************************************************************
424 : * \brief Returns the data_ptr and sizes of a Torch tensor of floats.
425 : * The returned pointer is only valide during the tensor's lifetime!
426 : * \author Ole Schuett
427 : ******************************************************************************/
428 66 : void torch_c_tensor_data_ptr_float(const torch_c_tensor_t *tensor,
429 : const int ndims, int64_t sizes[],
430 : float **data_ptr) {
431 66 : *data_ptr = (float *)get_data_ptr(tensor, torch::kFloat32, ndims, sizes);
432 66 : }
433 :
434 : /*******************************************************************************
435 : * \brief Returns the data_ptr and sizes of a Torch tensor of int64s.
436 : * The returned pointer is only valide during the tensor's live time!
437 : * \author Ole Schuett
438 : ******************************************************************************/
439 0 : void torch_c_tensor_data_ptr_int64(const torch_c_tensor_t *tensor,
440 : const int ndims, int64_t sizes[],
441 : int64_t **data_ptr) {
442 0 : *data_ptr = (int64_t *)get_data_ptr(tensor, torch::kInt64, ndims, sizes);
443 0 : }
444 :
445 : /*******************************************************************************
446 : * \brief Returns the data_ptr and sizes of a Torch tensor of doubles.
447 : * The returned pointer is only valide during the tensor's live time!
448 : * \author Ole Schuett
449 : ******************************************************************************/
450 1664 : void torch_c_tensor_data_ptr_double(const torch_c_tensor_t *tensor,
451 : const int ndims, int64_t sizes[],
452 : double **data_ptr) {
453 1664 : *data_ptr = (double *)get_data_ptr(tensor, torch::kFloat64, ndims, sizes);
454 1664 : }
455 :
456 : /*******************************************************************************
457 : * \brief Runs autograd on a Torch tensor.
458 : * \author Ole Schuett
459 : ******************************************************************************/
460 6 : void torch_c_tensor_backward(const torch_c_tensor_t *tensor,
461 : const torch_c_tensor_t *outer_grad) {
462 6 : TorchFloatingPointMaskGuard fpe_guard;
463 6 : c10::OptionalDeviceGuard guard;
464 6 : get_device_with_guard(guard);
465 6 : tensor->backward(*outer_grad);
466 6 : }
467 :
468 : /*******************************************************************************
469 : * \brief Runs autograd on a scalar Torch tensor.
470 : ******************************************************************************/
471 362 : void torch_c_tensor_backward_scalar(const torch_c_tensor_t *tensor) {
472 362 : TorchFloatingPointMaskGuard fpe_guard;
473 362 : c10::OptionalDeviceGuard guard;
474 362 : get_device_with_guard(guard);
475 724 : tensor->backward();
476 362 : }
477 :
478 : /*******************************************************************************
479 : * \brief Moves a tensor to the active device and makes it an autograd leaf.
480 : ******************************************************************************/
481 2180 : void torch_c_tensor_to_device_leaf(torch_c_tensor_t **tensor,
482 : const bool req_grad) {
483 2180 : c10::OptionalDeviceGuard guard;
484 2180 : const auto device = get_device_with_guard(guard);
485 4360 : auto moved = (*tensor)->to(device).detach();
486 2180 : moved.set_requires_grad(req_grad);
487 4360 : delete (*tensor);
488 4360 : *tensor = new torch_c_tensor_t(moved);
489 2180 : }
490 :
491 : /*******************************************************************************
492 : * \brief Select whether Torch wrappers should use CUDA when available.
493 : ******************************************************************************/
494 772 : void torch_c_use_cuda(const bool use_cuda) { use_cuda_if_available = use_cuda; }
495 :
496 : /*******************************************************************************
497 : * \brief Returns the gradient of a Torch tensor which was computed by autograd.
498 : * \author Ole Schuett
499 : ******************************************************************************/
500 552 : void torch_c_tensor_grad(const torch_c_tensor_t *tensor,
501 : torch_c_tensor_t **grad) {
502 552 : c10::OptionalDeviceGuard guard;
503 552 : get_device_with_guard(guard);
504 552 : const torch::Tensor maybe_grad = tensor->grad();
505 552 : assert(maybe_grad.defined());
506 1104 : torch::Tensor host_grad = maybe_grad.detach().cpu().contiguous();
507 552 : if (maybe_grad.is_cpu()) {
508 552 : host_grad = host_grad.clone();
509 : }
510 552 : *grad = new torch_c_tensor_t(std::move(host_grad));
511 552 : }
512 :
513 : /*******************************************************************************
514 : * \brief Copies three autograd gradients to CPU memory.
515 : ******************************************************************************/
516 360 : void torch_c_tensor_grad_batch3(const torch_c_tensor_t *tensor1,
517 : const torch_c_tensor_t *tensor2,
518 : const torch_c_tensor_t *tensor3,
519 : torch_c_tensor_t **grad1,
520 : torch_c_tensor_t **grad2,
521 : torch_c_tensor_t **grad3) {
522 360 : c10::OptionalDeviceGuard guard;
523 360 : const auto device = get_device_with_guard(guard);
524 360 : assert(*grad1 == nullptr && *grad2 == nullptr && *grad3 == nullptr);
525 :
526 360 : const torch::Tensor source1 = tensor1->grad();
527 360 : const torch::Tensor source2 = tensor2->grad();
528 360 : const torch::Tensor source3 = tensor3->grad();
529 360 : assert(source1.defined() && source2.defined() && source3.defined());
530 360 : assert(source1.device() == source2.device());
531 360 : assert(source1.device() == source3.device());
532 360 : if (use_batched_gradient_readback(source1)) {
533 0 : auto host1 =
534 : torch::empty(source1.sizes(),
535 0 : source1.options().device(torch::kCPU).pinned_memory(true));
536 0 : auto host2 =
537 : torch::empty(source2.sizes(),
538 0 : source2.options().device(torch::kCPU).pinned_memory(true));
539 0 : auto host3 =
540 : torch::empty(source3.sizes(),
541 0 : source3.options().device(torch::kCPU).pinned_memory(true));
542 0 : host1.copy_(source1, true);
543 0 : host2.copy_(source2, true);
544 0 : host3.copy_(source3, true);
545 0 : torch::cuda::synchronize(device.index());
546 0 : *grad1 = new torch_c_tensor_t(std::move(host1));
547 0 : *grad2 = new torch_c_tensor_t(std::move(host2));
548 0 : *grad3 = new torch_c_tensor_t(std::move(host3));
549 0 : } else {
550 : // Materialize independent host buffers instead of aliasing gradients owned
551 : // by the autograd graph when they are already contiguous CPU tensors.
552 360 : *grad1 = new torch_c_tensor_t(source1.detach().cpu().contiguous().clone());
553 360 : *grad2 = new torch_c_tensor_t(source2.detach().cpu().contiguous().clone());
554 360 : *grad3 = new torch_c_tensor_t(source3.detach().cpu().contiguous().clone());
555 : }
556 360 : }
557 :
558 : /*******************************************************************************
559 : * \brief Releases a Torch tensor and all its ressources.
560 : * \author Ole Schuett
561 : ******************************************************************************/
562 8404 : void torch_c_tensor_release(torch_c_tensor_t *tensor) { delete (tensor); }
563 :
564 : /*******************************************************************************
565 : * \brief Creates an empty Torch dictionary.
566 : * \author Ole Schuett
567 : ******************************************************************************/
568 450 : void torch_c_dict_create(torch_c_dict_t **dict_out) {
569 450 : assert(*dict_out == NULL);
570 450 : *dict_out = new c10::Dict<std::string, torch::Tensor>();
571 450 : }
572 :
573 : /*******************************************************************************
574 : * \brief Clones a Torch dictionary.
575 : ******************************************************************************/
576 132 : void torch_c_dict_clone(const torch_c_dict_t *dict, torch_c_dict_t **dict_out) {
577 132 : assert(*dict_out == NULL);
578 132 : torch_c_dict_t *clone = new c10::Dict<std::string, torch::Tensor>();
579 792 : for (const auto &entry : *dict) {
580 660 : clone->insert(entry.key(), entry.value());
581 : }
582 132 : *dict_out = clone;
583 132 : }
584 :
585 : /*******************************************************************************
586 : * \brief Inserts a Torch tensor into a Torch dictionary.
587 : * \author Ole Schuett
588 : ******************************************************************************/
589 2934 : void torch_c_dict_insert(const torch_c_dict_t *dict, const char *key,
590 : const torch_c_tensor_t *tensor) {
591 2934 : c10::OptionalDeviceGuard guard;
592 2934 : const auto device = get_device_with_guard(guard);
593 5868 : dict->insert(key, tensor->to(device));
594 2934 : }
595 :
596 : /*******************************************************************************
597 : * \brief Retrieves a Torch tensor from a Torch dictionary.
598 : * \author Ole Schuett
599 : ******************************************************************************/
600 76 : void torch_c_dict_get(const torch_c_dict_t *dict, const char *key,
601 : torch_c_tensor_t **tensor) {
602 76 : assert(dict->contains(key));
603 76 : *tensor = new torch_c_tensor_t(dict->at(key).cpu().contiguous());
604 76 : }
605 :
606 : /*******************************************************************************
607 : * \brief Releases a Torch dictionary and all its ressources.
608 : * \author Ole Schuett
609 : ******************************************************************************/
610 688 : void torch_c_dict_release(torch_c_dict_t *dict) { delete (dict); }
611 :
612 : /*******************************************************************************
613 : * \brief Loads a Torch model from given "*.pth" file.
614 : * In Torch lingo models are called modules.
615 : * \author Ole Schuett
616 : ******************************************************************************/
617 16 : void torch_c_model_load(torch_c_model_t **model_out, const char *filename) {
618 16 : assert(*model_out == NULL);
619 16 : c10::OptionalDeviceGuard guard;
620 16 : const auto device = get_device_with_guard(guard);
621 16 : set_jit_fusion_strategy();
622 16 : torch::jit::Module *model = new torch::jit::Module();
623 16 : *model = load_module_for_device(filename, device);
624 16 : model->eval(); // Set inference behavior for modules such as dropout.
625 16 : *model_out = model;
626 16 : }
627 :
628 : /*******************************************************************************
629 : * \brief Loads a Torch model and reads two metadata entries.
630 : ******************************************************************************/
631 97 : void torch_c_model_load_with_metadata(torch_c_model_t **model_out,
632 : const char *filename, const char *key1,
633 : const char *key2, char **content1,
634 : int *length1, char **content2,
635 : int *length2) {
636 97 : assert(*model_out == NULL);
637 97 : c10::OptionalDeviceGuard guard;
638 97 : const auto device = get_device_with_guard(guard);
639 97 : std::unordered_map<std::string, std::string> extra_files = {{key1, ""},
640 388 : {key2, ""}};
641 97 : set_jit_fusion_strategy();
642 97 : torch::jit::Module *model = new torch::jit::Module();
643 97 : *model = load_module_for_device(filename, device, &extra_files);
644 97 : model->eval(); // Set inference behavior for modules such as dropout.
645 97 : *model_out = model;
646 194 : copy_string_to_c_buffer(extra_files[key1], content1, length1);
647 194 : copy_string_to_c_buffer(extra_files[key2], content2, length2);
648 194 : }
649 :
650 : /*******************************************************************************
651 : * \brief Maps serialized TorchScript device constants to the active device.
652 : ******************************************************************************/
653 97 : void torch_c_model_remap_device_constants(torch_c_model_t *model) {
654 97 : c10::OptionalDeviceGuard guard;
655 97 : const auto device = get_device_with_guard(guard);
656 97 : if (device.is_cuda()) {
657 0 : remap_model_device_constants(*model, device);
658 : }
659 97 : }
660 :
661 : /*******************************************************************************
662 : * \brief Disables gradients for inference-only model parameters.
663 : ******************************************************************************/
664 97 : void torch_c_model_disable_parameter_gradients(torch_c_model_t *model) {
665 97 : torch::NoGradGuard no_grad;
666 7857 : for (auto parameter : model->parameters()) {
667 15520 : parameter.set_requires_grad(false);
668 7857 : }
669 97 : }
670 :
671 : /*******************************************************************************
672 : * \brief Evaluates the given Torch model.
673 : * \author Ole Schuett
674 : ******************************************************************************/
675 62 : void torch_c_model_forward(torch_c_model_t *model, const torch_c_dict_t *inputs,
676 : torch_c_dict_t *outputs) {
677 :
678 62 : TorchFloatingPointMaskGuard fpe_guard;
679 62 : c10::OptionalDeviceGuard guard;
680 62 : get_device_with_guard(guard);
681 310 : auto untyped_output = model->forward({*inputs}).toGenericDict();
682 62 : outputs->clear();
683 232 : for (const auto &entry : untyped_output) {
684 170 : outputs->insert(entry.key().toStringView(), entry.value().toTensor());
685 : }
686 186 : }
687 :
688 : /*******************************************************************************
689 : * \brief Evaluates a TorchScript model method expecting argument "mol".
690 : ******************************************************************************/
691 386 : void torch_c_model_forward_mol_tensor(torch_c_model_t *model,
692 : const char *method_name,
693 : const torch_c_dict_t *inputs,
694 : torch_c_tensor_t **output) {
695 :
696 386 : c10::OptionalDeviceGuard guard;
697 386 : get_device_with_guard(guard);
698 386 : assert(*output == NULL);
699 386 : *output = new torch_c_tensor_t(
700 1930 : model->get_method(method_name)({*inputs}).toTensor());
701 1158 : }
702 :
703 : /*******************************************************************************
704 : * \brief Returns the weighted sum of two Torch tensors.
705 : ******************************************************************************/
706 386 : void torch_c_tensor_weighted_sum(const torch_c_tensor_t *values,
707 : const torch_c_tensor_t *weights,
708 : torch_c_tensor_t **result) {
709 386 : c10::OptionalDeviceGuard guard;
710 386 : get_device_with_guard(guard);
711 386 : const auto weights_on_device = weights->to(values->device());
712 772 : *result = new torch_c_tensor_t((*values * weights_on_device).sum());
713 386 : }
714 :
715 : /*******************************************************************************
716 : * \brief Returns a scalar double value from a Torch tensor.
717 : ******************************************************************************/
718 386 : double torch_c_tensor_item_double(const torch_c_tensor_t *tensor) {
719 386 : c10::OptionalDeviceGuard guard;
720 386 : get_device_with_guard(guard);
721 386 : return tensor->item<double>();
722 386 : }
723 :
724 : /*******************************************************************************
725 : * \brief Releases a Torch model and all its ressources.
726 : * \author Ole Schuett
727 : ******************************************************************************/
728 16 : void torch_c_model_release(torch_c_model_t *model) { delete (model); }
729 :
730 : /*******************************************************************************
731 : * \brief Reads metadata entry from given "*.pth" file.
732 : * In Torch lingo they are called extra files.
733 : * The returned char array has to be deallocated by caller!
734 : * \author Ole Schuett
735 : ******************************************************************************/
736 40 : void torch_c_model_read_metadata(const char *filename, const char *key,
737 : char **content, int *length) {
738 :
739 120 : std::unordered_map<std::string, std::string> extra_files = {{key, ""}};
740 40 : torch::jit::load(filename, torch::kCPU, extra_files);
741 80 : const std::string &content_str = extra_files[key];
742 40 : copy_string_to_c_buffer(content_str, content, length);
743 80 : }
744 :
745 : /*******************************************************************************
746 : * \brief Returns true iff the Torch CUDA backend is available.
747 : * \author Ole Schuett
748 : ******************************************************************************/
749 2 : bool torch_c_cuda_is_available() { return torch::cuda::is_available(); }
750 :
751 : /*******************************************************************************
752 : * \brief Return the number of CUDA devices visible to Torch.
753 : ******************************************************************************/
754 0 : int torch_c_cuda_device_count() {
755 0 : return torch::cuda::is_available() ? torch::cuda::device_count() : 0;
756 : }
757 :
758 : /*******************************************************************************
759 : * \brief Set whether to allow TF32.
760 : * Needed due to changes in defaults from pytorch 1.7 to 1.11 to >=1.12
761 : * See https://pytorch.org/docs/stable/notes/cuda.html
762 : * \author Gabriele Tocci
763 : ******************************************************************************/
764 6 : void torch_c_allow_tf32(const bool allow_tf32) {
765 :
766 6 : at::globalContext().setAllowTF32CuBLAS(allow_tf32);
767 6 : at::globalContext().setAllowTF32CuDNN(allow_tf32);
768 6 : }
769 :
770 : /******************************************************************************
771 : * \brief Freeze the Torch model: generic optimization that speeds up model.
772 : * See https://pytorch.org/docs/stable/generated/torch.jit.freeze.html
773 : * \author Gabriele Tocci
774 : ******************************************************************************/
775 6 : void torch_c_model_freeze(torch_c_model_t *model) {
776 :
777 6 : *model = torch::jit::freeze(*model);
778 6 : }
779 :
780 : /*******************************************************************************
781 : * \brief Retrieves an int64 attribute. Must be called before model freeze.
782 : * \author Ole Schuett
783 : ******************************************************************************/
784 40 : void torch_c_model_get_attr_int64(const torch_c_model_t *model, const char *key,
785 : int64_t *dest) {
786 40 : *dest = model->attr(key).toInt();
787 40 : }
788 :
789 : /*******************************************************************************
790 : * \brief Retrieves a double attribute. Must be called before model freeze.
791 : * \author Ole Schuett
792 : ******************************************************************************/
793 8 : void torch_c_model_get_attr_double(const torch_c_model_t *model,
794 : const char *key, double *dest) {
795 8 : *dest = model->attr(key).toDouble();
796 8 : }
797 :
798 : /*******************************************************************************
799 : * \brief Retrieves a string attribute. Must be called before model freeze.
800 : * \author Ole Schuett
801 : ******************************************************************************/
802 16 : void torch_c_model_get_attr_string(const torch_c_model_t *model,
803 : const char *key, char *dest) {
804 16 : const std::string &str = model->attr(key).toStringRef();
805 16 : assert(str.size() < 80); // default_string_length
806 144 : for (int i = 0; i < str.size(); i++) {
807 128 : dest[i] = str[i];
808 : }
809 16 : }
810 :
811 : /*******************************************************************************
812 : * \brief Retrieves a list attribute's size. Must be called before model freeze.
813 : * \author Ole Schuett
814 : ******************************************************************************/
815 8 : void torch_c_model_get_attr_list_size(const torch_c_model_t *model,
816 : const char *key, int *size) {
817 8 : *size = model->attr(key).toList().size();
818 8 : }
819 :
820 : /*******************************************************************************
821 : * \brief Retrieves a single item from a string list attribute.
822 : * \author Ole Schuett
823 : ******************************************************************************/
824 16 : void torch_c_model_get_attr_strlist(const torch_c_model_t *model,
825 : const char *key, const int index,
826 : char *dest) {
827 32 : const auto list = model->attr(key).toList();
828 16 : const std::string &str = list[index].toStringRef();
829 16 : assert(str.size() < 80); // default_string_length
830 32 : for (int i = 0; i < str.size(); i++) {
831 16 : dest[i] = str[i];
832 : }
833 16 : }
834 :
835 : #ifdef __cplusplus
836 : }
837 : #endif
838 :
839 : #endif // defined(__LIBTORCH)
840 :
841 : // EOF
|
