4 ms·

Optimizing a Rust GPU matmul kernel

97 points by lukastyrychtr 2y ago

billti 2y ago
I've been playing around with some low-level GPU stuff lately and trying to port some of my CPU-bound matrix multiplication. The challenge has been that in the space I work in (quantum) the matrices are of complex numbers, and for a certain size of problem floats don't cut it and you need doubles. Nearly everything I find for GPUs is targeted at matrices of real floats. Any pointers to examples? I'd be fine sticking with floats as a first step, but would love to see some (reasonably optimized) low-level GPU code for working with matrices of complex numbers (preferably for Metal or WGPU, which is what I'm using).
- LegNeato 2y ago
  CUDA has support for complex numbers and some NVIDIA frameworks use and expose it (https://nvidia.github.io/cccl/thrust/api/structthrust_1_1complex.html#_CPPv4I0EN6thrust7complexE https://nvidia.github.io/cccl/thrust/api/structthrust_1_1com...). For Rust GPU, nothing built in but there are libraries like https://github.com/rust-num/num-complex https://github.com/rust-num/num-complex that support `no_std` and should work on the GPU. I've never used them so I don't know what (if any) the perf hit would be.
- pythomancer 2y ago
  In BLAS terminology this is usually called CGEMM (for single precision) or ZGEMM (for double precision). Both cuBLAS and rocBLAS support ZGEMM, and the latter is open source. rocBLAS is also pretty complicated though, and perhaps not such a good learning resource. This is a more readable library which implements CGEMM, or at least a similar operation: https://git.astron.nl/RD/recruit/ccglib https://git.astron.nl/RD/recruit/ccglib. The main issue is that double precision is not so interesting for AI and graphics, and so silicon is rather spent on more of these features and less double precision. Not so for HPC, though, and GPUs specialized for this usually have better throughput. For example, the AMD MI210 has the same performance in single and double precision (matrix) operations, while graphics GPUs either have something like 1/2, 1/4, 1/16 etc rate of fp64:fp16, or have no support at all.
taminka 2y ago
i’m always confused wrt how shader languages that target a bunch of other framework specific shader languages manage to map all the available features correctly? like tensor core support, non-uniform thread groups, different data type support, a bunch of simd group functions, stuff like that… i couldn’t find info wrt this in rust gpu, so i’m assuming it just tries to target the narrowest available feature set, that’s compatible across all shading languages?
- LegNeato 2y ago
  The goal is to handle things similar to how rust on the CPU does (shared traits with different implementations, intrinsics, platform-specific modules, user replaceable traits like alloc or panic handler and "asm" as a last ditch case). We already have a lot (we have many intrinsics and an `arch` module like `std::arch`, asm! to include raw spirv, etc). For example, here are the intrinsics: https://rust-gpu.github.io/rust-gpu/api/spirv_std/arch/index.html https://rust-gpu.github.io/rust-gpu/api/spirv_std/arch/index... and here is support for ray tracing (which obviously is not on every card: https://rust-gpu.github.io/rust-gpu/api/spirv_std/ray_tracing/index.html https://rust-gpu.github.io/rust-gpu/api/spirv_std/ray_tracin...). Vulkan has a way to query and specify different GPU capabilities and Rust-GPU uses that. Rust and Vulkan have many of the tools we need for progressive enhancement, we are not focused on lowest common denominator.
- cosmic_quanta 2y ago
  I'm not familiar with GPUs specifically, but I have seen this for ORMs that support multiple SQL dialects (e.g. [0]). A great technique is called 'tagless final encoding' [1]. Using this technique, you can specify capabilities of an embedded domain-specific language (eDSL) such that you can have a shared (but narrow) common set of features, while allowing specializations of this eDSL to support extra features. [0]: https://github.com/haskell-beam/beam https://github.com/haskell-beam/beam [1]: https://nrinaudo.github.io/articles/tagless_final.html https://nrinaudo.github.io/articles/tagless_final.html
- messe 2y ago
  It's mentioned near the beginning of the linked article: > These Rust GPU programs are then compiled into SPIR-V, a low-level format that most GPUs understand

97 points by lukastyrychtr 2y ago

I've been playing around with some low-level GPU stuff lately and trying to port some of my CPU-bound matrix multiplication. The challenge has been that in the space I work in (quantum) the matrices are of complex numbers, and for a certain size of problem floats don't cut it and you need doubles. Nearly everything I find for GPUs is targeted at matrices of real floats. Any pointers to examples? I'd be fine sticking with floats as a first step, but would love to see some (reasonably optimized) low-level GPU code for working with matrices of complex numbers (preferably for Metal or WGPU, which is what I'm using).

