使用 cuTensorNet 执行张量 SVD 采用了与 QR 示例 非常相似的工作流程。在这里,我们重点介绍两个 API 之间值得注意的差异。完整代码可以在 NVIDIA/cuQuantum 仓库中找到 (此处)。

定义 SVD 分解

与 QR 分解一样,我们首先定义要执行的 SVD 分解,包括数据类型、模式划分和范围。

 97   /******************************************************
 98   * Tensor SVD: T_{i,j,m,n} -> U_{i,x,m} S_{x} V_{n,x,j}  
 99   *******************************************************/
100
101   typedef float floatType;
102   cudaDataType_t typeData = CUDA_R_32F;
103
104   // Create vector of modes
105   int32_t sharedMode = 'x';
106
107   std::vector<int32_t> modesT{'i','j','m','n'}; // input
108   std::vector<int32_t> modesU{'i', sharedMode,'m'};
109   std::vector<int32_t> modesV{'n', sharedMode,'j'};  // SVD output
110
111   // Extents
112   std::unordered_map<int32_t, int64_t> extentMap;
113   extentMap['i'] = 16;
114   extentMap['j'] = 16;
115   extentMap['m'] = 16;
116   extentMap['n'] = 16;
117
118   int64_t rowExtent = computeCombinedExtent(extentMap, modesU);
119   int64_t colExtent = computeCombinedExtent(extentMap, modesV);
120   // cuTensorNet tensor SVD operates in reduced mode expecting k <= min(m, n)
121   int64_t fullSharedExtent = rowExtent <= colExtent? rowExtent: colExtent;
122   const int64_t maxExtent = fullSharedExtent / 2;  //fix extent truncation with half of the singular values trimmed out
123   extentMap[sharedMode] = maxExtent;
124
125   // Create a vector of extents for each tensor
126   std::vector<int64_t> extentT;
127   for (auto mode : modesT)
128      extentT.push_back(extentMap[mode]);
129   std::vector<int64_t> extentU;
130   for (auto mode : modesU)
131      extentU.push_back(extentMap[mode]);
132   std::vector<int64_t> extentV;
133   for (auto mode : modesV)
134      extentV.push_back(extentMap[mode]);

注意

要执行固定范围截断,我们直接将 maxExtent 设置为与精确 SVD 对应的完整范围的一半。

设置 SVD 截断参数

定义 SVD 分解后,我们可以按照与 QR 示例 相同的工作流程进行数据分配和张量描述符初始化。在查询工作区之前,我们可以在 cutensornetTensorSVDConfig_t 中选择不同的 SVD 选项。同时,我们可以创建 cutensornetTensorSVDInfo_t 以跟踪运行时截断信息。

221   /**********************************************
222   * Setup SVD algorithm and truncation parameters
223   ***********************************************/
224
225   cutensornetTensorSVDConfig_t svdConfig;
226   HANDLE_ERROR( cutensornetCreateTensorSVDConfig(handle, &svdConfig) );
227
228   // set up truncation parameters
229   double absCutoff = 1e-2;
230   HANDLE_ERROR( cutensornetTensorSVDConfigSetAttribute(handle, 
231                                          svdConfig, 
232                                          CUTENSORNET_TENSOR_SVD_CONFIG_ABS_CUTOFF, 
233                                          &absCutoff, 
234                                          sizeof(absCutoff)) );
235   double relCutoff = 4e-2;
236   HANDLE_ERROR( cutensornetTensorSVDConfigSetAttribute(handle, 
237                                          svdConfig, 
238                                          CUTENSORNET_TENSOR_SVD_CONFIG_REL_CUTOFF, 
239                                          &relCutoff, 
240                                          sizeof(relCutoff)) );
241   
242   // optional: choose gesvdj algorithm with customized parameters. Default is gesvd.
243   cutensornetTensorSVDAlgo_t svdAlgo = CUTENSORNET_TENSOR_SVD_ALGO_GESVDJ;
244   HANDLE_ERROR( cutensornetTensorSVDConfigSetAttribute(handle, 
245                                          svdConfig, 
246                                          CUTENSORNET_TENSOR_SVD_CONFIG_ALGO, 
247                                          &svdAlgo, 
248                                          sizeof(svdAlgo)) );
249   cutensornetGesvdjParams_t gesvdjParams{/*tol=*/1e-12, /*maxSweeps=*/80};
250   HANDLE_ERROR( cutensornetTensorSVDConfigSetAttribute(handle, 
251                                          svdConfig, 
252                                          CUTENSORNET_TENSOR_SVD_CONFIG_ALGO_PARAMS, 
253                                          &gesvdjParams, 
254                                          sizeof(gesvdjParams)) );
255   printf("Set up SVDConfig to use GESVDJ algorithm with truncation\n");
256   
257   /********************************************************
258   * Create SVDInfo to record runtime SVD truncation details
259   *********************************************************/
260
261   cutensornetTensorSVDInfo_t svdInfo; 
262   HANDLE_ERROR( cutensornetCreateTensorSVDInfo(handle, &svdInfo)) ;

