Skip to content

Add wrappers for syrk!, gemm_batched! and syrk_batched!#19

Open
alecarraro wants to merge 2 commits into
NextLinearAlgebra:mainfrom
alecarraro:wrap_syrk_gemm
Open

Add wrappers for syrk!, gemm_batched! and syrk_batched!#19
alecarraro wants to merge 2 commits into
NextLinearAlgebra:mainfrom
alecarraro:wrap_syrk_gemm

Conversation

@alecarraro

@alecarraro alecarraro commented Jun 3, 2026

Copy link
Copy Markdown
Collaborator

This PR adds backend-specific wrappers for syrk and batched gemm/syrk. Backend integrations are implemented through package extensions and weak dependencies, so users only need to install the GPU backend relevant to their system.

The goal is to provide a unified interface across CPU and GPU backends. NextLA now defines methods for:

  • gemm_batched! for multidimensional arrays (it maps to strided batched gemm) and vector of matrices (it maps to pointer batched gemm)
  • syrk!
  • syrk_batched!

for each supported backend.

Not all backends support all of these operations natively. When a backend provides an implementation, NextLA dispatches to it directly. Otherwise, it falls back to an equivalent implementation. For example, AMDGPU.jl provides batched syrk, while CUDA.jl and Metal.jl does not. For the latter two, syrk_batched! falls back to gemm_batched! and emits a warning. Similarly there is no batched gemm on CPU so it just fallsback to a loop of standard gemms.

One possible future direction is wrapping MAGMA batched routines. There is already a Julia package for this, Magma.jl, although it appears to be inactive.

GPU-accelerated gemm! is already overloaded from LinearAlgebra for all supported backends, so there is no need to wrap it. However, this PR adds wrappers for gemmEx! and gemmEx_batched! on CUDA and AMDGPU to support mixed-precision execution.

To provide a common interface, a simplified API is defined:

gemmEx!(
    transA, transB,
    alpha, A, B,
    beta, C;
    compute_type=default_compute_type(alpha, A, B, beta, C),
)

This allows selecting the compute type independently of the storage type of A, B, and C.
CUDA.jl already exposes a gemmEx! wrapper, but it does not provide a compute_type keyword and instead infers the compute type from the selected math mode. Here, I chose a simplified interface where the cuBLAS math mode is fixed to the default mode and the compute type is specified explicitly. This makes it possible to use the same API for both cuBLAS and rocBLAS.
The Ex methods are not exported for now. They were added because they are used by one of the TLR matrix algorithms. Although equivalent functionality is not available on all backends, keeping these wrappers available is useful for experimentation.

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant