Code example (distributed truncated SVD with MPI and NCCL)#
Starting with cuTensorNet v2.14.0, tensor QR and SVD can operate on tensors distributed across multiple
processes and GPUs. This example demonstrates a distributed truncated SVD on the bond-swapping step
of the anisotropic tensor renormalization group (ATRG), where two half-tensors are contracted over
mode i and the result is decomposed:
\[\Theta_{uvaxyb} = \sum_i B_{iuva}\, C_{ixyb}, \qquad
\Theta \approx \sum_g X_{axyg}\, s_g\, Y_{guvb}\]
Unlike Code example (Contraction with automatic distributed slicing), which enables distributed parallelization inside the contraction path finder and executor, this sample performs a collective tensor decomposition on distributed tensor descriptors created with cutensornetCreateDistributedTensorDescriptor(). The full code can be found in the NVIDIA/cuQuantum repository (here).
To run the sample:
Build with MPI support and execute with one GPU per MPI process, e.g.
mpirun -n 4 ./decompose_example_mpi_nccl.Set the environment variable
$CUTENSORNET_COMM_LIBto the path of the MPI distributed interface library (libcutensornet_distributed_interface_mpi.so), as described in the installation guide.NCCL and cuSOLVERMp are required; see distributed tensor decomposition requirements.
Note
Distributed SVD requires CUTENSORNET_TENSOR_SVD_ALGO_GESVDP; GESVD (including the
config default), GESVDJ, and GESVDR return CUTENSORNET_STATUS_NOT_SUPPORTED.
Gate split remains local-only.
Mixed local and distributed operands in a single call return CUTENSORNET_STATUS_INVALID_VALUE.
MPI and GPU setup#
As in the Code example (Contraction with automatic distributed slicing), we initialize MPI, assign one GPU per process, create a cuTensorNet library handle, and bind a duplicated MPI communicator with cutensornetDistributedResetConfiguration(). All subsequent distributed calls are collective over this communicator.
132 /*******************************************
133 * MPI setup: one GPU per process
134 ********************************************/
135
136 HANDLE_MPI_ERROR( MPI_Init(&argc, &argv) );
137
138 int rank{-1}, numRanks{0};
139 HANDLE_MPI_ERROR( MPI_Comm_rank(MPI_COMM_WORLD, &rank) );
140 HANDLE_MPI_ERROR( MPI_Comm_size(MPI_COMM_WORLD, &numRanks) );
141 const bool verbose = (rank == 0);
142
143 // The device must be selected before cutensornetCreate().
144 int numDevices{0};
145 HANDLE_CUDA_ERROR( cudaGetDeviceCount(&numDevices) );
146 HANDLE_CUDA_ERROR( cudaSetDevice(rank % numDevices) );
147
148 if (verbose)
149 printf("cuTensorNet-vers:%ld\n", cutensornetGetVersion());
150
151 cutensornetHandle_t handle;
152 HANDLE_ERROR( cutensornetCreate(&handle) );
153
154 // Bind a duplicated MPI communicator; it must stay alive until the next
155 // reset. All distributed calls below are collective over it.
156 MPI_Comm cutnComm;
157 HANDLE_MPI_ERROR( MPI_Comm_dup(MPI_COMM_WORLD, &cutnComm) );
158 HANDLE_ERROR( cutensornetDistributedResetConfiguration(handle, &cutnComm, sizeof(cutnComm)) );
Build the global tensor on the host#
The sample builds the global tensor Theta on every rank purely to produce input data for
the scatter step below.
161 /*******************************************************
162 * Build Theta = sum_i B[i,u,v,a] C[i,x,y,b] on the host
163 ********************************************************/
164
165 const int64_t D = 7; // spatial degrees of freedom (u, v, x, y); kept extent of g
166 const int64_t CHI = 13; // vertical bond dimension (a, b)
167 const int64_t CHI_T = 6; // temporal bond dimension (i), contracted away below
168
169 std::vector<double> bHalf(CHI_T * D * D * CHI);
170 std::vector<double> cHalf(CHI_T * D * D * CHI);
171 uint64_t rngState = 42;
172 for (auto& value: bHalf) value = pseudoRandom(rngState);
173 for (auto& value: cHalf) value = pseudoRandom(rngState);
174
175 // Theta in Fortran order over modes (u, v, a, x, y, b); B and C in
176 // Fortran order over (i, u, v, a) and (i, x, y, b).
177 const size_t elementsTheta = D * D * CHI * D * D * CHI;
178 std::vector<double> theta(elementsTheta, 0.0);
179 for (int64_t b = 0; b < CHI; ++b)
180 for (int64_t y = 0; y < D; ++y)
181 for (int64_t x = 0; x < D; ++x)
182 for (int64_t a = 0; a < CHI; ++a)
183 for (int64_t v = 0; v < D; ++v)
184 for (int64_t u = 0; u < D; ++u)
185 {
186 double sum = 0.0;
187 for (int64_t i = 0; i < CHI_T; ++i)
188 sum += bHalf[i + CHI_T * (u + D * (v + D * a))]
189 * cHalf[i + CHI_T * (x + D * (y + D * b))];
190 theta[u + D * (v + D * (a + CHI * (x + D * (y + D * b))))] = sum;
191 }
Create distributed tensor descriptors#
Each operand is described by its global extents, a per-mode process grid (nranksPerMode), and
an optional block-cyclic layout. Along each distributed mode, a positive block size assigns
consecutive blocks round-robin across process-grid coordinates, while block size zero selects one
near-even contiguous slab per coordinate. A process-grid extent of one leaves that mode undistributed.
Every rank must create the same descriptors in the same order.
194 /*******************************************************
195 * Create the distributed tensor descriptors (collective)
196 ********************************************************/
197
198 const cudaDataType_t typeData = CUDA_R_64F;
199
200 std::vector<int32_t> modesTheta{'u','v','a','x','y','b'};
201 std::vector<int32_t> modesX{'a','x','y','g'};
202 std::vector<int32_t> modesY{'g','u','v','b'};
203
204 std::vector<int64_t> extentsTheta{D, D, CHI, D, D, CHI};
205 // The shared extent of the outputs at creation time is the truncation
206 // cap: the bond swap keeps g = D of the D*D*CHI available singular
207 // values. This example truncates to that cap; a value-based cutoff
208 // (CUTENSORNET_TENSOR_SVD_CONFIG_ABS_CUTOFF / _REL_CUTOFF /
209 // _DISCARDED_WEIGHT_CUTOFF) would instead let the retained extent be
210 // decided by the spectrum, bounded above by this cap.
211 std::vector<int64_t> extentsX{CHI, D, D, D};
212 std::vector<int64_t> extentsY{D, D, D, CHI};
213
214 // Theta is split over u, X over x, and Y over u. The shared mode g
215 // remains undistributed. Positive block sizes select round-robin blocks;
216 // zero selects slabs. NULL strides select compact Fortran-order shards.
217 const int64_t BS = 2;
218 std::vector<int64_t> gridTheta{numRanks, 1, 1, 1, 1, 1};
219 std::vector<int64_t> gridX{1, numRanks, 1, 1};
220 std::vector<int64_t> gridY{1, numRanks, 1, 1};
221 std::vector<int64_t> blocksTheta{BS, 0, 0, 0, 0, 0};
222 std::vector<int64_t> blocksX{0, BS, 0, 0};
223 std::vector<int64_t> blocksY{0, BS, 0, 0};
224
225 cutensornetTensorDescriptor_t descTheta, descX, descY;
226 // Every rank creates the same descriptors in the same order.
227 HANDLE_ERROR( cutensornetCreateDistributedTensorDescriptor(handle,
228 modesTheta.size(), extentsTheta.data(), NULL, blocksTheta.data(), NULL,
229 gridTheta.data(), modesTheta.data(), typeData, &descTheta) );
230 HANDLE_ERROR( cutensornetCreateDistributedTensorDescriptor(handle,
231 modesX.size(), extentsX.data(), NULL, blocksX.data(), NULL,
232 gridX.data(), modesX.data(), typeData, &descX) );
233 HANDLE_ERROR( cutensornetCreateDistributedTensorDescriptor(handle,
234 modesY.size(), extentsY.data(), NULL, blocksY.data(), NULL,
235 gridY.data(), modesY.data(), typeData, &descY) );
236
237 if (verbose)
238 printf("Created distributed descriptors: Theta split over u across %d rank(s); "
239 "bond swap keeps g = %ld of %ld singular values\n",
240 numRanks, D, D * D * CHI);
Prepare rank-local device buffers#
Rank-local layout details can be queried with cutensornetTensorDescriptorGetAttribute(). This sample scatters from the replicated host tensor built above: each rank packs its owned elements into compact Fortran-order host storage, then copies the buffer to the device. The decomposition APIs only require rank-local device memory in that layout; you can populate those buffers on the device directly, without assembling a global tensor on the host.
243 /*******************************************************
244 * Allocate the local shards and scatter Theta
245 ********************************************************/
246
247 // Rank-local queries return each compact Fortran-order shard's size.
248 size_t thetaLocalBytes{0}, xLocalBytes{0}, yLocalBytes{0};
249 HANDLE_ERROR( cutensornetTensorDescriptorGetAttribute(handle, descTheta,
250 CUTENSORNET_TENSOR_DESCRIPTOR_LOCAL_DATA_SIZE,
251 &thetaLocalBytes, sizeof(thetaLocalBytes)) );
252 HANDLE_ERROR( cutensornetTensorDescriptorGetAttribute(handle, descX,
253 CUTENSORNET_TENSOR_DESCRIPTOR_LOCAL_DATA_SIZE,
254 &xLocalBytes, sizeof(xLocalBytes)) );
255 HANDLE_ERROR( cutensornetTensorDescriptorGetAttribute(handle, descY,
256 CUTENSORNET_TENSOR_DESCRIPTOR_LOCAL_DATA_SIZE,
257 &yLocalBytes, sizeof(yLocalBytes)) );
258
259 // This rank's shard of Theta: the block-cyclically owned runs of the u
260 // mode (with BS=2 and two ranks, mode u of extent 7 has blocks [0:2)
261 // [2:4) [4:6) [6:7); rank 0 owns blocks 0 and 2, rank 1 owns 1 and 3),
262 // packed consecutively in Fortran order over the local extents.
263 const std::vector<OwnedRun> uRuns = ownedRuns(D, BS, numRanks, rank);
264 const int64_t uCount = ownedCount(uRuns);
265 std::vector<double> thetaLocal(thetaLocalBytes / sizeof(double));
266 for (int64_t b = 0; b < CHI; ++b)
267 for (int64_t y = 0; y < D; ++y)
268 for (int64_t x = 0; x < D; ++x)
269 for (int64_t a = 0; a < CHI; ++a)
270 for (int64_t v = 0; v < D; ++v)
271 for (const auto& run: uRuns)
272 for (int64_t o = 0; o < run.count; ++o)
273 {
274 const int64_t u = run.globalStart + o;
275 const int64_t uLoc = run.localStart + o;
276 thetaLocal[uLoc + uCount * (v + D * (a + CHI * (x + D * (y + D * b))))] =
277 theta[u + D * (v + D * (a + CHI * (x + D * (y + D * b))))];
278 }
279
280 void *dTheta{nullptr}, *dX{nullptr}, *dS{nullptr}, *dY{nullptr};
281 const size_t sBytes = D * sizeof(double); // replicated on every rank
282 if (thetaLocalBytes > 0)
283 HANDLE_CUDA_ERROR( cudaMalloc(&dTheta, thetaLocalBytes) );
284 if (xLocalBytes > 0)
285 HANDLE_CUDA_ERROR( cudaMalloc(&dX, xLocalBytes) );
286 HANDLE_CUDA_ERROR( cudaMalloc(&dS, sBytes) );
287 if (yLocalBytes > 0)
288 HANDLE_CUDA_ERROR( cudaMalloc(&dY, yLocalBytes) );
289 if (thetaLocalBytes > 0)
290 HANDLE_CUDA_ERROR( cudaMemcpy(dTheta, thetaLocal.data(), thetaLocalBytes,
291 cudaMemcpyHostToDevice) );
292
293 // Initialize the output buffers before the SVD overwrites them.
294 {
295 std::vector<double> init(std::max(xLocalBytes, yLocalBytes) / sizeof(double));
296 for (auto& value: init) value = pseudoRandom(rngState);
297 if (xLocalBytes > 0)
298 HANDLE_CUDA_ERROR( cudaMemcpy(dX, init.data(), xLocalBytes,
299 cudaMemcpyHostToDevice) );
300 if (yLocalBytes > 0)
301 HANDLE_CUDA_ERROR( cudaMemcpy(dY, init.data(), yLocalBytes,
302 cudaMemcpyHostToDevice) );
303 }
SVD configuration and workspace query#
The workspace and execution workflow matches the local SVD example: create an SVD config and info object, query workspace sizes collectively, allocate scratch memory, and call cutensornetTensorSVD(). Fixed-extent truncation is specified by the shared extent of the output descriptors at creation time.
306 /*******************************************************
307 * SVD config/info and workspace (same calls as local)
308 ********************************************************/
309
310 // No cutoffs are set here, so the SVD truncates to the cap the output
311 // descriptors were created with. Distributed SVD requires GESVDP; the
312 // config default (GESVD) returns CUTENSORNET_STATUS_NOT_SUPPORTED.
313 cutensornetTensorSVDConfig_t svdConfig;
314 HANDLE_ERROR( cutensornetCreateTensorSVDConfig(handle, &svdConfig) );
315 const cutensornetTensorSVDAlgo_t svdAlgo = CUTENSORNET_TENSOR_SVD_ALGO_GESVDP;
316 HANDLE_ERROR( cutensornetTensorSVDConfigSetAttribute(handle, svdConfig,
317 CUTENSORNET_TENSOR_SVD_CONFIG_ALGO, &svdAlgo, sizeof(svdAlgo)) );
318 cutensornetTensorSVDInfo_t svdInfo;
319 HANDLE_ERROR( cutensornetCreateTensorSVDInfo(handle, &svdInfo) );
320
321 cutensornetWorkspaceDescriptor_t workDesc;
322 HANDLE_ERROR( cutensornetCreateWorkspaceDescriptor(handle, &workDesc) );
323 // Collective: every rank queries with the same descriptor trio.
324 HANDLE_ERROR( cutensornetWorkspaceComputeSVDSizes(handle, descTheta, descX, descY,
325 svdConfig, workDesc) );
326 int64_t deviceWorkspaceSize{0}, hostWorkspaceSize{0};
327 HANDLE_ERROR( cutensornetWorkspaceGetMemorySize(handle, workDesc,
328 CUTENSORNET_WORKSIZE_PREF_MIN, CUTENSORNET_MEMSPACE_DEVICE,
329 CUTENSORNET_WORKSPACE_SCRATCH, &deviceWorkspaceSize) );
330 HANDLE_ERROR( cutensornetWorkspaceGetMemorySize(handle, workDesc,
331 CUTENSORNET_WORKSIZE_PREF_MIN, CUTENSORNET_MEMSPACE_HOST,
332 CUTENSORNET_WORKSPACE_SCRATCH, &hostWorkspaceSize) );
333
334 void* devWork{nullptr};
335 void* hostWork{nullptr};
336 if (deviceWorkspaceSize > 0)
337 HANDLE_CUDA_ERROR( cudaMalloc(&devWork, deviceWorkspaceSize) );
338 if (hostWorkspaceSize > 0)
339 hostWork = malloc(hostWorkspaceSize);
340 HANDLE_ERROR( cutensornetWorkspaceSetMemory(handle, workDesc,
341 CUTENSORNET_MEMSPACE_DEVICE, CUTENSORNET_WORKSPACE_SCRATCH,
342 devWork, deviceWorkspaceSize) );
343 HANDLE_ERROR( cutensornetWorkspaceSetMemory(handle, workDesc,
344 CUTENSORNET_MEMSPACE_HOST, CUTENSORNET_WORKSPACE_SCRATCH,
345 hostWork, hostWorkspaceSize) );
Execution#
cutensornetTensorSVD() updates the output descriptors to the realized reduced extent after truncation. Recreate them before a later call if truncation can change that extent. Singular values are replicated on every rank.
348 /*******************************************************
349 * Execution: the distributed bond-swap SVD
350 ********************************************************/
351
352 cudaStream_t stream;
353 HANDLE_CUDA_ERROR( cudaStreamCreate(&stream) );
354
355 // Fixed-extent truncation: Kred == Kcap. After value-based truncation,
356 // descX/descY describe Kred; recreate them at Kcap before calling again.
357 HANDLE_ERROR( cutensornetTensorSVD(handle,
358 descTheta, dTheta,
359 descX, dX,
360 dS,
361 descY, dY,
362 svdConfig, svdInfo, workDesc, stream) );
363
364 // Outputs become visible to work enqueued on `stream` after the call.
365 std::vector<double> xLocal(xLocalBytes / sizeof(double));
366 std::vector<double> sHost(D);
367 std::vector<double> yLocal(yLocalBytes / sizeof(double));
368 if (xLocalBytes > 0)
369 HANDLE_CUDA_ERROR( cudaMemcpyAsync(xLocal.data(), dX, xLocalBytes,
370 cudaMemcpyDeviceToHost, stream) );
371 HANDLE_CUDA_ERROR( cudaMemcpyAsync(sHost.data(), dS, sBytes,
372 cudaMemcpyDeviceToHost, stream) );
373 if (yLocalBytes > 0)
374 HANDLE_CUDA_ERROR( cudaMemcpyAsync(yLocal.data(), dY, yLocalBytes,
375 cudaMemcpyDeviceToHost, stream) );
376 HANDLE_CUDA_ERROR( cudaStreamSynchronize(stream) );
377
378 int64_t reducedExtent{0};
379 double discardedWeight{0.0};
380 HANDLE_ERROR( cutensornetTensorSVDInfoGetAttribute(handle, svdInfo,
381 CUTENSORNET_TENSOR_SVD_INFO_REDUCED_EXTENT,
382 &reducedExtent, sizeof(reducedExtent)) );
383 HANDLE_ERROR( cutensornetTensorSVDInfoGetAttribute(handle, svdInfo,
384 CUTENSORNET_TENSOR_SVD_INFO_DISCARDED_WEIGHT,
385 &discardedWeight, sizeof(discardedWeight)) );
Verification#
The sample reassembles the distributed factors on the host and checks that the reconstruction residual matches the reported discarded weight, and that singular values are identical on every rank.
388 /*******************************************************
389 * Verification
390 ********************************************************/
391
392 // The singular values are replicated: every rank must hold the values
393 // rank 0 holds.
394 std::vector<double> sRoot(sHost);
395 HANDLE_MPI_ERROR( MPI_Bcast(sRoot.data(), D, MPI_DOUBLE, 0, MPI_COMM_WORLD) );
396 double sReplicationError = 0.0;
397 for (int64_t g = 0; g < D; ++g)
398 sReplicationError = std::max(sReplicationError, std::fabs(sHost[g] - sRoot[g]));
399
400 // Reassemble X and Y on every rank (zero-fill + allreduce over the
401 // block-cyclically owned entries), then reconstruct the truncated
402 // Theta. X is split over its x mode, Y over its u mode, both with the
403 // same block size, so both reuse the u-mode ownership runs.
404 const size_t elementsX = CHI * D * D * D;
405 const size_t elementsY = D * D * D * CHI;
406 std::vector<double> xFull(elementsX, 0.0);
407 std::vector<double> yFull(elementsY, 0.0);
408 for (int64_t g = 0; g < D; ++g)
409 for (int64_t y = 0; y < D; ++y)
410 for (const auto& run: uRuns)
411 for (int64_t o = 0; o < run.count; ++o)
412 for (int64_t a = 0; a < CHI; ++a)
413 {
414 const int64_t x = run.globalStart + o;
415 const int64_t xLoc = run.localStart + o;
416 xFull[a + CHI * (x + D * (y + D * g))] =
417 xLocal[a + CHI * (xLoc + uCount * (y + D * g))];
418 }
419 for (int64_t b = 0; b < CHI; ++b)
420 for (int64_t v = 0; v < D; ++v)
421 for (const auto& run: uRuns)
422 for (int64_t o = 0; o < run.count; ++o)
423 for (int64_t g = 0; g < D; ++g)
424 {
425 const int64_t u = run.globalStart + o;
426 const int64_t uLoc = run.localStart + o;
427 yFull[g + D * (u + D * (v + D * b))] =
428 yLocal[g + D * (uLoc + uCount * (v + D * b))];
429 }
430 HANDLE_MPI_ERROR( MPI_Allreduce(MPI_IN_PLACE, xFull.data(), elementsX,
431 MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD) );
432 HANDLE_MPI_ERROR( MPI_Allreduce(MPI_IN_PLACE, yFull.data(), elementsY,
433 MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD) );
434
435 // The relative squared residual of the reconstruction must equal the
436 // reported discarded weight: both are the weight of the singular values
437 // beyond the cap.
438 double residual2 = 0.0, thetaNorm2 = 0.0;
439 for (int64_t b = 0; b < CHI; ++b)
440 for (int64_t y = 0; y < D; ++y)
441 for (int64_t x = 0; x < D; ++x)
442 for (int64_t a = 0; a < CHI; ++a)
443 for (int64_t v = 0; v < D; ++v)
444 for (int64_t u = 0; u < D; ++u)
445 {
446 double rec = 0.0;
447 for (int64_t g = 0; g < D; ++g)
448 rec += xFull[a + CHI * (x + D * (y + D * g))] * sHost[g]
449 * yFull[g + D * (u + D * (v + D * b))];
450 const double ref = theta[u + D * (v + D * (a + CHI * (x + D * (y + D * b))))];
451 residual2 += (rec - ref) * (rec - ref);
452 thetaNorm2 += ref * ref;
453 }
454 const double relativeResidual2 = residual2 / thetaNorm2;
455
456 if (verbose)
457 {
458 printf("reduced extent: %ld (cap %ld)\n", reducedExtent, D);
459 printf("leading singular values:");
460 for (int64_t g = 0; g < std::min<int64_t>(D, 4); ++g)
461 printf(" %.6f", sHost[g]);
462 printf(" ...\n");
463 printf("truncation: relative residual^2 = %.6f, reported discarded weight = %.6f\n",
464 relativeResidual2, discardedWeight);
465 }
466 if (sReplicationError > 1e-12)
467 {
468 printf("Error: singular values differ across ranks (%e)\n", sReplicationError);
469 MPI_Abort(MPI_COMM_WORLD, 1);
470 }
471 if (std::fabs(relativeResidual2 - discardedWeight) > 1e-10)
472 {
473 printf("Error: reconstruction inconsistent with discarded weight\n");
474 MPI_Abort(MPI_COMM_WORLD, 1);
475 }
476 if (verbose)
477 printf("Distributed bond-swap SVD verified.\n");
Free resources#
After the computation, free CUDA, library, and MPI resources.
480 /*******************************************************
481 * Free resources
482 ********************************************************/
483
484 HANDLE_CUDA_ERROR( cudaStreamDestroy(stream) );
485 HANDLE_ERROR( cutensornetDestroyTensorDescriptor(descTheta) );
486 HANDLE_ERROR( cutensornetDestroyTensorDescriptor(descX) );
487 HANDLE_ERROR( cutensornetDestroyTensorDescriptor(descY) );
488 HANDLE_ERROR( cutensornetDestroyTensorSVDConfig(svdConfig) );
489 HANDLE_ERROR( cutensornetDestroyTensorSVDInfo(svdInfo) );
490 HANDLE_ERROR( cutensornetDestroyWorkspaceDescriptor(workDesc) );
491 HANDLE_ERROR( cutensornetDestroy(handle) );
492
493 if (dTheta) HANDLE_CUDA_ERROR( cudaFree(dTheta) );
494 if (dX) HANDLE_CUDA_ERROR( cudaFree(dX) );
495 HANDLE_CUDA_ERROR( cudaFree(dS) );
496 if (dY) HANDLE_CUDA_ERROR( cudaFree(dY) );
497 if (devWork) HANDLE_CUDA_ERROR( cudaFree(devWork) );
498 if (hostWork) free(hostWork);
499
500 HANDLE_MPI_ERROR( MPI_Comm_free(&cutnComm) );
501 HANDLE_MPI_ERROR( MPI_Finalize() );
502
503 if (verbose)
504 printf("Free resource and exit.\n");
505
506 return 0;
507}