使用 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 的张量描述符。
计算完成后,我们仍然需要释放所有资源。