You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
GPU recoupling kernels scale poorly for sectors with large recoupling matrices (e.g. SU(3)) #609
The GPU kernels for generic (non-UniqueFusion) tree transformers compute every destination element separately, as a sum over all source fusion trees of the block. This affects generic_batched_permute_kernel! (permute, braid, transpose) and, since #605, trace_recouple_kernel! (planar traces), both in ext/TensorKitGPUArraysExt.jl. With one thread per destination element (i, e) of a block whose recoupling matrix U has size rows × cols, the thread computes
Over the whole block this is a (rows × cols) * (cols × blocksize) matrix product done without any data reuse. Each source element is read again from global memory by all rows destination trees. For sectors whose blocks hold many fusion trees, such as SU(3) (outer multiplicities, large N-j symbols) or many legs of SU(2), rows and cols become large. The kernel then does rows * cols * blocksize scattered reads, while the CPU implementation extracts the subblocks into a buffer and does the recoupling as one BLAS mul!.
Possible directions, to be decided with benchmarks:
above some size of U, gather the source subblocks into a contiguous cols × blocksize buffer, recouple with a (batched) GEMM from the vendor BLAS, and scatter the result, as on the CPU
a tiled kernel that keeps a tile of U and of the sources in shared memory
at least, coalesce the reads of U, which are now strided by rows (see the TODO in generic_batched_permute_kernel!)
A first step would be benchmarks of permutes and planar traces with SU3Irrep (and SU2Irrep with many legs) on CUDA, comparing against the CPU path.
Activity