# GPU parallelization for a very large system of ODEs

**URL:** <https://discourse.julialang.org/t/gpu-parallelization-for-a-very-large-system-of-odes/107187>\
**Category:** Performance\
**Tags:** gpu, parallel\
**Created:** [December 5, 2023, 10:28pm UTC](https://discourse.julialang.org/t/gpu-parallelization-for-a-very-large-system-of-odes/107187 "2023-12-05T22:28:38Z")\
**Posts on this page:** 6\
**Page:** 1

<div class="post-metadata">

**Author:** ![ian](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/ian/32/205324_2.png) [@ian](https://discourse.julialang.org/u/ian)\
**Post date:** [December 5, 2023, 10:28pm UTC](https://discourse.julialang.org/t/gpu-parallelization-for-a-very-large-system-of-odes/107187/1 "2023-12-05T22:28:38Z")

</div>

Hi all,

I need to simulate a very large system of ODEs (~10^5 states). At every time step in the simulation I need to update a matrix using the current value of the state vector.

Here is an MWE of the matrix updating function. This needs to be run at every timestep in the ODE integration.

```julia
using BenchmarkTools

n = 1000
coefficients = rand(n^2, 4);
state = ones(n);
randidxs = rand(1:n, n^2, 4);
result = zeros(n^2);

function viewmultsum!(result, coefficients, state, randidxs)
    @views sum!(result, coefficients .* state[randidxs])
end;

@benchmark viewmultsum!(result, coefficients, state, randidxs)

```

```julia
BenchmarkTools.Trial: 357 samples with 1 evaluation.
 Range (min … max): 10.177 ms … 54.770 ms ┊ GC (min … max): 0.00% … 30.50%
 Time (median): 14.702 ms ┊ GC (median): 0.00%
 Time (mean ± σ): 14.001 ms ± 3.392 ms ┊ GC (mean ± σ): 14.57% ± 12.51%

      ▁▄▆▄▄ ▂█▄▂                       
  ▄▆█▆█████▆▄▄▂▅▂▄▄▃▄▃▂▃▁▃▄▃▁▃▆▆▆▄▇█████▆▆▄▃▃▃▁▄▃▄▃▁▁▅▁▂▁▂▁▁▄ ▄
  10.2 ms Histogram: frequency by time 19.7 ms <

 Memory estimate: 30.52 MiB, allocs estimate: 2.

```

In the real simulation `result` is an `n x n` matrix and I calculate `result * state` at every time step (again `n > 1E5`).

Given the huge size of the state vector I want to accelerate the computation of `result` as much as I can. I’ve read the [Introduction to GPU programming](https://cuda.juliagpu.org/stable/tutorials/introduction/) in the `CUDA.jl` docs but I’m still unsure if GPU acceleration is the right approach.

In particular I’m concerned that array slicing with `state[randidxs]` won’t be very performant on the GPU hardware. Unfortunately the equations I am simulating don’t seem to allow a more structured access pattern of the `state` vector. I’m also concerned that with `n~1e5` and hence `result` (a `n x n` matrix of Float 32s) being at least several dozen GB, I will have a significant memory transfer overhead that slows down parallel computation on GPUs which have relatively little RAM.

So my specific questions are:

1. Does the unstructured access of `state` using `state[randidxs]` necessarily mean performance will be signifcantly degraded on GPUs?
2. Does the large size of the `result` array (10s of GB) mean that data transfer overhead will kill performance gains from moving to GPUs? It seems the largest GPU RAM is about 80GB and `result` may be bigger than that as `n` gets very large.
3. Are there alternative approaches I should investigate before fully committing to the GPU route?

Thank you!

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [December 7, 2023, 4:56am UTC](https://discourse.julialang.org/t/gpu-parallelization-for-a-very-large-system-of-odes/107187/2 "2023-12-07T04:56:10Z")

</div>

> [@ian](#):
>
> In particular I’m concerned that array slicing with `state[randidxs]` won’t be very performant on the GPU hardware.

That’s fine.

> [@ian](#):
>
> - Does the large size of the `result` array (10s of GB) mean that data transfer overhead will kill performance gains from moving to GPUs? It seems the largest GPU RAM is about 80GB and `result` may be bigger than that as `n` gets very large.

If you keep the entire computation on the GPU you should be fine. That said, if your operations are all only O(n) like shown in your example, then you’re not likely to see that much of a speedup.

> [@ian](#):
>
> 1. Are there alternative approaches I should investigate before fully committing to the GPU route?

What is your actual equation?

---

<div class="post-metadata">

**Author:** ![ian](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/ian/32/205324_2.png) [@ian](https://discourse.julialang.org/u/ian)\
**Post date:** [December 7, 2023, 4:25pm UTC](https://discourse.julialang.org/t/gpu-parallelization-for-a-very-large-system-of-odes/107187/3 "2023-12-07T16:25:18Z")

</div>

Thanks for the response!

> If you keep the entire computation on the GPU you should be fine. That said, if your operations are all only O(n) like shown in your example, then you’re not likely to see that much of a speedup.

Regarding the computations being only O(n) and hence not likely to be sped up. My understanding was that moving the computation to the GPU would parallelize `viewmultsum!`

```julia
function viewmultsum!(result, coefficients, state, randidxs)
    @views sum!(result, coefficients .* state[randidxs])
end;

```

and therefore roughly reduce the computational time by a factor proportional to the number of GPU cores. Are you suggesting I need to write some kind of kernel that is explicitly parallel instead of using broadcasting? For example something like the following (using 4 threads) but with a large number of threads

```julia
n = 1000
coefficients = rand(n^2, 4);
state = ones(n);
randidxs = rand(1:n, n^2, 4);
result = zeros(n^2);

function loopmultsum!(result, coefficients, state, randidxs)
    is, js = axes(coefficients)
    for j in js
        Threads.@threads for i in is
            c, r = coefficients[i, j], randidxs[i, j]
            s = state[r]
            result[i] = (j == first(js)) ? c * s : muladd(c, s, result[i])
        end
    end
end;

```

```julia
BenchmarkTools.Trial: 1193 samples with 1 evaluation.
 Range (min … max): 3.011 ms … 41.851 ms ┊ GC (min … max): 0.00% … 0.00%
 Time (median): 3.878 ms ┊ GC (median): 0.00%
 Time (mean ± σ): 4.181 ms ± 1.893 ms ┊ GC (mean ± σ): 0.00% ± 0.00%

          ▂█▄                                                 
  ▃▆▆▅▅▃▄▆████▅▄▄▄▅▃▄▃▃▃▃▃▃▂▂▂▃▂▂▂▂▂▂▂▃▂▂▁▂▂▂▂▁▂▁▂▂▁▂▁▁▂▂▂▁▂ ▃
  3.01 ms Histogram: frequency by time 8.05 ms <

 Memory estimate: 8.59 KiB, allocs estimate: 95.

```

I wasn’t sure if the nested looping would degrade GPU performance.

And regarding

> What is your actual equation?

The actual equation is ` du/dt = (D + R) u + c`. Where `D` is a constant diagonal matrix, `R` is the `result` from the code above reshaped to be `n x n`, and `c` is a constant vector. `R` is not actually filled by randomly indexing `state` (i.e. `u`) as above, but I don’t think getting into the details of how `state` is indexed would be helpful. Suffice it to say that instead of the `randidxs` array above, I have another integer filled array `trueidxs` that I use to index `state`. In short the equation is a large, nonlinear (due to `R`) system of first order ODEs.

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [December 10, 2023, 1:08pm UTC](https://discourse.julialang.org/t/gpu-parallelization-for-a-very-large-system-of-odes/107187/4 "2023-12-10T13:08:59Z")

</div>

> [@ian](#):
>
> Regarding the computations being only O(n) and hence not likely to be sped up. My understanding was that moving the computation to the GPU would parallelize `viewmultsum!`
> 
> ```julia
> 
> ```

It would, but you need quite a bit of compute for GPUs to be faster than CPUs for O(n) calculations. They just aren’t great at things with that scaling.

> [@ian](#):
>
> and therefore roughly reduce the computational time by a factor proportional to the number of GPU cores.

No, GPU cores are much slower than CPU cores, so that mental math is mis calibrated. For a larger discussion, see [Doing small network scientific machine learning in Julia 5x faster than PyTorch](https://julialang.org/blog/2022/04/simple-chains/).

![image](https://global.discourse-cdn.com/julialang/original/3X/8/d/8da4aad1f674b20e6412a8bf03ee2f00e7d6e5c2.png)

At N=1000 you start to break even with CPU on Float32 operations. With Float64, the breakeven is closer to N=10000 because of how much slower the FPUs are on GPUs.

> [@ian](#):
>
> The actual equation is ` du/dt = (D + R) u + c`. Where `D` is a constant diagonal matrix, `R` is the `result` from the code above reshaped to be `n x n`, and `c` is a constant vector. `R` is not actually filled by randomly indexing `state` (i.e. `u`) as above, but I don’t think getting into the details of how `state` is indexed would be helpful. Suffice it to say that instead of the `randidxs` array above, I have another integer filled array `trueidxs` that I use to index `state`. In short the equation is a large, nonlinear (due to `R`) system of first order ODEs.

Can you express R as a sparse matrix or more directly as code? Looks like a PDE discretization, using arrays is generally slow for that. How dense is R?

---

<div class="post-metadata">

**Author:** ![ian](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/ian/32/205324_2.png) [@ian](https://discourse.julialang.org/u/ian)\
**Post date:** [December 10, 2023, 3:18pm UTC](https://discourse.julialang.org/t/gpu-parallelization-for-a-very-large-system-of-odes/107187/5 "2023-12-10T15:18:58Z")

</div>

Thanks for the reply!

> It would, but you need quite a bit of compute for GPUs to be faster than CPUs for O(n) calculations. They just aren’t great at things with that scaling.

I think I see what you mean with GPUs not being great on calculations with O(n) scaling. In my case I thought the algorithm to fill `result` is O(n^2) since `result` has a shape of `n x n`. Maybe the for-loop version of the code makes that clearer (copied below for reference). Let me know if I am mistaken.

```julia
n = 1000
coefficients = rand(n^2, 4);
state = ones(n);
idxs = rand(1:n, n^2, 4);
result = zeros(n^2);

function loopmultsum!(result, coefficients, state, randidxs)
    is, js = axes(coefficients)
    for j in js
        Threads.@threads for i in is
            c, r = coefficients[i, j], randidxs[i, j]
            s = state[r]
            result[i] = (j == first(js)) ? c * s : muladd(c, s, result[i])
        end
    end
end;

```

Regarding the sparsity of `R` and how to construct it as code. It is constructed very similarly to `result` in `viewmultsum!` or `loopmultsum!` above. The only difference is reshaping and the specific values in `coefficients` and `randidxs`. Concretely to get `R` I write

```julia
n = 1000
coefficients = rand(n^2, 4);
state = ones(n);
randidxs = rand(1:n, n^2, 4);
R= zeros(n, n);

function viewmultsum!(R, coefficients, state, randidxs)
    n = size(R)[1]
    R = reshape(R, :)
    @views sum!(R, coefficients .* state[randidxs])
    R = reshape(R, n, n)
end;

```

It doesn’t seem there is much overhead associated with reshaping so I didn’t include that part in the original post.

> How dense is R?

It turns out that for the real problem the `R` matrix is pretty sparse, e.g.

`count(iszero, R)/length(R) = .9778738175089268`

for `n ~ 9000` (and hence `R` having around `81M` elements since it is `n x n`). It also seems that `R` gets more sparse as `n` grows. Are you suggesting that there is some way to greatly reduce computational demands using the sparsity of `R`? Do SparseArrays have GPU support?

Thanks again.

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [December 30, 2023, 10:46am UTC](https://discourse.julialang.org/t/gpu-parallelization-for-a-very-large-system-of-odes/107187/6 "2023-12-30T10:46:48Z")

</div>

> [@ian](#):
>
> It turns out that for the real problem the `R` matrix is pretty sparse, e.g.

That’s not really sparse enough to warrant a sparse matrix. You need it to be like less than 1 in a thousand. So on the larger end it would make sense.

> [@ian](#):
>
> Are you suggesting that there is some way to greatly reduce computational demands using the sparsity of `R`? Do SparseArrays have GPU support?

If it’s sparse enough then yes. And yes SparseArrays has GPU support, there’s GPU sparse solvers. But GPU sparse solvers do not scale as well as they do in the dense case, so you need some rather large and structured sparsity patterns for them out perform the CPU case.