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