Code example (serial contraction)#
The following code example illustrates the common steps necessary to use cuTensorNet and also introduces typical tensor network operations. Specifically, the example performs the following tensor contraction:
\[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}\]
The full sample code can be found in the NVIDIA/cuQuantum repository (here).
Headers and data types#
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// cutensornetGetErrorString() returns the enum-name form of the status code,
19// while cutensornetGetLastError() returns a human-readable description of the
20// most recent error captured on the calling thread (when available). Print
21// both to give the user maximally actionable feedback on failure.
22#define HANDLE_ERROR(x) \
23 do { \
24 const auto err = x; \
25 if (err != CUTENSORNET_STATUS_SUCCESS) \
26 { \
27 printf("Error: %s in line %d\n", cutensornetGetErrorString(err), __LINE__); \
28 const char* details = cutensornetGetLastError(); \
29 if (details != nullptr && details[0] != '\0') \
30 printf("Details: %s\n", details); \
31 fflush(stdout); \
32 exit(EXIT_FAILURE); \
33 } \
34 } while (0)
35
36#define HANDLE_CUDA_ERROR(x) \
37 do { \
38 const auto err = x; \
39 if (err != cudaSuccess) \
40 { \
41 printf("CUDA Error: %s in line %d\n", cudaGetErrorString(err), __LINE__); \
42 fflush(stdout); \
43 exit(EXIT_FAILURE); \
44 } \
45 } while (0)
46
47// Usage: DEV_ATTR(cudaDevAttrClockRate, deviceId)
48#define DEV_ATTR(ENUMCONST, DID) \
49 ({ int v; \
50 HANDLE_CUDA_ERROR(cudaDeviceGetAttribute(&v, ENUMCONST, DID)); \
51 v; })
52
53
54struct GPUTimer
55{
56 GPUTimer(cudaStream_t stream) : stream_(stream)
57 {
58 HANDLE_CUDA_ERROR(cudaEventCreate(&start_));
59 HANDLE_CUDA_ERROR(cudaEventCreate(&stop_));
60 }
61
62 ~GPUTimer()
63 {
64 HANDLE_CUDA_ERROR(cudaEventDestroy(start_));
65 HANDLE_CUDA_ERROR(cudaEventDestroy(stop_));
66 }
67
68 void start() { HANDLE_CUDA_ERROR(cudaEventRecord(start_, stream_)); }
69
70 float seconds()
71 {
72 HANDLE_CUDA_ERROR(cudaEventRecord(stop_, stream_));
73 HANDLE_CUDA_ERROR(cudaEventSynchronize(stop_));
74 float time;
75 HANDLE_CUDA_ERROR(cudaEventElapsedTime(&time, start_, stop_));
76 return time * 1e-3;
77 }
78
79private:
80 cudaEvent_t start_, stop_;
81 cudaStream_t stream_;
82};
83
84int main()
85{
86 static_assert(sizeof(size_t) == sizeof(int64_t), "Please build this sample on a 64-bit architecture!");
87
88 bool verbose = true;
89
90 // Check cuTensorNet version
91 const size_t cuTensornetVersion = cutensornetGetVersion();
92 if (verbose) printf("cuTensorNet version: %ld\n", cuTensornetVersion);
93
94 // Set GPU device
95 int numDevices{0};
96 HANDLE_CUDA_ERROR(cudaGetDeviceCount(&numDevices));
97 const int deviceId = 0;
98 HANDLE_CUDA_ERROR(cudaSetDevice(deviceId));
99 cudaDeviceProp prop;
100 HANDLE_CUDA_ERROR(cudaGetDeviceProperties(&prop, deviceId));
101
102 if (verbose)
103 {
104 printf("===== device info ======\n");
105 printf("GPU-local-id:%d\n", deviceId);
106 printf("GPU-name:%s\n", prop.name);
107 printf("GPU-clock:%d\n", DEV_ATTR(cudaDevAttrClockRate, deviceId));
108 printf("GPU-memoryClock:%d\n", DEV_ATTR(cudaDevAttrMemoryClockRate, deviceId));
109 printf("GPU-nSM:%d\n", prop.multiProcessorCount);
110 printf("GPU-major:%d\n", prop.major);
111 printf("GPU-minor:%d\n", prop.minor);
112 printf("========================\n");
113 }
114
115 typedef float floatType;
116 cudaDataType_t typeData = CUDA_R_32F;
117 cutensornetComputeType_t typeCompute = CUTENSORNET_COMPUTE_32F;
118
119 if (verbose) printf("Included headers and defined data types\n");
Define tensor network and tensor sizes#
Next, we define the topology of the tensor network (i.e., the modes of the tensors, their extents, and their connectivity).
123 /**************************************************************************************
124 * 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}
125 **************************************************************************************/
126
127 constexpr int32_t numInputs = 4;
128
129 // Create vectors of tensor modes
130 std::vector<std::vector<int32_t>> tensorModes{ // for input tensors & output tensor
131 // input tensors
132 {'a', 'b', 'c', 'd', 'e', 'f'}, // tensor A
133 {'b', 'g', 'h', 'e', 'i', 'j'}, // tensor B
134 {'m', 'a', 'g', 'f', 'i', 'k'}, // tensor C
135 {'l', 'c', 'h', 'd', 'j', 'm'}, // tensor D
136 // output tensor
137 {'k', 'l'}, // tensor R
138 };
139
140 // Set mode extents
141 int64_t sameExtent = 16; // setting same extent for simplicity. In principle extents can differ.
142 std::unordered_map<int32_t, int64_t> extent;
143 for (auto& vec : tensorModes)
144 {
145 for (auto& mode : vec)
146 {
147 extent[mode] = sameExtent;
148 }
149 }
150
151 // Create a vector of extents for each tensor
152 std::vector<std::vector<int64_t>> tensorExtents; // for input tensors & output tensor
153 tensorExtents.resize(numInputs + 1); // hold inputs + output tensors
154 for (int32_t t = 0; t < numInputs + 1; ++t)
155 {
156 for (auto& mode : tensorModes[t]) tensorExtents[t].push_back(extent[mode]);
157 }
158
159 if (verbose) printf("Defined tensor network, modes, and extents\n");
Allocate memory and initialize data#
Next, we allocate memory for the tensor network operands and initialize them to random values.
162 /*****************
163 * Allocating data
164 *****************/
165
166 std::vector<size_t> tensorElements(numInputs + 1); // for input tensors & output tensor
167 std::vector<size_t> tensorSizes(numInputs + 1); // for input tensors & output tensor
168 size_t totalSize = 0;
169 for (int32_t t = 0; t < numInputs + 1; ++t)
170 {
171 size_t numElements = 1;
172 for (auto& mode : tensorModes[t]) numElements *= extent[mode];
173 tensorElements[t] = numElements;
174
175 tensorSizes[t] = sizeof(floatType) * numElements;
176 totalSize += tensorSizes[t];
177 }
178
179 if (verbose) printf("Total GPU memory used for tensor storage: %.2f GiB\n", (totalSize) / 1024. / 1024. / 1024);
180
181 void* tensorData_d[numInputs + 1]; // for input tensors & output tensor
182 for (int32_t t = 0; t < numInputs + 1; ++t)
183 {
184 HANDLE_CUDA_ERROR(cudaMalloc((void**)&tensorData_d[t], tensorSizes[t]));
185 }
186
187 floatType* tensorData_h[numInputs + 1]; // for input tensors & output tensor
188 for (int32_t t = 0; t < numInputs + 1; ++t)
189 {
190 tensorData_h[t] = (floatType*)malloc(tensorSizes[t]);
191 if (tensorData_h[t] == NULL)
192 {
193 printf("Error: Host memory allocation failed!\n");
194 return -1;
195 }
196 }
197
198 /*****************
199 * Initialize data
200 *****************/
201
202 // init output tensor to all 0s
203 memset(tensorData_h[numInputs], 0, tensorSizes[numInputs]);
204 // init input tensors to random values
205 for (int32_t t = 0; t < numInputs; ++t)
206 {
207 for (size_t e = 0; e < tensorElements[t]; ++e) tensorData_h[t][e] = ((floatType)rand()) / RAND_MAX;
208 }
209 // copy input data to device buffers
210 for (int32_t t = 0; t < numInputs; ++t)
211 {
212 HANDLE_CUDA_ERROR(cudaMemcpy(tensorData_d[t], tensorData_h[t], tensorSizes[t], cudaMemcpyHostToDevice));
213 }
cuTensorNet handle and network descriptor#
Next, we initialize the cuTensorNet library via cutensornetCreate().
Note that the created library context will be associated with the currently active GPU.
We create the network descriptor, and append the input tensors with the desired tensor modes and extents, as well as the data type.
We can, optionally, set the output tensor modes (if skipped, the output modes will be inferred).
Also, we can, optionally, set the compute mode on the network (if skipped, a default compute type, corresponding to the data type, will be used.
Refer to cutensornetCreateNetwork()).
216 /*************
217 * cuTensorNet
218 *************/
219
220 cudaStream_t stream;
221 HANDLE_CUDA_ERROR(cudaStreamCreate(&stream));
222
223 cutensornetHandle_t handle;
224 HANDLE_ERROR(cutensornetCreate(&handle));
225
226 if (verbose) printf("Allocated GPU memory for data, initialized data, and created library handle\n");
227
228 /****************
229 * Create Network
230 ****************/
231
232 // Set up tensor network
233 cutensornetNetworkDescriptor_t networkDesc;
234 HANDLE_ERROR(cutensornetCreateNetwork(handle, &networkDesc));
235
236 int64_t tensorIDs[numInputs]; // for input tensors
237 // attach the input tensors to the network
238 for (int32_t t = 0; t < numInputs; ++t)
239 {
240 HANDLE_ERROR(cutensornetNetworkAppendTensor(handle,
241 networkDesc,
242 tensorModes[t].size(),
243 tensorExtents[t].data(),
244 tensorModes[t].data(),
245 NULL,
246 typeData,
247 &tensorIDs[t]));
248 }
249
250 // set the output tensor
251 HANDLE_ERROR(cutensornetNetworkSetOutputTensor(handle,
252 networkDesc,
253 tensorModes[numInputs].size(),
254 tensorModes[numInputs].data(),
255 typeData));
256
257 // set the network compute type
258 HANDLE_ERROR(cutensornetNetworkSetAttribute(handle,
259 networkDesc,
260 CUTENSORNET_NETWORK_COMPUTE_TYPE,
261 &typeCompute,
262 sizeof(typeCompute)));
263 if (verbose) printf("Initialized the cuTensorNet library and created a tensor network descriptor\n");
Optimal contraction order and slicing#
At this stage, we can deploy the cuTensorNet optimizer to find an optimized contraction path and slicing combination.
We choose a limit for the workspace needed to perform the contraction based on the available memory resources, and provide it
to the optimizer as a constraint. We then create an optimizer configuration object of type cutensornetContractionOptimizerConfig_t
to specify various optimizer options and provide it to the optimizer, which is invoked via cutensornetContractionOptimize().
The results from the optimizer will be returned in an optimizer info object of type cutensornetContractionOptimizerInfo_t.
266 /******************************************************
267 * Choose workspace limit based on available resources.
268 ******************************************************/
269
270 size_t freeMem, totalMem;
271 HANDLE_CUDA_ERROR(cudaMemGetInfo(&freeMem, &totalMem));
272 uint64_t workspaceLimit = (uint64_t)((double)freeMem * 0.9);
273 if (verbose) printf("Workspace limit = %lu\n", workspaceLimit);
274
275 /*******************************
276 * Find "optimal" contraction order and slicing
277 *******************************/
278
279 cutensornetContractionOptimizerConfig_t optimizerConfig;
280 HANDLE_ERROR(cutensornetCreateContractionOptimizerConfig(handle, &optimizerConfig));
281
282 // Set the desired number of hyper-samples (defaults to 0)
283 int32_t num_hypersamples = 8;
284 HANDLE_ERROR(cutensornetContractionOptimizerConfigSetAttribute(handle,
285 optimizerConfig,
286 CUTENSORNET_CONTRACTION_OPTIMIZER_CONFIG_HYPER_NUM_SAMPLES,
287 &num_hypersamples,
288 sizeof(num_hypersamples)));
289
290 // Create contraction optimizer info and find an optimized contraction path
291 cutensornetContractionOptimizerInfo_t optimizerInfo;
292 HANDLE_ERROR(cutensornetCreateContractionOptimizerInfo(handle, networkDesc, &optimizerInfo));
293
294 HANDLE_ERROR(cutensornetContractionOptimize(handle,
295 networkDesc,
296 optimizerConfig,
297 workspaceLimit,
298 optimizerInfo));
299
300 // Query the number of slices the tensor network execution will be split into
301 int64_t numSlices = 0;
302 HANDLE_ERROR(cutensornetContractionOptimizerInfoGetAttribute(handle,
303 optimizerInfo,
304 CUTENSORNET_CONTRACTION_OPTIMIZER_INFO_NUM_SLICES,
305 &numSlices,
306 sizeof(numSlices)));
307 assert(numSlices > 0);
308
309 if (verbose) printf("Found an optimized contraction path using cuTensorNet optimizer\n");
It is also possible to bypass the cuTensorNet optimizer and import a pre-determined contraction path, as well as slicing information,
directly to the optimizer info object via cutensornetContractionOptimizerInfoSetAttribute(),
then attach it to the network via cutensornetNetworkSetOptimizerInfo().
Create workspace descriptor and allocate workspace memory#
Next, we create a workspace descriptor, compute the workspace sizes, and query the minimum workspace size needed to contract the network. We then allocate device memory for the workspace and set this in the workspace descriptor. The workspace descriptor will be provided to the contraction preparation and computation calls.
312 /*******************************
313 * Create workspace descriptor, allocate workspace, and set it.
314 *******************************/
315
316 cutensornetWorkspaceDescriptor_t workDesc;
317 HANDLE_ERROR(cutensornetCreateWorkspaceDescriptor(handle, &workDesc));
318
319 int64_t requiredWorkspaceSize = 0;
320 HANDLE_ERROR(cutensornetWorkspaceComputeContractionSizes(handle,
321 networkDesc,
322 optimizerInfo,
323 workDesc));
324
325 HANDLE_ERROR(cutensornetWorkspaceGetMemorySize(handle,
326 workDesc,
327 CUTENSORNET_WORKSIZE_PREF_MIN,
328 CUTENSORNET_MEMSPACE_DEVICE,
329 CUTENSORNET_WORKSPACE_SCRATCH,
330 &requiredWorkspaceSize));
331
332 void* work = nullptr;
333 HANDLE_CUDA_ERROR(cudaMalloc(&work, requiredWorkspaceSize));
334
335 HANDLE_ERROR(cutensornetWorkspaceSetMemory(handle,
336 workDesc,
337 CUTENSORNET_MEMSPACE_DEVICE,
338 CUTENSORNET_WORKSPACE_SCRATCH,
339 work,
340 requiredWorkspaceSize));
341
342 if (verbose) printf("Allocated and set up the GPU workspace\n");
Contraction preparation and auto-tuning#
We prepare the tensor network contraction, via cutensornetNetworkPrepareContraction(),
and set tensor’s data buffers and strides via cutensornetNetworkSetInputTensorMemory() and cutensornetNetworkSetOutputTensorMemory().
Optionally, we can auto-tune the contraction, via cutensornetNetworkAutotuneContraction(),
such that cuTENSOR selects the best kernel for each pairwise contraction.
This prepared network contraction can be reused for many (possibly different) data inputs, avoiding
the cost of initializing it redundantly.
345 /**************************
346 * Prepare the contraction.
347 **************************/
348
349 HANDLE_ERROR(cutensornetNetworkPrepareContraction(handle,
350 networkDesc,
351 workDesc));
352
353 // set tensor's data buffers and strides
354 for (int32_t t = 0; t < numInputs; ++t)
355 {
356 HANDLE_ERROR(cutensornetNetworkSetInputTensorMemory(handle,
357 networkDesc,
358 tensorIDs[t],
359 tensorData_d[t],
360 NULL));
361 }
362 HANDLE_ERROR(cutensornetNetworkSetOutputTensorMemory(handle,
363 networkDesc,
364 tensorData_d[numInputs],
365 NULL));
366 /****************************************************************
367 * Optional: Auto-tune the contraction to pick the fastest kernel
368 * for each pairwise tensor contraction.
369 ****************************************************************/
370 cutensornetNetworkAutotunePreference_t autotunePref;
371 HANDLE_ERROR(cutensornetCreateNetworkAutotunePreference(handle,
372 &autotunePref));
373
374 const int numAutotuningIterations = 5; // may be 0
375 HANDLE_ERROR(cutensornetNetworkAutotunePreferenceSetAttribute(handle,
376 autotunePref,
377 CUTENSORNET_NETWORK_AUTOTUNE_MAX_ITERATIONS,
378 &numAutotuningIterations,
379 sizeof(numAutotuningIterations)));
380
381 // Modify the network again to find the best pair-wise contractions
382 HANDLE_ERROR(cutensornetNetworkAutotuneContraction(handle,
383 networkDesc,
384 workDesc,
385 autotunePref,
386 stream));
387
388 HANDLE_ERROR(cutensornetDestroyNetworkAutotunePreference(autotunePref));
389
390 if (verbose) printf("Prepared the network contraction for cuTensorNet and optionally auto-tuned it\n");
Tensor network contraction execution#
Finally, we contract the tensor network as many times as needed, via cutensornetNetworkContract(),
possibly with different input each time (reset the data buffers and strides via cutensornetNetworkSetInputTensorMemory()).
Tensor network slices, captured as a cutensornetSliceGroup_t object, are computed using the same prepared network contraction.
For convenience, NULL can be provided to the cutensornetNetworkContract() function instead of a slice group when
the goal is to contract all the slices in the network. We also clean up and free allocated resources.
393 /****************************************
394 * Execute the tensor network contraction
395 ****************************************/
396
397 // Create a cutensornetSliceGroup_t object from a range of slice IDs
398 cutensornetSliceGroup_t sliceGroup{};
399 HANDLE_ERROR(cutensornetCreateSliceGroupFromIDRange(handle, 0, numSlices, 1, &sliceGroup));
400
401 GPUTimer timer{stream};
402 double minTimeCUTENSORNET = 1e100;
403 const int numRuns = 3; // number of repeats to get stable performance results
404 for (int i = 0; i < numRuns; ++i)
405 {
406 // reset the output tensor data
407 HANDLE_CUDA_ERROR(cudaMemcpy(tensorData_d[numInputs], tensorData_h[numInputs], tensorSizes[numInputs], cudaMemcpyHostToDevice));
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(cutensornetNetworkContract(handle,
417 networkDesc,
418 accumulateOutput,
419 workDesc,
420 sliceGroup, // alternatively, NULL can also be used to contract over all slices instead of specifying a sliceGroup object
421 stream));
422
423 // Synchronize and measure best timing
424 auto time = timer.seconds();
425 minTimeCUTENSORNET = (time > minTimeCUTENSORNET) ? minTimeCUTENSORNET : time;
426 }
427
428 if (verbose) printf("Contracted the tensor network, each slice used the same prepared contraction\n");
429
430 // Print the 1-norm of the output tensor (verification)
431 HANDLE_CUDA_ERROR(cudaStreamSynchronize(stream));
432 // restore the output tensor on Host
433 HANDLE_CUDA_ERROR(cudaMemcpy(tensorData_h[numInputs], tensorData_d[numInputs], tensorSizes[numInputs], cudaMemcpyDeviceToHost));
434 double norm1 = 0.0;
435 for (int64_t i = 0; i < tensorElements[numInputs]; ++i)
436 {
437 norm1 += std::abs(tensorData_h[numInputs][i]);
438 }
439 if (verbose) printf("Computed the 1-norm of the output tensor: %e\n", norm1);
440
441 /*************************/
442
443 // Query the total Flop count for the tensor network contraction
444 double flops{0.0};
445 HANDLE_ERROR(cutensornetContractionOptimizerInfoGetAttribute(handle,
446 optimizerInfo,
447 CUTENSORNET_CONTRACTION_OPTIMIZER_INFO_FLOP_COUNT,
448 &flops,
449 sizeof(flops)));
450
451 if (verbose)
452 {
453 printf("Number of tensor network slices = %ld\n", numSlices);
454 printf("Tensor network contraction time (ms) = %.3f\n", minTimeCUTENSORNET * 1000.f);
455 }
456
457 // Sphinx: #9
458 /****************
459 * Free resources
460 ****************/
461
462 // Free cuTensorNet resources
463 HANDLE_ERROR(cutensornetDestroySliceGroup(sliceGroup));
464 HANDLE_ERROR(cutensornetDestroyWorkspaceDescriptor(workDesc));
465 HANDLE_ERROR(cutensornetDestroyContractionOptimizerInfo(optimizerInfo));
466 HANDLE_ERROR(cutensornetDestroyContractionOptimizerConfig(optimizerConfig));
467 HANDLE_ERROR(cutensornetDestroyNetwork(networkDesc));
468 HANDLE_ERROR(cutensornetDestroy(handle));
469
470 HANDLE_CUDA_ERROR(cudaStreamDestroy(stream));
471
472 // Free Host and GPU memory resources
473 for (int32_t t = 0; t < numInputs + 1; ++t)
474 {
475 if (tensorData_h[t]) free(tensorData_h[t]);
476 if (tensorData_d[t]) cudaFree(tensorData_d[t]);
477 }
478 if (work) cudaFree(work);
479
480 if (verbose) printf("Freed resources and exited\n");
481
482 return 0;
483}
Recall that the full sample code can be found in the NVIDIA/cuQuantum repository (here).