Add wrappers for syrk!, gemm_batched! and syrk_batched!#19
Open
alecarraro wants to merge 2 commits into
Open
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR adds backend-specific wrappers for
syrkand batchedgemm/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.jlprovides batchedsyrk, whileCUDA.jlandMetal.jldoes not. For the latter two,syrk_batched!falls back togemm_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 fromLinearAlgebrafor all supported backends, so there is no need to wrap it. However, this PR adds wrappers forgemmEx!andgemmEx_batched!on CUDA and AMDGPU to support mixed-precision execution.To provide a common interface, a simplified API is defined:
This allows selecting the compute type independently of the storage type of
A,B, andC.CUDA.jlalready exposes agemmEx!wrapper, but it does not provide acompute_typekeyword 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
Exmethods 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.