TRIQS/nda 2.0.0
Multi-dimensional array library for C++
Loading...
Searching...
No Matches
cutensor_interface.cpp
Go to the documentation of this file.
1// Copyright (c) 2024--present, The Simons Foundation
2// This file is part of TRIQS/nda and is licensed under the Apache License, Version 2.0.
3// SPDX-License-Identifier: Apache-2.0
4// See LICENSE in the root of this distribution for details.
5
10
12#include "../tools.hpp"
13#include "../../concepts.hpp"
14#include "../../device.hpp"
15#include "../../exceptions.hpp"
16
17#include <cutensor.h>
18
19#include <algorithm>
20#include <bit>
21#include <complex>
22#include <cstdint>
23#include <string_view>
24#include <type_traits>
25#include <vector>
26
27namespace nda::tensor::device {
28
29 // File-local synchronization flag, exposed via the get/set functions below (matches the BLAS interface pattern).
30 thread_local bool synchronize = true; // NOLINT (per-thread option is on purpose)
31 void set_synchronization(bool do_sync) noexcept { synchronize = do_sync; }
32 bool get_synchronization() noexcept { return synchronize; }
33
34 // Get the cutensor handle.
35 cutensorHandle_t &get_handle() {
36 struct handle_storage_t { // RAII for handle
37 handle_storage_t() { cutensorCreate(&handle); }
38 ~handle_storage_t() { cutensorDestroy(handle); }
39 cutensorHandle_t handle = {};
40 };
41 static auto sto = handle_storage_t{};
42 return sto.handle;
43 }
44
45 // Anonymous namespace for local functions.
46 namespace {
47
48 // Check the success of a cutensor operation.
49 void cutensor_error_check(cutensorStatus_t status, std::string_view func) {
50 if (status != CUTENSOR_STATUS_SUCCESS) {
51 NDA_RUNTIME_ERROR << "cuTENSOR runtime error in function " << func << "\n"
52 << " cutensorStatus_t: " << status << "\n"
53 << " cutensorGetErrorString: " << cutensorGetErrorString(status) << "\n";
54 }
55 }
56
57 // Cuda data type conversion.
58 template <typename T, typename U = std::remove_const_t<T>>
59 constexpr auto cuda_data_type() {
60 if constexpr (std::is_same_v<U, float>) {
61 return CUTENSOR_R_32F;
62 } else if constexpr (std::is_same_v<U, double>) {
63 return CUTENSOR_R_64F;
64 } else if constexpr (std::is_same_v<U, std::complex<float>>) {
65 return CUTENSOR_C_32F;
66 } else if constexpr (std::is_same_v<U, std::complex<double>>) {
67 return CUTENSOR_C_64F;
68 }
69 }
70
71 // Cutensor compute type conversion.
72 template <typename T, typename U = std::remove_const_t<T>>
73 constexpr auto cutensor_compute_type() {
74 if constexpr (AnyOf<U, float, std::complex<float>>) {
75 return CUTENSOR_COMPUTE_DESC_32F;
76 } else if constexpr (AnyOf<U, double, std::complex<double>>) {
77 return CUTENSOR_COMPUTE_DESC_64F;
78 }
79 }
80
81 // Largest power-of-two divisor of the pointer address, capped at 256 (cudaMalloc default alignment).
82 // Capping the shift count also covers p == nullptr, where countr_zero returns the full width.
83 template <typename T>
84 std::uint32_t find_alignment(T *p) {
85 auto const x = reinterpret_cast<std::uintptr_t>(p); // NOLINT (reinterpret_cast is necessary here)
86 return std::uint32_t{1} << std::min(std::countr_zero(x), 8);
87 }
88
89 // Convert an index string to a vector of int32_t mode labels for cuTENSOR.
90 auto to_modes(std::string_view idx) { return std::vector<std::int32_t>(idx.begin(), idx.end()); }
91
92 // Map unary_op enum to cuTENSOR unary operator.
93 // clang-format off
94 cutensorOperator_t to_cutensor_unary_op(unary_op op) {
95 switch (op) {
96 case unary_op::IDENTITY: return CUTENSOR_OP_IDENTITY;
97 case unary_op::SQRT: return CUTENSOR_OP_SQRT;
98 case unary_op::RELU: return CUTENSOR_OP_RELU;
99 case unary_op::CONJ: return CUTENSOR_OP_CONJ;
100 case unary_op::RCP: return CUTENSOR_OP_RCP;
101 case unary_op::SIGMOID: return CUTENSOR_OP_SIGMOID;
102 case unary_op::TANH: return CUTENSOR_OP_TANH;
103 case unary_op::EXP: return CUTENSOR_OP_EXP;
104 case unary_op::LOG: return CUTENSOR_OP_LOG;
105 case unary_op::ABS: return CUTENSOR_OP_ABS;
106 case unary_op::NEG: return CUTENSOR_OP_NEG;
107 case unary_op::SIN: return CUTENSOR_OP_SIN;
108 case unary_op::COS: return CUTENSOR_OP_COS;
109 case unary_op::TAN: return CUTENSOR_OP_TAN;
110 case unary_op::SINH: return CUTENSOR_OP_SINH;
111 case unary_op::COSH: return CUTENSOR_OP_COSH;
112 case unary_op::ASIN: return CUTENSOR_OP_ASIN;
113 case unary_op::ACOS: return CUTENSOR_OP_ACOS;
114 case unary_op::ATAN: return CUTENSOR_OP_ATAN;
115 case unary_op::ASINH: return CUTENSOR_OP_ASINH;
116 case unary_op::ACOSH: return CUTENSOR_OP_ACOSH;
117 case unary_op::ATANH: return CUTENSOR_OP_ATANH;
118 case unary_op::CEIL: return CUTENSOR_OP_CEIL;
119 case unary_op::FLOOR: return CUTENSOR_OP_FLOOR;
120 case unary_op::MISH: return CUTENSOR_OP_MISH;
121 case unary_op::SWISH: return CUTENSOR_OP_SWISH;
122 case unary_op::SOFT_PLUS: return CUTENSOR_OP_SOFT_PLUS;
123 case unary_op::SOFT_SIGN: return CUTENSOR_OP_SOFT_SIGN;
124 default: NDA_RUNTIME_ERROR << "nda::tensor::cutensor: unary_op has no cuTENSOR equivalent";
125 }
126 }
127 // clang-format on
128
129 // Map binary_op enum to cuTENSOR binary operator.
130 cutensorOperator_t to_cutensor_binary_op(binary_op op) {
131 switch (op) {
132 case binary_op::SUM: return CUTENSOR_OP_ADD;
133 case binary_op::PROD: return CUTENSOR_OP_MUL;
134 case binary_op::MAX: return CUTENSOR_OP_MAX;
135 case binary_op::MIN: return CUTENSOR_OP_MIN;
136 default: NDA_RUNTIME_ERROR << "nda::tensor::cutensor: binary_op has no cuTENSOR equivalent";
137 }
138 }
139
140 // Create a tensor descriptor from a tensor view.
141 template <typename T>
142 auto create_tensor_desc(tensor_view<T> tv) {
143 cutensorTensorDescriptor_t desc{};
144 auto status = cutensorCreateTensorDescriptor(get_handle(), &desc, static_cast<std::uint32_t>(tv.ndim), tv.extents, tv.strides,
145 cuda_data_type<T>(), find_alignment(tv.data));
146 cutensor_error_check(status, "cutensorCreateTensorDescriptor");
147 return desc;
148 }
149
150 // Destroy a given tensor descriptor.
151 void destroy_tensor_desc(cutensorTensorDescriptor_t desc) {
152 cutensor_error_check(cutensorDestroyTensorDescriptor(desc), "cutensorDestroyTensorDescriptor");
153 }
154
155 // Create a plan preference with default algorithm and no JIT.
156 cutensorPlanPreference_t create_plan_pref() {
157 cutensorPlanPreference_t pref{};
158 cutensor_error_check(cutensorCreatePlanPreference(get_handle(), &pref, CUTENSOR_ALGO_DEFAULT, CUTENSOR_JIT_MODE_NONE),
159 "cutensorCreatePlanPreference");
160 return pref;
161 }
162
163 // Destroy a given plan preference.
164 void destroy_plan_pref(cutensorPlanPreference_t pref) {
165 cutensor_error_check(cutensorDestroyPlanPreference(pref), "cutensorDestroyPlanPreference");
166 }
167
168 // Create an execution plan from an operation descriptor, plan preference and workspace size limit.
169 cutensorPlan_t create_plan(cutensorOperationDescriptor_t op_desc, cutensorPlanPreference_t pref, std::uint64_t workspace_limit = 0) {
170 cutensorPlan_t plan{};
171 cutensor_error_check(cutensorCreatePlan(get_handle(), &plan, op_desc, pref, workspace_limit), "cutensorCreatePlan");
172 return plan;
173 }
174
175 // Destroy a given execution plan.
176 void destroy_plan(cutensorPlan_t plan) { cutensor_error_check(cutensorDestroyPlan(plan), "cutensorDestroyPlan"); }
177
178 // Destroy a given operation descriptor.
179 void destroy_op_desc(cutensorOperationDescriptor_t op_desc) {
180 cutensor_error_check(cutensorDestroyOperationDescriptor(op_desc), "cutensorDestroyOperationDescriptor");
181 }
182
183 // Estimate the workspace size for an operation.
184 std::uint64_t estimate_workspace(cutensorOperationDescriptor_t op_desc, cutensorPlanPreference_t pref,
185 cutensorWorksizePreference_t ws_pref = CUTENSOR_WORKSPACE_DEFAULT) {
186 std::uint64_t size = 0;
187 cutensor_error_check(cutensorEstimateWorkspaceSize(get_handle(), op_desc, pref, ws_pref, &size), "cutensorEstimateWorkspaceSize");
188 return size;
189 }
190
191 // Helper function to call the cuTENSOR permutation routine: B = alpha * opA(A).
192 template <typename T>
193 void permute_impl(T alpha, const_tensor_view<T> A, std::string_view idx_A, tensor_view<T> B, std::string_view idx_B) {
194 auto &handle = get_handle();
195
196 // create tensor descriptors
197 auto desc_A = create_tensor_desc(A);
198 auto desc_B = create_tensor_desc(B);
199
200 // convert index strings to mode arrays
201 auto modes_A = to_modes(idx_A);
202 auto modes_B = to_modes(idx_B);
203
204 // create operation descriptor
205 cutensorOperationDescriptor_t op_desc{};
206 cutensor_error_check(cutensorCreatePermutation(handle, &op_desc, desc_A, modes_A.data(), to_cutensor_unary_op(A.op), desc_B, modes_B.data(),
207 cutensor_compute_type<T>()),
208 "cutensorCreatePermutation");
209
210 // create plan preference, estimate workspace, and create plan
211 auto pref = create_plan_pref();
212 auto ws_limit = estimate_workspace(op_desc, pref);
213 auto plan = create_plan(op_desc, pref, ws_limit);
214
215 // execute permutation
216 cutensor_error_check(cutensorPermute(handle, plan, cuscalar(alpha), A.data, B.data, nullptr /*stream*/), "cutensorPermute");
217
218 // synchronize
219 cuda_device_sync(synchronize, "cutensorPermute");
220
221 // cleanup
222 destroy_plan(plan);
223 destroy_plan_pref(pref);
224 destroy_op_desc(op_desc);
225 destroy_tensor_desc(desc_A);
226 destroy_tensor_desc(desc_B);
227 }
228
229 // Helper function to call the cuTENSOR elementwise binary routine: D = op_AC(alpha * op_A(A), gamma * op_C(C)).
230 // D must have the same descriptor (shape/strides) as C but may point to different memory.
231 template <typename T>
232 void elementwise_binary_impl(T alpha, const_tensor_view<T> A, std::string_view idx_A, T gamma, const_tensor_view<T> C, std::string_view idx_C,
233 tensor_view<T> D, binary_op op_AC) {
234 auto &handle = get_handle();
235
236 // create tensor descriptors (D descriptor must match C in shape/modes)
237 auto desc_A = create_tensor_desc(A);
238 auto desc_C = create_tensor_desc(C);
239 auto desc_D = create_tensor_desc(D);
240
241 // convert index strings to mode arrays
242 auto modes_A = to_modes(idx_A);
243 auto modes_C = to_modes(idx_C);
244
245 // create operation descriptor (D has same modes as C)
246 cutensorOperationDescriptor_t op_desc{};
247 cutensor_error_check(cutensorCreateElementwiseBinary(handle, &op_desc, desc_A, modes_A.data(), to_cutensor_unary_op(A.op), desc_C,
248 modes_C.data(), to_cutensor_unary_op(C.op), desc_D, modes_C.data(),
249 to_cutensor_binary_op(op_AC), cutensor_compute_type<T>()),
250 "cutensorCreateElementwiseBinary");
251
252 // create plan preference, estimate workspace, and create plan
253 auto pref = create_plan_pref();
254 auto ws_limit = estimate_workspace(op_desc, pref);
255 auto plan = create_plan(op_desc, pref, ws_limit);
256
257 // execute elementwise binary
258 cutensor_error_check(
259 cutensorElementwiseBinaryExecute(handle, plan, cuscalar(alpha), A.data, cuscalar(gamma), C.data, D.data, nullptr /*stream*/),
260 "cutensorElementwiseBinaryExecute");
261
262 // synchronize
263 cuda_device_sync(synchronize, "cutensorElementwiseBinaryExecute");
264
265 // cleanup
266 destroy_plan(plan);
267 destroy_plan_pref(pref);
268 destroy_op_desc(op_desc);
269 destroy_tensor_desc(desc_A);
270 destroy_tensor_desc(desc_C);
271 destroy_tensor_desc(desc_D);
272 }
273
274 // Helper function to call the cuTENSOR elementwise trinary routine: D = op_ABC(op_AB(alpha * op_A(A), beta * op_B(B)), gamma * op_C(C)).
275 // D must have the same descriptor (shape/strides) as C but may point to different memory.
276 template <typename T>
277 void elementwise_trinary_impl(T alpha, const_tensor_view<T> A, std::string_view idx_A, T beta, const_tensor_view<T> B, std::string_view idx_B,
278 T gamma, const_tensor_view<T> C, std::string_view idx_C, tensor_view<T> D, binary_op op_AB, binary_op op_ABC) {
279 auto &handle = get_handle();
280
281 // create tensor descriptors (D descriptor must match C in shape/modes)
282 auto desc_A = create_tensor_desc(A);
283 auto desc_B = create_tensor_desc(B);
284 auto desc_C = create_tensor_desc(C);
285 auto desc_D = create_tensor_desc(D);
286
287 // convert index strings to mode arrays
288 auto modes_A = to_modes(idx_A);
289 auto modes_B = to_modes(idx_B);
290 auto modes_C = to_modes(idx_C);
291
292 // create operation descriptor (D has same modes as C)
293 cutensorOperationDescriptor_t op_desc{};
294 cutensor_error_check(cutensorCreateElementwiseTrinary(handle, &op_desc, desc_A, modes_A.data(), to_cutensor_unary_op(A.op), desc_B,
295 modes_B.data(), to_cutensor_unary_op(B.op), desc_C, modes_C.data(),
296 to_cutensor_unary_op(C.op), desc_D, modes_C.data(), to_cutensor_binary_op(op_AB),
297 to_cutensor_binary_op(op_ABC), cutensor_compute_type<T>()),
298 "cutensorCreateElementwiseTrinary");
299
300 // create plan preference, estimate workspace, and create plan
301 auto pref = create_plan_pref();
302 auto ws_limit = estimate_workspace(op_desc, pref);
303 auto plan = create_plan(op_desc, pref, ws_limit);
304
305 // execute elementwise trinary
306 cutensor_error_check(cutensorElementwiseTrinaryExecute(handle, plan, cuscalar(alpha), A.data, cuscalar(beta), B.data, cuscalar(gamma), C.data,
307 D.data, nullptr /*stream*/),
308 "cutensorElementwiseTrinaryExecute");
309
310 // synchronize
311 cuda_device_sync(synchronize, "cutensorElementwiseTrinaryExecute");
312
313 // cleanup
314 destroy_plan(plan);
315 destroy_plan_pref(pref);
316 destroy_op_desc(op_desc);
317 destroy_tensor_desc(desc_A);
318 destroy_tensor_desc(desc_B);
319 destroy_tensor_desc(desc_C);
320 destroy_tensor_desc(desc_D);
321 }
322
323 // Helper function to call the cuTENSOR reduction routine: D = alpha * opReduce(op_A(A)) + beta * op_C(C).
324 // D must have the same descriptor (shape/strides) as C but may point to different memory.
325 // The modes of C/D must be a subset of the modes of A. The modes in A but not in C are reduced.
326 template <typename T>
327 void reduce_impl(T alpha, const_tensor_view<T> A, std::string_view idx_A, T beta, const_tensor_view<T> C, std::string_view idx_C,
328 tensor_view<T> D, binary_op op_reduce) {
329 auto &handle = get_handle();
330
331 // create tensor descriptors (D descriptor must match C in shape/modes)
332 auto desc_A = create_tensor_desc(A);
333 auto desc_C = create_tensor_desc(C);
334 auto desc_D = create_tensor_desc(D);
335
336 // convert index strings to mode arrays
337 auto modes_A = to_modes(idx_A);
338 auto modes_C = to_modes(idx_C);
339
340 // create operation descriptor (D has same modes as C)
341 cutensorOperationDescriptor_t op_desc{};
342 cutensor_error_check(cutensorCreateReduction(handle, &op_desc, desc_A, modes_A.data(), to_cutensor_unary_op(A.op), desc_C, modes_C.data(),
343 to_cutensor_unary_op(C.op), desc_D, modes_C.data(), to_cutensor_binary_op(op_reduce),
344 cutensor_compute_type<T>()),
345 "cutensorCreateReduction");
346
347 // create plan preference, estimate workspace, and create plan
348 auto pref = create_plan_pref();
349 auto ws_limit = estimate_workspace(op_desc, pref);
350 auto plan = create_plan(op_desc, pref, ws_limit);
351
352 // query the actual required workspace size from the plan
353 std::uint64_t ws_size = 0;
354 cutensor_error_check(cutensorPlanGetAttribute(handle, plan, CUTENSOR_PLAN_REQUIRED_WORKSPACE, &ws_size, sizeof(ws_size)),
355 "cutensorPlanGetAttribute");
356
357 // allocate workspace
358 void *workspace = nullptr;
359 if (ws_size > 0) { device_error_check(cudaMalloc(&workspace, ws_size), "cudaMalloc"); }
360
361 // execute reduction
362 cutensor_error_check(
363 cutensorReduce(handle, plan, cuscalar(alpha), A.data, cuscalar(beta), C.data, D.data, workspace, ws_size, nullptr /*stream*/),
364 "cutensorReduce");
365
366 // synchronize
367 cuda_device_sync(synchronize, "cutensorReduce");
368
369 // free workspace
370 if (workspace) { device_error_check(cudaFree(workspace), "cudaFree"); }
371
372 // cleanup
373 destroy_plan(plan);
374 destroy_plan_pref(pref);
375 destroy_op_desc(op_desc);
376 destroy_tensor_desc(desc_A);
377 destroy_tensor_desc(desc_C);
378 destroy_tensor_desc(desc_D);
379 }
380
381 // Helper function to call the cuTENSOR contraction routine: D = alpha * op_A(A) * op_B(B) + beta * op_C(C).
382 // D must have the same descriptor (shape/strides) as C but may point to different memory.
383 template <typename T>
384 void contract_impl(T alpha, const_tensor_view<T> A, std::string_view idx_A, const_tensor_view<T> B, std::string_view idx_B, T beta,
385 const_tensor_view<T> C, std::string_view idx_C, tensor_view<T> D) {
386 auto &handle = get_handle();
387
388 // create tensor descriptors (D descriptor must match C in shape/modes)
389 auto desc_A = create_tensor_desc(A);
390 auto desc_B = create_tensor_desc(B);
391 auto desc_C = create_tensor_desc(C);
392 auto desc_D = create_tensor_desc(D);
393
394 // convert index strings to mode arrays
395 auto modes_A = to_modes(idx_A);
396 auto modes_B = to_modes(idx_B);
397 auto modes_C = to_modes(idx_C);
398
399 // create operation descriptor (D has same modes as C)
400 cutensorOperationDescriptor_t op_desc{};
401 cutensor_error_check(cutensorCreateContraction(handle, &op_desc, desc_A, modes_A.data(), to_cutensor_unary_op(A.op), desc_B, modes_B.data(),
402 to_cutensor_unary_op(B.op), desc_C, modes_C.data(), to_cutensor_unary_op(C.op), desc_D,
403 modes_C.data(), cutensor_compute_type<T>()),
404 "cutensorCreateContraction");
405
406 // create plan preference, estimate workspace, and create plan
407 auto pref = create_plan_pref();
408 auto ws_limit = estimate_workspace(op_desc, pref);
409 auto plan = create_plan(op_desc, pref, ws_limit);
410
411 // query the actual required workspace size from the plan
412 std::uint64_t ws_size = 0;
413 cutensor_error_check(cutensorPlanGetAttribute(handle, plan, CUTENSOR_PLAN_REQUIRED_WORKSPACE, &ws_size, sizeof(ws_size)),
414 "cutensorPlanGetAttribute");
415
416 // allocate workspace
417 void *workspace = nullptr;
418 if (ws_size > 0) { device_error_check(cudaMalloc(&workspace, ws_size), "cudaMalloc"); }
419
420 // execute contraction
421 cutensor_error_check(
422 cutensorContract(handle, plan, cuscalar(alpha), A.data, B.data, cuscalar(beta), C.data, D.data, workspace, ws_size, nullptr /*stream*/),
423 "cutensorContract");
424
425 // synchronize
426 cuda_device_sync(synchronize, "cutensorContract");
427
428 // free workspace
429 if (workspace) { device_error_check(cudaFree(workspace), "cudaFree"); }
430
431 // cleanup
432 destroy_plan(plan);
433 destroy_plan_pref(pref);
434 destroy_op_desc(op_desc);
435 destroy_tensor_desc(desc_A);
436 destroy_tensor_desc(desc_B);
437 destroy_tensor_desc(desc_C);
438 destroy_tensor_desc(desc_D);
439 }
440
441 } // namespace
442
443 // permute
444 void permute(float alpha, const_tensor_view<float> A, std::string_view idx_A, tensor_view<float> B, std::string_view idx_B) {
445 permute_impl(alpha, A, idx_A, B, idx_B);
446 }
447 void permute(double alpha, const_tensor_view<double> A, std::string_view idx_A, tensor_view<double> B, std::string_view idx_B) {
448 permute_impl(alpha, A, idx_A, B, idx_B);
449 }
450 void permute(std::complex<float> alpha, const_tensor_view<std::complex<float>> A, std::string_view idx_A, tensor_view<std::complex<float>> B,
451 std::string_view idx_B) {
452 permute_impl(alpha, A, idx_A, B, idx_B);
453 }
454 void permute(std::complex<double> alpha, const_tensor_view<std::complex<double>> A, std::string_view idx_A, tensor_view<std::complex<double>> B,
455 std::string_view idx_B) {
456 permute_impl(alpha, A, idx_A, B, idx_B);
457 }
458
459 // elementwise_binary
460 void elementwise_binary(float alpha, const_tensor_view<float> A, std::string_view idx_A, float gamma, const_tensor_view<float> C,
461 std::string_view idx_C, tensor_view<float> D, binary_op op_AC) {
462 elementwise_binary_impl(alpha, A, idx_A, gamma, C, idx_C, D, op_AC);
463 }
464 void elementwise_binary(double alpha, const_tensor_view<double> A, std::string_view idx_A, double gamma, const_tensor_view<double> C,
465 std::string_view idx_C, tensor_view<double> D, binary_op op_AC) {
466 elementwise_binary_impl(alpha, A, idx_A, gamma, C, idx_C, D, op_AC);
467 }
468 void elementwise_binary(std::complex<float> alpha, const_tensor_view<std::complex<float>> A, std::string_view idx_A, std::complex<float> gamma,
469 const_tensor_view<std::complex<float>> C, std::string_view idx_C, tensor_view<std::complex<float>> D, binary_op op_AC) {
470 elementwise_binary_impl(alpha, A, idx_A, gamma, C, idx_C, D, op_AC);
471 }
472 void elementwise_binary(std::complex<double> alpha, const_tensor_view<std::complex<double>> A, std::string_view idx_A, std::complex<double> gamma,
473 const_tensor_view<std::complex<double>> C, std::string_view idx_C, tensor_view<std::complex<double>> D, binary_op op_AC) {
474 elementwise_binary_impl(alpha, A, idx_A, gamma, C, idx_C, D, op_AC);
475 }
476
477 // elementwise_trinary
478 void elementwise_trinary(float alpha, const_tensor_view<float> A, std::string_view idx_A, float beta, const_tensor_view<float> B,
479 std::string_view idx_B, float gamma, const_tensor_view<float> C, std::string_view idx_C, tensor_view<float> D,
480 binary_op op_AB, binary_op op_ABC) {
481 elementwise_trinary_impl(alpha, A, idx_A, beta, B, idx_B, gamma, C, idx_C, D, op_AB, op_ABC);
482 }
483 void elementwise_trinary(double alpha, const_tensor_view<double> A, std::string_view idx_A, double beta, const_tensor_view<double> B,
484 std::string_view idx_B, double gamma, const_tensor_view<double> C, std::string_view idx_C, tensor_view<double> D,
485 binary_op op_AB, binary_op op_ABC) {
486 elementwise_trinary_impl(alpha, A, idx_A, beta, B, idx_B, gamma, C, idx_C, D, op_AB, op_ABC);
487 }
488 void elementwise_trinary(std::complex<float> alpha, const_tensor_view<std::complex<float>> A, std::string_view idx_A, std::complex<float> beta,
489 const_tensor_view<std::complex<float>> B, std::string_view idx_B, std::complex<float> gamma,
490 const_tensor_view<std::complex<float>> C, std::string_view idx_C, tensor_view<std::complex<float>> D, binary_op op_AB,
491 binary_op op_ABC) {
492 elementwise_trinary_impl(alpha, A, idx_A, beta, B, idx_B, gamma, C, idx_C, D, op_AB, op_ABC);
493 }
494 void elementwise_trinary(std::complex<double> alpha, const_tensor_view<std::complex<double>> A, std::string_view idx_A, std::complex<double> beta,
495 const_tensor_view<std::complex<double>> B, std::string_view idx_B, std::complex<double> gamma,
496 const_tensor_view<std::complex<double>> C, std::string_view idx_C, tensor_view<std::complex<double>> D, binary_op op_AB,
497 binary_op op_ABC) {
498 elementwise_trinary_impl(alpha, A, idx_A, beta, B, idx_B, gamma, C, idx_C, D, op_AB, op_ABC);
499 }
500
501 // reduce
502 void reduce(float alpha, const_tensor_view<float> A, std::string_view idx_A, float beta, const_tensor_view<float> C, std::string_view idx_C,
503 tensor_view<float> D, binary_op op_reduce) {
504 reduce_impl(alpha, A, idx_A, beta, C, idx_C, D, op_reduce);
505 }
506 void reduce(double alpha, const_tensor_view<double> A, std::string_view idx_A, double beta, const_tensor_view<double> C, std::string_view idx_C,
507 tensor_view<double> D, binary_op op_reduce) {
508 reduce_impl(alpha, A, idx_A, beta, C, idx_C, D, op_reduce);
509 }
510 void reduce(std::complex<float> alpha, const_tensor_view<std::complex<float>> A, std::string_view idx_A, std::complex<float> beta,
511 const_tensor_view<std::complex<float>> C, std::string_view idx_C, tensor_view<std::complex<float>> D, binary_op op_reduce) {
512 reduce_impl(alpha, A, idx_A, beta, C, idx_C, D, op_reduce);
513 }
514 void reduce(std::complex<double> alpha, const_tensor_view<std::complex<double>> A, std::string_view idx_A, std::complex<double> beta,
515 const_tensor_view<std::complex<double>> C, std::string_view idx_C, tensor_view<std::complex<double>> D, binary_op op_reduce) {
516 reduce_impl(alpha, A, idx_A, beta, C, idx_C, D, op_reduce);
517 }
518
519 // contract
520 void contract(float alpha, const_tensor_view<float> A, std::string_view idx_A, const_tensor_view<float> B, std::string_view idx_B, float beta,
521 const_tensor_view<float> C, std::string_view idx_C, tensor_view<float> D) {
522 contract_impl(alpha, A, idx_A, B, idx_B, beta, C, idx_C, D);
523 }
524 void contract(double alpha, const_tensor_view<double> A, std::string_view idx_A, const_tensor_view<double> B, std::string_view idx_B, double beta,
525 const_tensor_view<double> C, std::string_view idx_C, tensor_view<double> D) {
526 contract_impl(alpha, A, idx_A, B, idx_B, beta, C, idx_C, D);
527 }
528 void contract(std::complex<float> alpha, const_tensor_view<std::complex<float>> A, std::string_view idx_A, const_tensor_view<std::complex<float>> B,
529 std::string_view idx_B, std::complex<float> beta, const_tensor_view<std::complex<float>> C, std::string_view idx_C,
530 tensor_view<std::complex<float>> D) {
531 contract_impl(alpha, A, idx_A, B, idx_B, beta, C, idx_C, D);
532 }
533 void contract(std::complex<double> alpha, const_tensor_view<std::complex<double>> A, std::string_view idx_A,
534 const_tensor_view<std::complex<double>> B, std::string_view idx_B, std::complex<double> beta,
535 const_tensor_view<std::complex<double>> C, std::string_view idx_C, tensor_view<std::complex<double>> D) {
536 contract_impl(alpha, A, idx_A, B, idx_B, beta, C, idx_C, D);
537 }
538
539} // namespace nda::tensor::device
Provides concepts for the nda library.
Provides a C++ interface for various cuTENSOR routines.
Provides GPU and non-GPU specific functionality.
Provides a custom runtime error class and macros to assert conditions and throw exceptions.
#define device_error_check(ARG1, ARG2)
Trigger a compilation error every time the nda::device_error_check function is called.
Definition device.hpp:237
void cuda_device_sync(bool do_sync=true, std::string_view func="")
Empty function if CudaSupport is not enabled.
Definition device.hpp:246
unary_op
Unary element-wise operations for tensor operations.
Definition tools.hpp:103
binary_op
Binary operations for tensor operations.
Definition tools.hpp:67
tensor_view< const T > const_tensor_view
Alias for a tensor_view with const value type.
Definition tools.hpp:234
Provides various traits and utilities for the tensor interface.