# Batch matrix/vector operations with CUDA.jl

**URL:** <https://discourse.julialang.org/t/batch-matrix-vector-operations-with-cuda-jl/117869>\
**Category:** GPU\
**Tags:** question\
**Created:** [August 6, 2024, 6:28am UTC](https://discourse.julialang.org/t/batch-matrix-vector-operations-with-cuda-jl/117869 "2024-08-06T06:28:39Z")\
**Posts on this page:** 6\
**Page:** 1

<div class="post-metadata">

**Author:** ![TimHargreaves](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/timhargreaves/32/207284_2.png) [@TimHargreaves](https://discourse.julialang.org/u/TimHargreaves)\
**Post date:** [August 6, 2024, 6:28am UTC](https://discourse.julialang.org/t/batch-matrix-vector-operations-with-cuda-jl/117869/1 "2024-08-06T06:28:39Z")

</div>

I’m trying to perform various operations such as multiplication, inversion, solving, Cholesky decomposition in batches with CUDA.jl.

As far as I’m aware, there is no high-level API for doing this, as is the case with Tensorflow/PyTorch, e.g. [tf.linalg.inv](https://www.tensorflow.org/api_docs/python/tf/linalg/inv) which natively supports batched operations.

Instead, I have been trying to use the low-level CUBLAS wrappers. I would have expected these to accept multi-dimensional arrays […, N, N] but instead most of them seem to accept vectors of CuArrays.

[Here](https://github.com/JuliaGPU/CUDA.jl/blob/d7077da2b7df32f9d0a2bced56511cdd778ab4ed/test/libraries/cublas.jl#L146) is an example using gemv\_batched!.

I wanted to confirm two things:

1. Is this the correct way to be performing batched operations with CUDA.jl?
2. Are these operations actually parallelised over the batch or just within each single operation (it’s the vector of CuArrays that is making me doubt this)?

As a follow-up question, are there any plans to implement a Tensorflow/PyTorch style high-level interface for batch operations. If so, I would be happy to help where I can as my research would benefit greatly from such implementations.

---

<div class="post-metadata">

**Author:** ![maleadt](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/maleadt/32/10097_2.png) [@maleadt](https://discourse.julialang.org/u/maleadt)\
**Post date:** [August 6, 2024, 7:50pm UTC](https://discourse.julialang.org/t/batch-matrix-vector-operations-with-cuda-jl/117869/2 "2024-08-06T19:50:08Z")

</div>

> [@TimHargreaves](#):
>
> I would have expected these to accept multi-dimensional arrays […, N, N] but instead most of them seem to accept vectors of CuArrays.

The `strided_batched` methods are the ones accepting multi-dimensional arrays (as per NVIDIA naming), and are supposed to be faster than the ones using vectors of GPU arrays: [https://developer.nvidia.com/blog/cublas-strided-batched-matrix-multiply/](https://developer.nvidia.com/blog/cublas-strided-batched-matrix-multiply/)

Other than that, I’m not terribly familiar with the use or design of batched APIs, so help is always appreciated. There is some existing work though, like `batched_mul!` in NNlib.jl, Batched.jl, BatchedBLAS.jl, etc.

---

<div class="post-metadata">

**Author:** ![amontoison](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/amontoison/32/218741_2.png) [@amontoison](https://discourse.julialang.org/u/amontoison)\
**Post date:** [August 8, 2024, 2:10am UTC](https://discourse.julialang.org/t/batch-matrix-vector-operations-with-cuda-jl/117869/3 "2024-08-08T02:10:33Z")

</div>

We have a few low-level wrappers for batch operations in `CUBLAS` and `CUSOLVER`:

- [CUDA.jl/lib/cublas/wrappers.jl at master · JuliaGPU/CUDA.jl · GitHub](https://github.com/JuliaGPU/CUDA.jl/blob/master/lib/cublas/wrappers.jl#L402-L486)
- [CUDA.jl/lib/cusolver/dense.jl at master · JuliaGPU/CUDA.jl · GitHub](https://github.com/JuliaGPU/CUDA.jl/blob/master/lib/cusolver/dense.jl#L821-L898)

I confirm that these operations are parallelized over the batch.  
I am reluctant to adding an high-level dispatch (mul!, cholesky, etc…) because we don’t have a “batch” version for CPU.

---

<div class="post-metadata">

**Author:** ![TimHargreaves](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/timhargreaves/32/207284_2.png) [@TimHargreaves](https://discourse.julialang.org/u/TimHargreaves)\
**Post date:** [August 27, 2024, 9:51pm UTC](https://discourse.julialang.org/t/batch-matrix-vector-operations-with-cuda-jl/117869/4 "2024-08-27T21:51:30Z")

</div>

Thank you both for clarifying!

---

<div class="post-metadata">

**Author:** ![TimHargreaves](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/timhargreaves/32/207284_2.png) [@TimHargreaves](https://discourse.julialang.org/u/TimHargreaves)\
**Post date:** [September 4, 2024, 1:03pm UTC](https://discourse.julialang.org/t/batch-matrix-vector-operations-with-cuda-jl/117869/5 "2024-09-04T13:03:36Z")

</div>

> [@amontoison](#):
>
> I am reluctant to adding an high-level dispatch (mul!, cholesky, etc…) because we don’t have a “batch” version for CPU.

This is a very fair point and I can see how this could lead to confusion. How would you feel about adding a high-level interface specifically for the non-strided batch case?

I believe all of these are already supported by Julia through automatic broadcasting. E.g.

```julia
As = [rand(2, 2) for i in 1:10]
Bs = [rand(2, 3) for i in 1:10]
As .* Bs

```

We could just intercept the broadcasting and replace it with a call to `gemm_batched`.

---

<div class="post-metadata">

**Author:** ![amontoison](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/amontoison/32/218741_2.png) [@amontoison](https://discourse.julialang.org/u/amontoison)\
**Post date:** [September 4, 2024, 1:34pm UTC](https://discourse.julialang.org/t/batch-matrix-vector-operations-with-cuda-jl/117869/6 "2024-09-04T13:34:39Z")

</div>

It’s an excellent idea @TimHargreaves, if we can intercept the `.*`, it makes sense to dispatch to `gemm_batched`.

We just need to check if a recent CUDA release is installed (12.4) because the matrices in `As` and `Bs` needed to have the same shape before.  
It also explains why we never think of that, the CUDA routine was too restrictive on `As` and `Bs`.