From f5ef424daa0720d5fadd0b4ab268fc0fd5dda4f4 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Thu, 5 Nov 2026 09:53:27 -0500 Subject: [PATCH] Grouped gemm (and 1m for complex) in `mul!` for CUDA tensors Small blocks are multiplied in a single `cublasGemmGroupedBatchedEx` call, using the real 1m embedding for complex element types. Large blocks use individual `cublasGemmEx` calls with the scalars uploaded once. Co-Authored-By: Claude Opus 5.5 --- docs/src/Changelog.md | 2 + ext/TensorKitCUDAExt/TensorKitCUDAExt.jl | 11 ++ ext/TensorKitCUDAExt/mul.jl | 226 +++++++++++++++++++++++ test/cuda/tensors.jl | 18 ++ 4 files changed, 257 insertions(+) create mode 100644 ext/TensorKitCUDAExt/mul.jl diff --git a/docs/src/Changelog.md b/docs/src/Changelog.md index bc2dc34ad..fdc2fee14 100644 --- a/docs/src/Changelog.md +++ b/docs/src/Changelog.md @@ -36,6 +36,8 @@ When releasing a new version, move the "Unreleased" changes to a new version sec ### Performance +- `mul!` for CUDA tensors multiplies all small blocks in a single grouped gemm call (cuBLAS ≥ 12.5), using a real 1m embedding for complex element types; large blocks (`m * n * k` above `TensorKitCUDAExt.GROUPED_GEMM_THRESHOLD[]`, separately for real and complex element types) use individual gemm calls without per-call scalar uploads + ## [0.17.3](https://github.com/QuantumKitHub/TensorKit.jl/compare/v0.17.2...v0.17.3) - 2026-10-30 ### Added diff --git a/ext/TensorKitCUDAExt/TensorKitCUDAExt.jl b/ext/TensorKitCUDAExt/TensorKitCUDAExt.jl index a54d2490d..d49731299 100644 --- a/ext/TensorKitCUDAExt/TensorKitCUDAExt.jl +++ b/ext/TensorKitCUDAExt/TensorKitCUDAExt.jl @@ -11,6 +11,7 @@ using TensorKit.Factorizations using TensorKit.Strided using TensorKit.Factorizations: AbstractAlgorithm using TensorKit: SectorDict, tensormaptype, scalar, similarstoragetype, AdjointTensorMap, scalartype, project_symmetric_and_check +using TensorKit: @cached, TensorMapSpace, blockstructure, compose import TensorKit: randisometry, rand, randn, fill_braidingsubblock! using TensorKit: MatrixAlgebraKit @@ -18,5 +19,15 @@ using TensorKit: MatrixAlgebraKit using Random include("cutensormap.jl") +include("mul.jl") + +function __init__() + # extensions can't modify GLOBAL_CACHES during precompilation, so register explicitly + push!( + TensorKit.GLOBAL_CACHES, + Symbol("TensorKitCUDAExt.GLOBAL_GROUPEDGEMMPLAN_CACHE") => GLOBAL_GROUPEDGEMMPLAN_CACHE + ) + return nothing +end end diff --git a/ext/TensorKitCUDAExt/mul.jl b/ext/TensorKitCUDAExt/mul.jl new file mode 100644 index 000000000..2d2851ce5 --- /dev/null +++ b/ext/TensorKitCUDAExt/mul.jl @@ -0,0 +1,226 @@ +const CuStoredTensorMap{T} = TensorMap{T, S, N₁, N₂, A} where {S, N₁, N₂, A <: CuVector{T}} +const CuMulOperand{T} = Union{ + CuStoredTensorMap{T}, + AdjointTensorMap{T, S, N₁, N₂, TT} where {S, N₁, N₂, TT <: CuStoredTensorMap{T}}, +} + +""" + GROUPED_GEMM_THRESHOLD + +Blocks with `m * n * k` at least `GROUPED_GEMM_THRESHOLD[].real` (real element types) or +`GROUPED_GEMM_THRESHOLD[].complex` (complex element types) are multiplied with individual gemm +calls in `mul!` for CUDA tensors, all smaller blocks are combined into a single grouped gemm call. +Complex blocks benefit from individual calls much earlier, since the grouped call works on their +real 1m embedding. +""" +const GROUPED_GEMM_THRESHOLD = Ref((real = 1024^3, complex = 256^3)) + +const _HAS_GROUPED_GEMM = Ref{Union{Nothing, Bool}}(nothing) +function has_grouped_gemm() + isnothing(_HAS_GROUPED_GEMM[]) && (_HAS_GROUPED_GEMM[] = cuBLAS.version() >= v"12.5") + return _HAS_GROUPED_GEMM[]::Bool +end + +_parent_adj(t::TensorMap) = (t, false) +_parent_adj(t::AdjointTensorMap) = (parent(t), true) + +""" + GroupedGemmPlan + +Everything about `mul!(C, op(A), op(B))` for CUDA tensors that only depends on the spaces, with +`op` either `identity` or `adjoint`: + +- `dims` (`m, n, k`) and `offsets` (into the data of `A`, `B` and `C`), one row per multiplied + block, where `op(A)` is `m × k` and `op(B)` is `k × n`. The first `ngrouped` rows go into a + single grouped gemm call, the remaining ones are large blocks for individual gemm calls. +- `emptyC`: ranges of blocks of `C` without a matching block in `A` or `B`, which are only scaled. +- `transa`, `transb` and `gdims`: the arguments of the grouped gemm call. For complex element + types, this call acts on real 1m embeddings of the blocks: `op(A)` is expanded into a buffer of + length `lenAR`, using the per-block columns `offA, m, k, offAR` of `meta1m`. + +Plans are cached by `groupedgemmplan`, since building one requires the block structures of three +spaces and allocates many small arrays, which would otherwise dominate for tensors with many +small blocks. +""" +struct GroupedGemmPlan + ngrouped::Int + dims::Matrix{Int64} + offsets::Matrix{Int64} + emptyC::Vector{UnitRange{Int}} + transa::cuBLAS.cublasOperation_t + transb::cuBLAS.cublasOperation_t + gdims::Matrix{Int64} + meta1m::Matrix{Int64} + lenAR::Int +end + +@cached function groupedgemmplan( + VA::TensorMapSpace, VB::TensorMapSpace, adjA::Bool, adjB::Bool, cplx::Bool, threshold::Int + )::GroupedGemmPlan + VC = compose(adjA ? VA' : VA, adjB ? VB' : VB) + bsA, bsB, bsC = blockstructure(VA), blockstructure(VB), blockstructure(VC) + blocks = NTuple{6, Int}[] + emptyC = UnitRange{Int}[] + for (c, (_, rC)) in pairs(bsC) + if !(haskey(bsA, c) && haskey(bsB, c)) + push!(emptyC, rC) + continue + end + (dA, rA), (dB, rB) = bsA[c], bsB[c] + m, k = adjA ? reverse(dA) : dA + n = adjB ? first(dB) : last(dB) + push!(blocks, (m, n, k, first(rA) - 1, first(rB) - 1, first(rC) - 1)) + end + isdirect(b) = b[1] * b[2] * b[3] >= threshold + sort!(blocks; by = isdirect) + G = count(!isdirect, blocks) + table = [b[j] for b in blocks, j in 1:6] + dims, offsets = table[:, 1:3], table[:, 4:6] + + gdims = dims[1:G, :] + if cplx # 1m: op(A) is expanded explicitly, B and C are reinterpreted + gdims[:, 1] .*= 2 + gdims[:, 3] .*= 2 + adjA = adjB = false + end + lengths = 4 .* dims[1:G, 1] .* dims[1:G, 3] + offAR = cumsum(lengths) .- lengths + meta1m = cplx ? [offsets[1:G, 1] dims[1:G, 1] dims[1:G, 3] offAR] : zeros(Int64, 0, 4) + return GroupedGemmPlan( + G, dims, offsets, emptyC, + adjA ? cuBLAS.CUBLAS_OP_T : cuBLAS.CUBLAS_OP_N, + adjB ? cuBLAS.CUBLAS_OP_T : cuBLAS.CUBLAS_OP_N, + gdims, meta1m, sum(lengths; init = 0) + ) +end + +function LinearAlgebra.mul!( + tC::CuStoredTensorMap{T}, tA::CuMulOperand{T}, tB::CuMulOperand{T}, α::Number, β::Number + ) where {T <: LinearAlgebra.BlasFloat} + compose(space(tA), space(tB)) == space(tC) || + throw(SpaceMismatch(lazy"$(space(tC)) ≠ $(space(tA)) * $(space(tB))")) + pA, adjA = _parent_adj(tA) + pB, adjB = _parent_adj(tB) + α′, β′ = T(TensorKit._blasscalar(α)), T(TensorKit._blasscalar(β)) + + # the 1m embedding does not support a conjugated B + threshold = !has_grouped_gemm() || (T <: Complex && adjB) ? 0 : + T <: Complex ? GROUPED_GEMM_THRESHOLD[].complex : GROUPED_GEMM_THRESHOLD[].real + plan = groupedgemmplan(space(pA), space(pB), adjA, adjB, T <: Complex, threshold) + + # the grouped 1m gemm is real, so complex β is applied upfront + if !isreal(β′) + rmul!(tC.data, β′) + β′ = one(T) + else + for r in plan.emptyC + v = view(tC.data, r) + iszero(β′) ? fill!(v, zero(T)) : isone(β′) || rmul!(v, β′) + end + end + + _direct_gemm!(tC.data, pA.data, pB.data, plan, adjA, adjB, α′, β′) + _grouped_gemm!(tC.data, pA.data, pB.data, plan, adjA, α′, β′) + return tC +end + +function _direct_gemm!( + C::CuVector{T}, A::CuVector{T}, B::CuVector{T}, plan::GroupedGemmPlan, + adjA::Bool, adjB::Bool, α::T, β::T + ) where {T} + plan.ngrouped == size(plan.dims, 1) && return C + # upload scalars once instead of on every gemm call + dα, dβ = CuRef{T}(α), CuRef{T}(β) + adjop = T <: Complex ? cuBLAS.CUBLAS_OP_C : cuBLAS.CUBLAS_OP_T + opA = adjA ? adjop : cuBLAS.CUBLAS_OP_N + opB = adjB ? adjop : cuBLAS.CUBLAS_OP_N + dt = convert(cuBLAS.cudaDataType_t, T) + ct = cuBLAS.gemmExComputeType(T, T, T, 1, 1, 1) + h = cuBLAS.handle() + for i in (plan.ngrouped + 1):size(plan.dims, 1) + m, n, k = view(plan.dims, i, :) + offA, offB, offC = view(plan.offsets, i, :) + cuBLAS.cublasGemmEx_64( + h, opA, opB, m, n, k, dα, + pointer(A, offA + 1), dt, adjA ? k : m, + pointer(B, offB + 1), dt, adjB ? n : k, + dβ, pointer(C, offC + 1), dt, m, ct, cuBLAS.CUBLAS_GEMM_DEFAULT + ) + end + return C +end + +# reduced precision compute types (e.g. TF32) crash on H100 with misaligned blocks +_grouped_computetype(::Type{Float32}) = cuBLAS.CUBLAS_COMPUTE_32F +_grouped_computetype(::Type{Float64}) = cuBLAS.CUBLAS_COMPUTE_64F + +function _grouped_gemm!( + C::CuVector{T}, A::CuVector{T}, B::CuVector{T}, plan::GroupedGemmPlan, + adjA::Bool, α::T, β::T + ) where {T} + G = plan.ngrouped + G == 0 && return C + R = real(T) + AR = T <: Complex ? CuVector{R}(undef, plan.lenAR) : A + offAR = T <: Complex ? view(plan.meta1m, :, 4) : view(plan.offsets, :, 1) + # device pointers to the blocks of A (or AR), B and C, followed by the 1m metadata + ptrs = Matrix{Int64}(undef, G, 3 + size(plan.meta1m, 2)) + for i in 1:G + ptrs[i, 1] = UInt(pointer(AR, offAR[i] + 1)) + ptrs[i, 2] = UInt(pointer(B, plan.offsets[i, 2] + 1)) + ptrs[i, 3] = UInt(pointer(C, plan.offsets[i, 3] + 1)) + end + T <: Complex && (ptrs[:, 4:end] = plan.meta1m) + dptrs = CuArray(ptrs) + if T <: Complex + maxlen = maximum(i -> plan.dims[i, 1] * plan.dims[i, 3], 1:G) + blocks = (min(cld(maxlen, 256), 1024), min(G, 65535)) + @cuda threads = 256 blocks = blocks _expand1m_kernel!(AR, A, dptrs, G, α, adjA) + α = one(T) + end + + base = UInt(pointer(dptrs)) + gm, gn, gk = view(plan.gdims, :, 1), view(plan.gdims, :, 2), view(plan.gdims, :, 3) + lda = plan.transa == cuBLAS.CUBLAS_OP_T ? gk : gm + ldb = plan.transb == cuBLAS.CUBLAS_OP_T ? gn : gk + dt = convert(cuBLAS.cudaDataType_t, R) + h = cuBLAS.handle() + # grouped gemm only supports host pointers for the scalars + cuBLAS.cublasSetPointerMode_v2(h, cuBLAS.CUBLAS_POINTER_MODE_HOST) + try + cuBLAS.cublasGemmGroupedBatchedEx_64( + h, fill(plan.transa, G), fill(plan.transb, G), gm, gn, gk, fill(real(α), G), + Ptr{Ptr{Cvoid}}(base), dt, lda, + Ptr{Ptr{Cvoid}}(base + 8G), dt, ldb, + fill(real(β), G), Ptr{Ptr{Cvoid}}(base + 16G), dt, gm, + G, ones(Int64, G), _grouped_computetype(R) + ) + finally + cuBLAS.cublasSetPointerMode_v2(h, cuBLAS.CUBLAS_POINTER_MODE_DEVICE) + end + CUDA.unsafe_free!(dptrs) + T <: Complex && CUDA.unsafe_free!(AR) + return C +end + +# writes the 1m embedding [ar -ai; ai ar] of α * op(A) for every grouped block +function _expand1m_kernel!(AR, A, meta, G, α, adjA) + b = blockIdx().y + while b <= G + @inbounds offA, m, k, offAR = meta[b, 4], meta[b, 5], meta[b, 6], meta[b, 7] + idx = (blockIdx().x - 1) * blockDim().x + threadIdx().x - 1 + while idx < m * k + r, s = idx % m, idx ÷ m + @inbounds a = adjA ? conj(A[offA + s + r * k + 1]) : A[offA + idx + 1] + re, im = reim(α * a) + p = offAR + 2r + 4m * s + 1 + @inbounds AR[p] = re + @inbounds AR[p + 1] = im + @inbounds AR[p + 2m] = -im + @inbounds AR[p + 2m + 1] = re + idx += gridDim().x * blockDim().x + end + b += gridDim().y + end + return nothing +end diff --git a/test/cuda/tensors.jl b/test/cuda/tensors.jl index 68e3bfa8f..715a94b42 100644 --- a/test/cuda/tensors.jl +++ b/test/cuda/tensors.jl @@ -426,6 +426,24 @@ for V in spacelist @test tp ≈ tp * tp end end + @timedtestset "Multiplication: grouped and direct gemm via CPU" begin + W1, W2 = V1 ⊗ V2 ⊗ V3, (V4 ⊗ V5)' + threshold = CUDAExt.GROUPED_GEMM_THRESHOLD[] + try + for T in (Float32, ComplexF64), thr in (typemax(Int), 0) + CUDAExt.GROUPED_GEMM_THRESHOLD[] = (real = thr, complex = thr) + for adjA in (false, true), adjB in (false, true), β in (-1.2, T <: Complex ? -1.2 + 0.4im : 0.5) + tA = adjA ? cuRAND.rand(T, W1, W1)' : cuRAND.rand(T, W1, W1) + tB = adjB ? cuRAND.rand(T, W2, W1)' : cuRAND.rand(T, W1, W2) + tC = cuRAND.rand(T, W1, W2) + hA, hB, hC = adapt(Array, tA), adapt(Array, tB), adapt(Array, tC) + @test adapt(Array, mul!(tC, tA, tB, 0.3, β)) ≈ mul!(hC, hA, hB, 0.3, β) + end + end + finally + CUDAExt.GROUPED_GEMM_THRESHOLD[] = threshold + end + end @timedtestset "Multiplication and inverse: test via CPU" begin W1 = V1 ⊗ V2 ⊗ V3 W2 = (V4 ⊗ V5)'