示例¶
在本节中，我们将展示如何使用 cuTensorNet 收缩张量网络。首先，我们描述如何编译示例代码。然后，我们展示一个示例代码，用于执行 cuTensorNet 中的常用步骤。在该示例中，我们执行以下张量收缩
\[R_{k,l} = A_{a,b,c,d,e,f} B_{b,g,h,e,i,j} C_{m,a,g,f,i,k} D_{l,c,h,d,j,m}\]
我们逐步构建代码，每个步骤都在末尾添加代码。这些步骤通过简洁的多行注释块分隔。
建议读者参考 概述 和 cuTENSOR 文档，以熟悉术语和 cuTENSOR 操作。
编译代码¶
假设 cuQuantum 已在 CUQUANTUM_ROOT 中解压，cuTENSOR 已在 CUTENSOR_ROOT 中解压，我们按如下方式更新库路径
export LD_LIBRARY_PATH=${CUQUANTUM_ROOT}/lib:${CUTENSOR_ROOT}/lib/11:${LD_LIBRARY_PATH}
根据您的 CUDA 工具包，您可能需要选择不同的库版本（例如，${CUTENSOR_ROOT}/lib/11.0）。
下面讨论的串行示例代码 (tensornet_example.cu) 可以通过以下命令编译
nvcc tensornet_example.cu -I${CUQUANTUM_ROOT}/include -I${CUTENSOR_ROOT}/include -L${CUQUANTUM_ROOT}/lib -L${CUTENSOR_ROOT}/lib/11 -lcutensornet -lcutensor -o tensornet_example
对于与 cuTensorNet 库的静态链接，请使用以下命令（请注意，libmetis_static.a 需要显式链接，假设它通过 NVIDIA CUDA 工具包安装并且可以通过 $LIBRARY_PATH 访问）
nvcc tensornet_example.cu -I${CUQUANTUM_ROOT}/include -I${CUTENSOR_ROOT}/include ${CUQUANTUM_ROOT}/lib/libcutensornet_static.a -L${CUTENSOR_DIR}/lib/11 -lcutensor libmetis_static.a -o tensornet_example
为了构建示例的并行 (MPI) 版本（tensornet_example_mpi_auto.cu 和 tensornet_example_mpi.cu），您将需要安装 MPI 库（例如，最新的 Open MPI、MVAPICH 或 MPICH）。特别是，自动并行示例需要CUDA-aware MPI，请参阅下面的 代码示例（自动切片分布式并行化）。在这种情况下，您将需要将 -I${MPI_PATH}/include 和 -L${MPI_PATH}/lib -lmpi 添加到构建命令
nvcc tensornet_example_mpi_auto.cu -I${CUQUANTUM_ROOT}/include -I${CUTENSOR_ROOT}/include -I${MPI_PATH}/include -L${CUQUANTUM_ROOT}/lib -L${CUTENSOR_ROOT}/lib/11 -lcutensornet -lcutensor -L${MPI_PATH}/lib -lmpi -o tensornet_example_mpi_auto
nvcc tensornet_example_mpi.cu -I${CUQUANTUM_ROOT}/include -I${CUTENSOR_ROOT}/include -I${MPI_PATH}/include -L${CUQUANTUM_ROOT}/lib -L${CUTENSOR_ROOT}/lib/11 -lcutensornet -lcutensor -L${MPI_PATH}/lib -lmpi -o tensornet_example_mpi
警告
当在没有 CUDA-aware MPI 的情况下运行 tensornet_example_mpi_auto.cu 时，程序将崩溃。
注意
根据 cuQuantum 包的来源，您可能需要将上面的 lib 替换为 lib64。
代码示例（串行）¶
以下代码示例说明了使用 cuTensorNet 所需的常用步骤，并介绍了典型的张量网络操作。完整的示例代码可以在 NVIDIA/cuQuantum 存储库中找到（此处）。
头文件和数据类型¶
8#include <stdlib.h>
9#include <stdio.h>
10
11#include <unordered_map>
12#include <vector>
13#include <cassert>
14
15#include <cuda_runtime.h>
16#include <cutensornet.h>
17
18
19#define HANDLE_ERROR(x) \
20{ const auto err = x; \
21 if( err != CUTENSORNET_STATUS_SUCCESS ) \
22 { printf("Error: %s in line %d\n", cutensornetGetErrorString(err), __LINE__); \
23 fflush(stdout); \
24 } \
25};
26
27#define HANDLE_CUDA_ERROR(x) \
28{ const auto err = x; \
29 if( err != cudaSuccess ) \
30 { printf("CUDA Error: %s in line %d\n", cudaGetErrorString(err), __LINE__); \
31 fflush(stdout); \
32 } \
33};
34
35
36struct GPUTimer
37{
38 GPUTimer(cudaStream_t stream): stream_(stream)
39 {
40 HANDLE_CUDA_ERROR(cudaEventCreate(&start_));
41 HANDLE_CUDA_ERROR(cudaEventCreate(&stop_));
42 }
43
44 ~GPUTimer()
45 {
46 HANDLE_CUDA_ERROR(cudaEventDestroy(start_));
47 HANDLE_CUDA_ERROR(cudaEventDestroy(stop_));
48 }
49
50 void start()
51 {
52 HANDLE_CUDA_ERROR(cudaEventRecord(start_, stream_));
53 }
54
55 float seconds()
56 {
57 HANDLE_CUDA_ERROR(cudaEventRecord(stop_, stream_));
58 HANDLE_CUDA_ERROR(cudaEventSynchronize(stop_));
59 float time;
60 HANDLE_CUDA_ERROR(cudaEventElapsedTime(&time, start_, stop_));
61 return time * 1e-3;
62 }
63
64 private:
65 cudaEvent_t start_, stop_;
66 cudaStream_t stream_;
67};
68
69
70int main()
71{
72 static_assert(sizeof(size_t) == sizeof(int64_t), "Please build this sample on a 64-bit architecture!");
73
74 bool verbose = true;
75
76 // Check cuTensorNet version
77 const size_t cuTensornetVersion = cutensornetGetVersion();
78 if(verbose)
79 printf("cuTensorNet version: %ld\n", cuTensornetVersion);
80
81 // Set GPU device
82 int numDevices {0};
83 HANDLE_CUDA_ERROR( cudaGetDeviceCount(&numDevices) );
84 const int deviceId = 0;
85 HANDLE_CUDA_ERROR( cudaSetDevice(deviceId) );
86 cudaDeviceProp prop;
87 HANDLE_CUDA_ERROR( cudaGetDeviceProperties(&prop, deviceId) );
88
89 if(verbose) {
90 printf("===== device info ======\n");
91 printf("GPU-local-id:%d\n", deviceId);
92 printf("GPU-name:%s\n", prop.name);
93 printf("GPU-clock:%d\n", prop.clockRate);
94 printf("GPU-memoryClock:%d\n", prop.memoryClockRate);
95 printf("GPU-nSM:%d\n", prop.multiProcessorCount);
96 printf("GPU-major:%d\n", prop.major);
97 printf("GPU-minor:%d\n", prop.minor);
98 printf("========================\n");
99 }
100
101 typedef float floatType;
102 cudaDataType_t typeData = CUDA_R_32F;
103 cutensornetComputeType_t typeCompute = CUTENSORNET_COMPUTE_32F;
104
105 if(verbose)
106 printf("Included headers and defined data types\n");
定义张量网络和张量大小¶
接下来，我们定义张量网络的拓扑结构（即，张量的模式、它们的范围以及它们的连通性）。
110 /**********************
111 * Computing: R_{k,l} = A_{a,b,c,d,e,f} B_{b,g,h,e,i,j} C_{m,a,g,f,i,k} D_{l,c,h,d,j,m}
112 **********************/
113
114 constexpr int32_t numInputs = 4;
115
116 // Create vectors of tensor modes
117 std::vector<int32_t> modesA{'a','b','c','d','e','f'};
118 std::vector<int32_t> modesB{'b','g','h','e','i','j'};
119 std::vector<int32_t> modesC{'m','a','g','f','i','k'};
120 std::vector<int32_t> modesD{'l','c','h','d','j','m'};
121 std::vector<int32_t> modesR{'k','l'};
122
123 // Set mode extents
124 std::unordered_map<int32_t, int64_t> extent;
125 extent['a'] = 16;
126 extent['b'] = 16;
127 extent['c'] = 16;
128 extent['d'] = 16;
129 extent['e'] = 16;
130 extent['f'] = 16;
131 extent['g'] = 16;
132 extent['h'] = 16;
133 extent['i'] = 16;
134 extent['j'] = 16;
135 extent['k'] = 16;
136 extent['l'] = 16;
137 extent['m'] = 16;
138
139 // Create a vector of extents for each tensor
140 std::vector<int64_t> extentA;
141 for (auto mode : modesA)
142 extentA.push_back(extent[mode]);
143 std::vector<int64_t> extentB;
144 for (auto mode : modesB)
145 extentB.push_back(extent[mode]);
146 std::vector<int64_t> extentC;
147 for (auto mode : modesC)
148 extentC.push_back(extent[mode]);
149 std::vector<int64_t> extentD;
150 for (auto mode : modesD)
151 extentD.push_back(extent[mode]);
152 std::vector<int64_t> extentR;
153 for (auto mode : modesR)
154 extentR.push_back(extent[mode]);
155
156 if(verbose)
157 printf("Defined tensor network, modes, and extents\n");
分配内存并初始化数据¶
接下来，我们为张量网络操作数分配内存并将它们初始化为随机值。
160 /**********************
161 * Allocating data
162 **********************/
163
164 size_t elementsA = 1;
165 for (auto mode : modesA)
166 elementsA *= extent[mode];
167 size_t elementsB = 1;
168 for (auto mode : modesB)
169 elementsB *= extent[mode];
170 size_t elementsC = 1;
171 for (auto mode : modesC)
172 elementsC *= extent[mode];
173 size_t elementsD = 1;
174 for (auto mode : modesD)
175 elementsD *= extent[mode];
176 size_t elementsR = 1;
177 for (auto mode : modesR)
178 elementsR *= extent[mode];
179
180 size_t sizeA = sizeof(floatType) * elementsA;
181 size_t sizeB = sizeof(floatType) * elementsB;
182 size_t sizeC = sizeof(floatType) * elementsC;
183 size_t sizeD = sizeof(floatType) * elementsD;
184 size_t sizeR = sizeof(floatType) * elementsR;
185 if(verbose)
186 printf("Total GPU memory used for tensor storage: %.2f GiB\n",
187 (sizeA + sizeB + sizeC + sizeD + sizeR) / 1024. /1024. / 1024);
188
189 void* rawDataIn_d[numInputs];
190 void* R_d;
191 HANDLE_CUDA_ERROR( cudaMalloc((void**) &rawDataIn_d[0], sizeA) );
192 HANDLE_CUDA_ERROR( cudaMalloc((void**) &rawDataIn_d[1], sizeB) );
193 HANDLE_CUDA_ERROR( cudaMalloc((void**) &rawDataIn_d[2], sizeC) );
194 HANDLE_CUDA_ERROR( cudaMalloc((void**) &rawDataIn_d[3], sizeD) );
195 HANDLE_CUDA_ERROR( cudaMalloc((void**) &R_d, sizeR));
196
197 floatType *A = (floatType*) malloc(sizeof(floatType) * elementsA);
198 floatType *B = (floatType*) malloc(sizeof(floatType) * elementsB);
199 floatType *C = (floatType*) malloc(sizeof(floatType) * elementsC);
200 floatType *D = (floatType*) malloc(sizeof(floatType) * elementsD);
201 floatType *R = (floatType*) malloc(sizeof(floatType) * elementsR);
202
203 if (A == NULL || B == NULL || C == NULL || D == NULL || R == NULL)
204 {
205 printf("Error: Host memory allocation failed!\n");
206 return -1;
207 }
208
209 /*******************
210 * Initialize data
211 *******************/
212
213 memset(R, 0, sizeof(floatType) * elementsR);
214 for (uint64_t i = 0; i < elementsA; i++)
215 A[i] = ((floatType) rand()) / RAND_MAX;
216 for (uint64_t i = 0; i < elementsB; i++)
217 B[i] = ((floatType) rand()) / RAND_MAX;
218 for (uint64_t i = 0; i < elementsC; i++)
219 C[i] = ((floatType) rand()) / RAND_MAX;
220 for (uint64_t i = 0; i < elementsD; i++)
221 D[i] = ((floatType) rand()) / RAND_MAX;
222
223 HANDLE_CUDA_ERROR( cudaMemcpy(rawDataIn_d[0], A, sizeA, cudaMemcpyHostToDevice) );
224 HANDLE_CUDA_ERROR( cudaMemcpy(rawDataIn_d[1], B, sizeB, cudaMemcpyHostToDevice) );
225 HANDLE_CUDA_ERROR( cudaMemcpy(rawDataIn_d[2], C, sizeC, cudaMemcpyHostToDevice) );
226 HANDLE_CUDA_ERROR( cudaMemcpy(rawDataIn_d[3], D, sizeD, cudaMemcpyHostToDevice) );
227
228 if(verbose)
229 printf("Allocated GPU memory for data, and initialize data\n");
cuTensorNet 句柄和网络描述符¶
接下来，我们通过 cutensornetCreate() 初始化 cuTensorNet 库，并使用所需的张量模式、范围和步幅以及数据和计算类型创建网络描述符。请注意，创建的库上下文将与当前活动的 GPU 关联。
232 /*************************
233 * cuTensorNet
234 *************************/
235
236 cudaStream_t stream;
237 HANDLE_CUDA_ERROR( cudaStreamCreate(&stream) );
238
239 cutensornetHandle_t handle;
240 HANDLE_ERROR( cutensornetCreate(&handle) );
241
242 const int32_t nmodeA = modesA.size();
243 const int32_t nmodeB = modesB.size();
244 const int32_t nmodeC = modesC.size();
245 const int32_t nmodeD = modesD.size();
246 const int32_t nmodeR = modesR.size();
247
248 /*******************************
249 * Create Network Descriptor
250 *******************************/
251
252 const int32_t* modesIn[] = {modesA.data(), modesB.data(), modesC.data(), modesD.data()};
253 int32_t const numModesIn[] = {nmodeA, nmodeB, nmodeC, nmodeD};
254 const int64_t* extentsIn[] = {extentA.data(), extentB.data(), extentC.data(), extentD.data()};
255 const int64_t* stridesIn[] = {NULL, NULL, NULL, NULL}; // strides are optional; if no stride is provided, cuTensorNet assumes a generalized column-major data layout
256
257 // Set up tensor network
258 cutensornetNetworkDescriptor_t descNet;
259 HANDLE_ERROR( cutensornetCreateNetworkDescriptor(handle,
260 numInputs, numModesIn, extentsIn, stridesIn, modesIn, NULL,
261 nmodeR, extentR.data(), /*stridesOut = */NULL, modesR.data(),
262 typeData, typeCompute,
263 &descNet) );
264
265 if(verbose)
266 printf("Initialized the cuTensorNet library and created a tensor network descriptor\n");
最佳收缩顺序和切片¶
在此阶段，我们可以部署 cuTensorNet 优化器来查找优化的收缩路径和切片组合。我们根据可用内存资源选择执行收缩所需的工作区限制，并将其作为约束提供给优化器。然后，我们创建 cutensornetContractionOptimizerConfig_t 类型的优化器配置对象，以指定各种优化器选项，并将其提供给通过 cutensornetContractionOptimize() 调用的优化器。优化器的结果将在 cutensornetContractionOptimizerInfo_t 类型的优化器信息对象中返回。
269 /*******************************
270 * Choose workspace limit based on available resources.
271 *******************************/
272
273 size_t freeMem, totalMem;
274 HANDLE_CUDA_ERROR( cudaMemGetInfo(&freeMem, &totalMem) );
275 uint64_t workspaceLimit = (uint64_t)((double)freeMem * 0.9);
276 if(verbose)
277 printf("Workspace limit = %lu\n", workspaceLimit);
278
279 /*******************************
280 * Find "optimal" contraction order and slicing
281 *******************************/
282
283 cutensornetContractionOptimizerConfig_t optimizerConfig;
284 HANDLE_ERROR( cutensornetCreateContractionOptimizerConfig(handle, &optimizerConfig) );
285
286 // Set the desired number of hyper-samples (defaults to 0)
287 int32_t num_hypersamples = 8;
288 HANDLE_ERROR( cutensornetContractionOptimizerConfigSetAttribute(handle,
289 optimizerConfig,
290 CUTENSORNET_CONTRACTION_OPTIMIZER_CONFIG_HYPER_NUM_SAMPLES,
291 &num_hypersamples,
292 sizeof(num_hypersamples)) );
293
294 // Create contraction optimizer info and find an optimized contraction path
295 cutensornetContractionOptimizerInfo_t optimizerInfo;
296 HANDLE_ERROR( cutensornetCreateContractionOptimizerInfo(handle, descNet, &optimizerInfo) );
297
298 HANDLE_ERROR( cutensornetContractionOptimize(handle,
299 descNet,
300 optimizerConfig,
301 workspaceLimit,
302 optimizerInfo) );
303
304 // Query the number of slices the tensor network execution will be split into
305 int64_t numSlices = 0;
306 HANDLE_ERROR( cutensornetContractionOptimizerInfoGetAttribute(
307 handle,
308 optimizerInfo,
309 CUTENSORNET_CONTRACTION_OPTIMIZER_INFO_NUM_SLICES,
310 &numSlices,
311 sizeof(numSlices)) );
312 assert(numSlices > 0);
313
314 if(verbose)
315 printf("Found an optimized contraction path using cuTensorNet optimizer\n");
也可以绕过 cuTensorNet 优化器，并通过 cutensornetContractionOptimizerInfoSetAttribute() 将预定的收缩路径以及切片信息直接导入到优化器信息对象。
创建工作区描述符并分配工作区内存¶
接下来，我们创建一个工作区描述符，计算工作区大小，并查询收缩网络所需的最小工作区大小。然后，我们为工作区分配设备内存，并在工作区描述符中设置它。工作区描述符将提供给收缩计划。
318 /*******************************
319 * Create workspace descriptor, allocate workspace, and set it.
320 *******************************/
321
322 cutensornetWorkspaceDescriptor_t workDesc;
323 HANDLE_ERROR( cutensornetCreateWorkspaceDescriptor(handle, &workDesc) );
324
325 int64_t requiredWorkspaceSize = 0;
326 HANDLE_ERROR( cutensornetWorkspaceComputeContractionSizes(handle,
327 descNet,
328 optimizerInfo,
329 workDesc) );
330
331 HANDLE_ERROR( cutensornetWorkspaceGetMemorySize(handle,
332 workDesc,
333 CUTENSORNET_WORKSIZE_PREF_MIN,
334 CUTENSORNET_MEMSPACE_DEVICE,
335 CUTENSORNET_WORKSPACE_SCRATCH,
336 &requiredWorkspaceSize) );
337
338 void* work = nullptr;
339 HANDLE_CUDA_ERROR( cudaMalloc(&work, requiredWorkspaceSize) );
340
341 HANDLE_ERROR( cutensornetWorkspaceSetMemory(handle,
342 workDesc,
343 CUTENSORNET_MEMSPACE_DEVICE,
344 CUTENSORNET_WORKSPACE_SCRATCH,
345 work,
346 requiredWorkspaceSize) );
347
348 if(verbose)
349 printf("Allocated and set up the GPU workspace\n");
收缩计划和自动调优¶
我们创建一个张量网络收缩计划，其中包含 cuTENSOR 的所有成对收缩计划。可选地，我们可以自动调整计划，以便 cuTENSOR 为每个成对收缩选择最佳内核。此收缩计划可以重用于许多（可能不同的）数据输入，从而避免冗余地初始化此计划的成本。
352 /*******************************
353 * Initialize the pairwise contraction plan (for cuTENSOR).
354 *******************************/
355
356 cutensornetContractionPlan_t plan;
357 HANDLE_ERROR( cutensornetCreateContractionPlan(handle,
358 descNet,
359 optimizerInfo,
360 workDesc,
361 &plan) );
362
363 /*******************************
364 * Optional: Auto-tune cuTENSOR's cutensorContractionPlan to pick the fastest kernel
365 * for each pairwise tensor contraction.
366 *******************************/
367 cutensornetContractionAutotunePreference_t autotunePref;
368 HANDLE_ERROR( cutensornetCreateContractionAutotunePreference(handle,
369 &autotunePref) );
370
371 const int numAutotuningIterations = 5; // may be 0
372 HANDLE_ERROR( cutensornetContractionAutotunePreferenceSetAttribute(
373 handle,
374 autotunePref,
375 CUTENSORNET_CONTRACTION_AUTOTUNE_MAX_ITERATIONS,
376 &numAutotuningIterations,
377 sizeof(numAutotuningIterations)) );
378
379 // Modify the plan again to find the best pair-wise contractions
380 HANDLE_ERROR( cutensornetContractionAutotune(handle,
381 plan,
382 rawDataIn_d,
383 R_d,
384 workDesc,
385 autotunePref,
386 stream) );
387
388 HANDLE_ERROR( cutensornetDestroyContractionAutotunePreference(autotunePref) );
389
390 if(verbose)
391 printf("Created a contraction plan for cuTensorNet and optionally auto-tuned it\n");
张量网络收缩执行¶
最后，我们根据需要多次收缩张量网络，每次可能使用不同的输入。张量网络切片（捕获为 cutensornetSliceGroup_t 对象）使用相同的收缩计划计算。为方便起见，当目标是收缩网络中的所有切片时，可以将 NULL 提供给 cutensornetContractSlices() 函数，而不是切片组。我们还清理并释放分配的资源。
394 /**********************
395 * Execute the tensor network contraction
396 **********************/
397
398 // Create a cutensornetSliceGroup_t object from a range of slice IDs
399 cutensornetSliceGroup_t sliceGroup{};
400 HANDLE_ERROR( cutensornetCreateSliceGroupFromIDRange(handle, 0, numSlices, 1, &sliceGroup) );
401
402 GPUTimer timer {stream};
403 double minTimeCUTENSORNET = 1e100;
404 const int numRuns = 3; // number of repeats to get stable performance results
405 for (int i = 0; i < numRuns; ++i)
406 {
407 HANDLE_CUDA_ERROR( cudaMemcpy(R_d, R, sizeR, cudaMemcpyHostToDevice) ); // restore the output tensor on GPU
408 HANDLE_CUDA_ERROR( cudaDeviceSynchronize() );
409
410 /*
411 * Contract all slices of the tensor network
412 */
413 timer.start();
414
415 int32_t accumulateOutput = 0; // output tensor data will be overwritten
416 HANDLE_ERROR( cutensornetContractSlices(handle,
417 plan,
418 rawDataIn_d,
419 R_d,
420 accumulateOutput,
421 workDesc,
422 sliceGroup, // alternatively, NULL can also be used to contract over all slices instead of specifying a sliceGroup object
423 stream) );
424
425 // Synchronize and measure best timing
426 auto time = timer.seconds();
427 minTimeCUTENSORNET = (time > minTimeCUTENSORNET) ? minTimeCUTENSORNET : time;
428 }
429
430 if(verbose)
431 printf("Contracted the tensor network, each slice used the same contraction plan\n");
432
433 // Print the 1-norm of the output tensor (verification)
434 HANDLE_CUDA_ERROR( cudaStreamSynchronize(stream) );
435 HANDLE_CUDA_ERROR( cudaMemcpy(R, R_d, sizeR, cudaMemcpyDeviceToHost) ); // restore the output tensor on Host
436 double norm1 = 0.0;
437 for (int64_t i = 0; i < elementsR; ++i) {
438 norm1 += std::abs(R[i]);
439 }
440 if(verbose)
441 printf("Computed the 1-norm of the output tensor: %e\n", norm1);
442
443 /*************************/
444
445 // Query the total Flop count for the tensor network contraction
446 double flops {0.0};
447 HANDLE_ERROR( cutensornetContractionOptimizerInfoGetAttribute(
448 handle,
449 optimizerInfo,
450 CUTENSORNET_CONTRACTION_OPTIMIZER_INFO_FLOP_COUNT,
451 &flops,
452 sizeof(flops)) );
453
454 if(verbose) {
455 printf("Number of tensor network slices = %ld\n", numSlices);
456 printf("Tensor network contraction time (ms) = %.3f\n", minTimeCUTENSORNET * 1000.f);
457 }
458
459 // Free cuTensorNet resources
460 HANDLE_ERROR( cutensornetDestroySliceGroup(sliceGroup) );
461 HANDLE_ERROR( cutensornetDestroyContractionPlan(plan) );
462 HANDLE_ERROR( cutensornetDestroyWorkspaceDescriptor(workDesc) );
463 HANDLE_ERROR( cutensornetDestroyContractionOptimizerInfo(optimizerInfo) );
464 HANDLE_ERROR( cutensornetDestroyContractionOptimizerConfig(optimizerConfig) );
465 HANDLE_ERROR( cutensornetDestroyNetworkDescriptor(descNet) );
466 HANDLE_ERROR( cutensornetDestroy(handle) );
467
468 // Free Host memory resources
469 if (R) free(R);
470 if (D) free(D);
471 if (C) free(C);
472 if (B) free(B);
473 if (A) free(A);
474
475 // Free GPU memory resources
476 if (work) cudaFree(work);
477 if (R_d) cudaFree(R_d);
478 if (rawDataIn_d[0]) cudaFree(rawDataIn_d[0]);
479 if (rawDataIn_d[1]) cudaFree(rawDataIn_d[1]);
480 if (rawDataIn_d[2]) cudaFree(rawDataIn_d[2]);
481 if (rawDataIn_d[3]) cudaFree(rawDataIn_d[3]);
482
483 if(verbose)
484 printf("Freed resources and exited\n");
485
486 return 0;
487}
回想一下，完整的示例代码可以在 NVIDIA/cuQuantum 存储库中找到（此处）。
代码示例（自动切片分布式并行化）¶
可以轻松地调整 代码示例（串行） 并启用跨多个/许多 GPU 设备（跨多个/许多节点）的自动并行执行。我们将通过使用消息传递接口 (MPI) 作为通信层的示例来说明这一点。下面我们展示了为了启用分布式并行执行而需要进行的少量添加，而无需对原始串行源代码进行任何更改。完整的 MPI 自动示例代码可以在 NVIDIA/cuQuantum 存储库中找到。为了启用自动并行性，cuTensorNet 要求
环境变量
$CUTENSORNET_COMM_LIB设置为包装器共享库libcutensornet_distributed_interface_mpi.so的路径，并且可执行文件链接到 CUDA-aware MPI 库
上面安装指南中给出了设置这些的详细说明。
首先，除了 头文件和数据类型 中提到的头文件和定义之外，我们还包括 MPI 头文件并定义一个宏来处理 MPI 错误。我们还需要初始化 MPI 服务，并为每个 MPI 进程分配一个唯一的 GPU 设备，该设备稍后将与 MPI 进程内部创建的 cuTensorNet 库句柄关联。
20#include <mpi.h>
44#define HANDLE_MPI_ERROR(x) \
45{ const auto err = x; \
46 if( err != MPI_SUCCESS ) \
47 { char error[MPI_MAX_ERROR_STRING]; int len; \
48 MPI_Error_string(err, error, &len); \
49 printf("MPI Error: %s in line %d\n", error, __LINE__); \
50 fflush(stdout); \
51 MPI_Abort(MPI_COMM_WORLD, err); \
52 } \
53};
MPI 服务初始化必须在第一个 cutensornetCreate() 调用之前进行，该调用创建 cuTensorNet 库句柄。尝试在初始化 MPI 服务之前调用 cutensornetCreate() 将导致错误。
97 // Initialize MPI
98 HANDLE_MPI_ERROR( MPI_Init(&argc, &argv) );
99 int rank {-1};
100 HANDLE_MPI_ERROR( MPI_Comm_rank(MPI_COMM_WORLD, &rank) );
101 int numProcs {0};
102 HANDLE_MPI_ERROR( MPI_Comm_size(MPI_COMM_WORLD, &numProcs) );
如果位于同一节点上的多个 GPU 设备对 MPI 进程可见，我们需要为每个 MPI 进程选择一个独占 GPU 设备。如果您的 MPI 库实现提供的 mpirun（或 mpiexec）命令设置了一个环境变量，该变量显示 MPI 进程在其调用期间的排名，则可以使用该环境变量来设置 CUDA_VISIBLE_DEVICES 以指向专门分配给 MPI 进程的特定单个 GPU 设备（例如，Open MPI 为此目的提供了 ${OMPI_COMM_WORLD_LOCAL_RANK}）。否则，可以手动设置 GPU 设备，如下所示。
122 // Set GPU device based on ranks and nodes
123 int numDevices {0};
124 HANDLE_CUDA_ERROR( cudaGetDeviceCount(&numDevices) );
125 const int deviceId = rank % numDevices; // we assume that the processes are mapped to nodes in contiguous chunks
126 HANDLE_CUDA_ERROR( cudaSetDevice(deviceId) );
127 cudaDeviceProp prop;
128 HANDLE_CUDA_ERROR( cudaGetDeviceProperties(&prop, deviceId) );
接下来，我们按照 定义张量网络和张量大小 中所述定义张量网络。在每个进程一个 GPU 设备模型中，张量网络（包括操作数和结果数据）在每个进程上复制。根进程初始化输入数据并将其广播到其他进程。
258 /*******************
259 * Initialize data
260 *******************/
261
262 memset(R, 0, sizeof(floatType) * elementsR);
263 if(rank == 0)
264 {
265 for (uint64_t i = 0; i < elementsA; i++)
266 A[i] = ((floatType) rand()) / RAND_MAX;
267 for (uint64_t i = 0; i < elementsB; i++)
268 B[i] = ((floatType) rand()) / RAND_MAX;
269 for (uint64_t i = 0; i < elementsC; i++)
270 C[i] = ((floatType) rand()) / RAND_MAX;
271 for (uint64_t i = 0; i < elementsD; i++)
272 D[i] = ((floatType) rand()) / RAND_MAX;
273 }
274
275 // Broadcast input data to all ranks
276 HANDLE_MPI_ERROR( MPI_Bcast(A, elementsA, floatTypeMPI, 0, MPI_COMM_WORLD) );
277 HANDLE_MPI_ERROR( MPI_Bcast(B, elementsB, floatTypeMPI, 0, MPI_COMM_WORLD) );
278 HANDLE_MPI_ERROR( MPI_Bcast(C, elementsC, floatTypeMPI, 0, MPI_COMM_WORLD) );
279 HANDLE_MPI_ERROR( MPI_Bcast(D, elementsD, floatTypeMPI, 0, MPI_COMM_WORLD) );
280
281 // Copy data to GPU
282 HANDLE_CUDA_ERROR( cudaMemcpy(rawDataIn_d[0], A, sizeA, cudaMemcpyHostToDevice) );
283 HANDLE_CUDA_ERROR( cudaMemcpy(rawDataIn_d[1], B, sizeB, cudaMemcpyHostToDevice) );
284 HANDLE_CUDA_ERROR( cudaMemcpy(rawDataIn_d[2], C, sizeC, cudaMemcpyHostToDevice) );
285 HANDLE_CUDA_ERROR( cudaMemcpy(rawDataIn_d[3], D, sizeD, cudaMemcpyHostToDevice) );
286
287 if(verbose)
288 printf("Allocated GPU memory for data, and initialize data\n");
一旦 MPI 服务被初始化并且 cuTensorNet 库句柄在之后被创建，就可以通过调用 cutensornetDistributedResetConfiguration() 来激活分布式并行执行。按照标准实践，用户的代码需要通过 MPI_Comm_dup 创建重复的 MPI 通信器。然后，通过将指向重复的 MPI 通信器的指针及其大小（以字节为单位）传递给 cutensornetDistributedResetConfiguration() 调用，将重复的 MPI 通信器与 cuTensorNet 库句柄关联。MPI 通信器将存储在 cuTensorNet 库句柄内部，以便对张量网络收缩路径查找器和张量网络收缩执行器的所有后续调用将在所有参与的 MPI 进程（每个 MPI 进程都与其自己的 GPU 关联）之间并行化。
342 /*******************************
343 * Activate distributed (parallel) execution prior to
344 * calling contraction path finder and contraction executor
345 *******************************/
346 // HANDLE_ERROR( cutensornetDistributedResetConfiguration(handle, NULL, 0) ); // resets back to serial execution
347 MPI_Comm cutnComm;
348 HANDLE_MPI_ERROR( MPI_Comm_dup(MPI_COMM_WORLD, &cutnComm) ); // duplicate MPI communicator
349 HANDLE_ERROR( cutensornetDistributedResetConfiguration(handle, &cutnComm, sizeof(cutnComm)) );
350 if(verbose)
351 printf("Reset distributed MPI configuration\n");
注意
cutensornetDistributedResetConfiguration() 是一个集体调用，必须由所有参与的 MPI 进程执行。
这种分布式并行化模型的 API 使得在多个 GPU/节点上运行为串行执行编写的源代码变得简单。本质上，所有 MPI 进程将执行完全相同的（串行）源代码，同时在张量网络收缩路径查找器和张量网络收缩执行器调用内部自动执行分布式并行化。张量网络收缩路径查找器的并行化仅在请求的超样本数大于零时才会发生。然而，无论如何，分布式并行化的激活必须先于张量网络收缩路径查找器的调用。也就是说，张量网络收缩路径查找器和张量网络收缩执行的调用必须严格在通过 cutensornetDistributedResetConfiguration() 激活分布式并行化之后完成。当分布式配置设置为并行模式时，通常期望用户通过调用 cutensornetContractSlices() 函数来调用张量网络收缩执行，该函数提供全范围的张量网络切片，这些切片将自动分布在所有 MPI 进程中。
由于张量网络的大小必须足够大，才能从分布式执行中获得加速益处，因此较小的张量网络（仅包含单个切片的张量网络）仍然可以在不进行分布式并行化的情况下进行处理。这可以通过调用 cutensornetDistributedResetConfiguration() 并使用 NULL 参数来代替 MPI 通信器指针来实现（与之前一样，这应该在调用张量网络收缩路径查找器之前完成）。也就是说，分布式并行化和冗余串行执行之间的切换可以基于每个张量网络进行。用户可以决定哪些（较大的）张量网络以并行方式处理，哪些（较小的）张量网络以冗余串行方式处理，只需适当地重置分布式配置即可。在这两种情况下，所有 MPI 进程都将在张量网络执行结束时生成相同的输出张量（结果）。
注意
在当前版本的 cuTensorNet 库中，由 cutensornetContractSlices() 调用触发的并行张量网络收缩执行将阻塞提供的 CUDA 流以及调用 CPU 线程，直到所有 MPI 进程上的执行完成。这是一个临时限制，将在未来版本的 cuTensorNet 库中解除。届时，对 cutensornetContractSlices() 的调用将是完全异步的，类似于串行执行的情况。此外，对于所有 MPI 进程的显式同步（屏障），可以集体调用 cutensornetDistributedSynchronize()。
在终止之前，需要最终确定 MPI 服务。
561 // Shut down MPI service
562 HANDLE_MPI_ERROR( MPI_Finalize() );
完整的 MPI 自动示例 可以在 NVIDIA/cuQuantum 仓库中找到。
代码示例（手动切片分布式并行化）¶
对于高级用户，也可以（但更复杂）调整 代码示例（串行） 以在多个 GPU 设备上显式地并行化张量网络收缩操作的执行。这里我们也将使用 MPI 作为通信层。为了简洁起见，我们将仅展示需要在串行示例之上进行的更改。完整的 MPI 手动示例代码 可以在 NVIDIA/cuQuantum 仓库中找到。请注意，此示例 **不** 需要 CUDA-aware MPI。
首先，除了 头文件和数据类型 中提到的头文件和定义之外，我们还需要包含 MPI 头文件并定义一个宏来处理 MPI 错误。我们还需要初始化 MPI 服务，并将每个 MPI 进程与其自己的 GPU 设备关联起来，如前所述。
20#include <mpi.h>
44#define HANDLE_MPI_ERROR(x) \
45{ const auto err = x; \
46 if( err != MPI_SUCCESS ) \
47 { char error[MPI_MAX_ERROR_STRING]; int len; \
48 MPI_Error_string(err, error, &len); \
49 printf("MPI Error: %s in line %d\n", error, __LINE__); \
50 fflush(stdout); \
51 MPI_Abort(MPI_COMM_WORLD, err); \
52 } \
53};
97 // Initialize MPI
98 HANDLE_MPI_ERROR( MPI_Init(&argc, &argv) );
99 int rank {-1};
100 HANDLE_MPI_ERROR( MPI_Comm_rank(MPI_COMM_WORLD, &rank) );
101 int numProcs {0};
102 HANDLE_MPI_ERROR( MPI_Comm_size(MPI_COMM_WORLD, &numProcs) );
122 // Set GPU device based on ranks and nodes
123 int numDevices {0};
124 HANDLE_CUDA_ERROR( cudaGetDeviceCount(&numDevices) );
125 const int deviceId = rank % numDevices; // we assume that the processes are mapped to nodes in contiguous chunks
126 HANDLE_CUDA_ERROR( cudaSetDevice(deviceId) );
127 cudaDeviceProp prop;
128 HANDLE_CUDA_ERROR( cudaGetDeviceProperties(&prop, deviceId) );
接下来，我们按照 定义张量网络和张量大小 中所述定义张量网络。在每个进程一个 GPU 设备的模型中，张量网络（包括操作数和结果数据）在每个进程上复制。根进程初始化输入数据并将其广播到其他进程。
258 /*******************
259 * Initialize data
260 *******************/
261
262 memset(R, 0, sizeof(floatType) * elementsR);
263 if(rank == 0)
264 {
265 for (uint64_t i = 0; i < elementsA; i++)
266 A[i] = ((floatType) rand()) / RAND_MAX;
267 for (uint64_t i = 0; i < elementsB; i++)
268 B[i] = ((floatType) rand()) / RAND_MAX;
269 for (uint64_t i = 0; i < elementsC; i++)
270 C[i] = ((floatType) rand()) / RAND_MAX;
271 for (uint64_t i = 0; i < elementsD; i++)
272 D[i] = ((floatType) rand()) / RAND_MAX;
273 }
274
275 // Broadcast input data to all ranks
276 HANDLE_MPI_ERROR( MPI_Bcast(A, elementsA, floatTypeMPI, 0, MPI_COMM_WORLD) );
277 HANDLE_MPI_ERROR( MPI_Bcast(B, elementsB, floatTypeMPI, 0, MPI_COMM_WORLD) );
278 HANDLE_MPI_ERROR( MPI_Bcast(C, elementsC, floatTypeMPI, 0, MPI_COMM_WORLD) );
279 HANDLE_MPI_ERROR( MPI_Bcast(D, elementsD, floatTypeMPI, 0, MPI_COMM_WORLD) );
280
281 // Copy data to GPU
282 HANDLE_CUDA_ERROR( cudaMemcpy(rawDataIn_d[0], A, sizeA, cudaMemcpyHostToDevice) );
283 HANDLE_CUDA_ERROR( cudaMemcpy(rawDataIn_d[1], B, sizeB, cudaMemcpyHostToDevice) );
284 HANDLE_CUDA_ERROR( cudaMemcpy(rawDataIn_d[2], C, sizeC, cudaMemcpyHostToDevice) );
285 HANDLE_CUDA_ERROR( cudaMemcpy(rawDataIn_d[3], D, sizeD, cudaMemcpyHostToDevice) );
286
287 if(verbose)
288 printf("Allocated GPU memory for data, and initialize data\n");
然后，我们在每个进程上创建库句柄和张量网络描述符，如 cuTensorNet 句柄和网络描述符 中所述。
接下来，我们为我们的张量网络找到最佳的收缩路径和切片组合。我们将在所有进程上运行 cuTensorNet 优化器，并确定哪个进程具有 FLOP 计数方面最佳的路径。然后，我们将在此进程上打包优化器信息对象，广播打包的缓冲区，并在所有其他进程上解包它。现在，每个进程都具有相同的优化器信息对象，我们使用它来计算每个进程的切片份额。
363 // Compute the path on all ranks so that we can choose the path with the lowest cost. Note that since this is a tiny
364 // example with 4 operands, all processes will compute the same globally optimal path. This is not the case for large
365 // tensor networks. For large networks, hyper-optimization does become beneficial.
366
367 // Enforce tensor network slicing (for parallelization)
368 const int32_t min_slices = numProcs;
369 HANDLE_ERROR( cutensornetContractionOptimizerConfigSetAttribute(handle,
370 optimizerConfig,
371 CUTENSORNET_CONTRACTION_OPTIMIZER_CONFIG_SLICER_MIN_SLICES,
372 &min_slices,
373 sizeof(min_slices)) );
374
375 // Find an optimized tensor network contraction path on each MPI process
376 HANDLE_ERROR( cutensornetContractionOptimize(handle,
377 descNet,
378 optimizerConfig,
379 workspaceLimit,
380 optimizerInfo) );
381
382 // Query the obtained Flop count
383 double flops{-1.};
384 HANDLE_ERROR( cutensornetContractionOptimizerInfoGetAttribute(handle,
385 optimizerInfo,
386 CUTENSORNET_CONTRACTION_OPTIMIZER_INFO_FLOP_COUNT,
387 &flops,
388 sizeof(flops)) );
389
390 // Choose the contraction path with the lowest Flop cost
391 struct {
392 double value;
393 int rank;
394 } in{flops, rank}, out;
395 HANDLE_MPI_ERROR( MPI_Allreduce(&in, &out, 1, MPI_DOUBLE_INT, MPI_MINLOC, MPI_COMM_WORLD) );
396 const int sender = out.rank;
397 flops = out.value;
398
399 if (verbose)
400 printf("Process %d has the path with the lowest FLOP count %lf\n", sender, flops);
401
402 // Get the buffer size for optimizerInfo and broadcast it
403 size_t bufSize {0};
404 if (rank == sender)
405 {
406 HANDLE_ERROR( cutensornetContractionOptimizerInfoGetPackedSize(handle, optimizerInfo, &bufSize) );
407 }
408 HANDLE_MPI_ERROR( MPI_Bcast(&bufSize, 1, MPI_INT64_T, sender, MPI_COMM_WORLD) );
409
410 // Allocate a buffer
411 std::vector<char> buffer(bufSize);
412
413 // Pack optimizerInfo on sender and broadcast it
414 if (rank == sender)
415 {
416 HANDLE_ERROR( cutensornetContractionOptimizerInfoPackData(handle, optimizerInfo, buffer.data(), bufSize) );
417 }
418 HANDLE_MPI_ERROR( MPI_Bcast(buffer.data(), bufSize, MPI_CHAR, sender, MPI_COMM_WORLD) );
419
420 // Unpack optimizerInfo from the buffer
421 if (rank != sender)
422 {
423 HANDLE_ERROR( cutensornetUpdateContractionOptimizerInfoFromPackedData(handle, buffer.data(), bufSize, optimizerInfo) );
424 }
425
426 // Query the number of slices the tensor network execution will be split into
427 int64_t numSlices = 0;
428 HANDLE_ERROR( cutensornetContractionOptimizerInfoGetAttribute(
429 handle,
430 optimizerInfo,
431 CUTENSORNET_CONTRACTION_OPTIMIZER_INFO_NUM_SLICES,
432 &numSlices,
433 sizeof(numSlices)) );
434 assert(numSlices > 0);
435
436 // Calculate each process's share of the slices
437 int64_t procChunk = numSlices / numProcs;
438 int extra = numSlices % numProcs;
439 int procSliceBegin = rank * procChunk + std::min(rank, extra);
440 int procSliceEnd = (rank == numProcs - 1) ? numSlices : (rank + 1) * procChunk + std::min(rank + 1, extra);
我们现在创建工作空间描述符并分配内存，如 创建工作空间描述符并分配工作空间内存 中所述，并创建 收缩计划和自动调优 张量网络。
接下来，在每个进程上，我们创建一个切片组（参见 cutensornetSliceGroup_t），该切片组对应于其张量网络切片份额。然后，我们将此切片组对象提供给 cutensornetContractSlices() 函数，以在每个进程上获得部分收缩结果。
530 // Create a cutensornetSliceGroup_t object from a range of slice IDs
531 cutensornetSliceGroup_t sliceGroup{};
532 HANDLE_ERROR( cutensornetCreateSliceGroupFromIDRange(handle, procSliceBegin, procSliceEnd, 1, &sliceGroup) );
553 HANDLE_ERROR( cutensornetContractSlices(handle,
554 plan,
555 rawDataIn_d,
556 R_d,
557 accumulateOutput,
558 workDesc,
559 sliceGroup,
560 stream) );
最后，我们对部分贡献求和，以获得张量网络收缩的结果。
566 // Perform Allreduce operation on the output tensor
567 HANDLE_CUDA_ERROR( cudaStreamSynchronize(stream) );
568 HANDLE_CUDA_ERROR( cudaMemcpy(R, R_d, sizeR, cudaMemcpyDeviceToHost) ); // restore the output tensor on Host
569 HANDLE_MPI_ERROR( MPI_Allreduce(MPI_IN_PLACE, R, elementsR, floatTypeMPI, MPI_SUM, MPI_COMM_WORLD) );
在终止之前，需要最终确定 MPI 服务。
631 // Shut down MPI service
632 HANDLE_MPI_ERROR( MPI_Finalize() );
完整的 MPI 手动示例 可以在 NVIDIA/cuQuantum 仓库中找到。
代码示例 (tensorQR)¶
代码示例 (tensorSVD)¶
代码示例 (MPS 分解)¶
代码示例 (中间张量重用)¶
代码示例 (梯度计算)¶
代码示例 (幅度切片)¶
代码示例 (期望值)¶
代码示例 (边际分布)¶
代码示例 (张量网络采样)¶
代码示例 (MPS 幅度切片)¶
代码示例 (MPS 期望值)¶
代码示例 (MPS 边际分布)¶
代码示例 (MPS 采样)¶
代码示例 (MPS 采样 QFT)¶
代码示例 (MPS 采样 MPO)¶
实用技巧¶
对于调试，可以设置环境变量
CUTENSORNET_LOG_LEVEL=n。级别n= 0, 1, …, 5 对应于cutensornetLoggerSetLevel()中描述和使用的记录器级别。环境变量CUTENSORNET_LOG_FILE=<filepath>可用于将日志输出重定向到<filepath>而不是stdout的自定义文件。