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