Skip to content

GPU recoupling kernels scale poorly for sectors with large recoupling matrices (e.g. SU(3)) #609

Open

Description

@lkdvos

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

dst_i[e] = α * sum(U[i, j] * src_j[e] for j in 1:cols) + β * dst_i[e]

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.

Raised in #605 (comment).

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions