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