LegNeato 2y ago
CUDA has support for complex numbers and some NVIDIA frameworks use and expose it (https://nvidia.github.io/cccl/thrust/api/structthrust_1_1complex.html#_CPPv4I0EN6thrust7complexE https://nvidia.github.io/cccl/thrust/api/structthrust_1_1com...). For Rust GPU, nothing built in but there are libraries like https://github.com/rust-num/num-complex https://github.com/rust-num/num-complex that support `no_std` and should work on the GPU. I've never used them so I don't know what (if any) the perf hit would be.
pythomancer 2y ago
In BLAS terminology this is usually called CGEMM (for single precision) or ZGEMM (for double precision). Both cuBLAS and rocBLAS support ZGEMM, and the latter is open source. rocBLAS is also pretty complicated though, and perhaps not such a good learning resource. This is a more readable library which implements CGEMM, or at least a similar operation: https://git.astron.nl/RD/recruit/ccglib https://git.astron.nl/RD/recruit/ccglib. The main issue is that double precision is not so interesting for AI and graphics, and so silicon is rather spent on more of these features and less double precision. Not so for HPC, though, and GPUs specialized for this usually have better throughput. For example, the AMD MI210 has the same performance in single and double precision (matrix) operations, while graphics GPUs either have something like 1/2, 1/4, 1/16 etc rate of fp64:fp16, or have no support at all.

taminka 2y ago

i’m always confused wrt how shader languages that target a bunch of other framework specific shader languages manage to map all the available features correctly? like tensor core support, non-uniform thread groups, different data type support, a bunch of simd group functions, stuff like that… i couldn’t find info wrt this in rust gpu, so i’m assuming it just tries to target the narrowest available feature set, that’s compatible across all shading languages?

LegNeato 2y ago
The goal is to handle things similar to how rust on the CPU does (shared traits with different implementations, intrinsics, platform-specific modules, user replaceable traits like alloc or panic handler and "asm" as a last ditch case). We already have a lot (we have many intrinsics and an `arch` module like `std::arch`, asm! to include raw spirv, etc). For example, here are the intrinsics: https://rust-gpu.github.io/rust-gpu/api/spirv_std/arch/index.html https://rust-gpu.github.io/rust-gpu/api/spirv_std/arch/index... and here is support for ray tracing (which obviously is not on every card: https://rust-gpu.github.io/rust-gpu/api/spirv_std/ray_tracing/index.html https://rust-gpu.github.io/rust-gpu/api/spirv_std/ray_tracin...). Vulkan has a way to query and specify different GPU capabilities and Rust-GPU uses that. Rust and Vulkan have many of the tools we need for progressive enhancement, we are not focused on lowest common denominator.
cosmic_quanta 2y ago
I'm not familiar with GPUs specifically, but I have seen this for ORMs that support multiple SQL dialects (e.g. [0]). A great technique is called 'tagless final encoding' [1]. Using this technique, you can specify capabilities of an embedded domain-specific language (eDSL) such that you can have a shared (but narrow) common set of features, while allowing specializations of this eDSL to support extra features. [0]: https://github.com/haskell-beam/beam https://github.com/haskell-beam/beam [1]: https://nrinaudo.github.io/articles/tagless_final.html https://nrinaudo.github.io/articles/tagless_final.html
messe 2y ago
It's mentioned near the beginning of the linked article: > These Rust GPU programs are then compiled into SPIR-V, a low-level format that most GPUs understand