执行

接下来,我们可以使用 cutensornetWorkspaceComputeSVDSizes() 查询和分配工作区,这与其 QR 对应部分非常相似。在此阶段,我们可以通过调用 cutensornetTensorSVD() 执行 SVD 分解。

308   /**********
309   * Execution
310   ***********/
311  
312   GPUTimer timer{stream};
313   double minTimeCUTENSOR = 1e100;
314   const int numRuns = 3; // to get stable perf results
315   for (int i=0; i < numRuns; ++i)
316   {  
317      // restore output
318      cudaMemsetAsync(D_U, 0, sizeU, stream);
319      cudaMemsetAsync(D_S, 0, sizeS, stream);
320      cudaMemsetAsync(D_V, 0, sizeV, stream);
321      cudaDeviceSynchronize();
322      
323      // With value-based truncation, `cutensornetTensorSVD` can potentially update the shared extent in descTensorU/V.
324      // We here restore descTensorU/V to the original problem.
325      HANDLE_ERROR( cutensornetDestroyTensorDescriptor(descTensorU) );
326      HANDLE_ERROR( cutensornetDestroyTensorDescriptor(descTensorV) );
327      HANDLE_ERROR( cutensornetCreateTensorDescriptor(handle, numModesU, extentU.data(), strides, modesU.data(), typeData, &descTensorU) );
328      HANDLE_ERROR( cutensornetCreateTensorDescriptor(handle, numModesV, extentV.data(), strides, modesV.data(), typeData, &descTensorV) );
329
330      timer.start();
331      HANDLE_ERROR( cutensornetTensorSVD(handle, 
332                        descTensorIn, D_T, 
333                        descTensorU, D_U, 
334                        D_S, 
335                        descTensorV, D_V, 
336                        svdConfig, 
337                        svdInfo,
338                        workDesc,
339                        stream) );
340      // Synchronize and measure timing
341      auto time = timer.seconds();
342      minTimeCUTENSOR = (minTimeCUTENSOR < time) ? minTimeCUTENSOR : time;
343   }
344
345   printf("Performing SVD\n");
346
347   HANDLE_CUDA_ERROR( cudaMemcpyAsync(U, D_U, sizeU, cudaMemcpyDeviceToHost) );
348   HANDLE_CUDA_ERROR( cudaMemcpyAsync(S, D_S, sizeS, cudaMemcpyDeviceToHost) );
349   HANDLE_CUDA_ERROR( cudaMemcpyAsync(V, D_V, sizeV, cudaMemcpyDeviceToHost) );

注意

由于我们在本示例中启用了加权截断选项,如果我们希望多次执行相同的计算,则需要恢复 U 和 V 的张量描述符。

计算完成后,我们仍然需要释放所有资源。