# Improving GPU performance for symbolic regression
**URL:** https://discourse.julialang.org/t/improving-gpu-performance-for-symbolic-regression/108800
**Category:** Performance
**Tags:** question, gpu, cuda, symbolic-regression
**Created:** [January 14, 2024, 9:29pm UTC](https://discourse.julialang.org/t/improving-gpu-performance-for-symbolic-regression/108800 "2024-01-14T21:29:48Z")
**Posts on this page:** 15
**Page:** 1
### Author:  [@MilesCranmer](https://discourse.julialang.org/u/MilesCranmer)
#### Post date: [January 14, 2024, 9:29pm UTC](https://discourse.julialang.org/t/improving-gpu-performance-for-symbolic-regression/108800/1 "2024-01-14T21:29:48Z")
I’m trying to improve the GPU performance for evaluation of an expression represented as a tree in DynamicExpressions.jl. The way it works is that an expression is created at runtime – here’s a movie of how this is used in the context of SymbolicRegression.jl:
The bottleneck has always been in the evaluation of these expressions, and therefore that part has gone through many development cycles to get better perf.
To try to speed things up further I’m planning to allow the user to pass a `CuArray` for the dataset, so that the evaluation of these trees is faster.
~~Right now I have a recursive evaluation scheme that does a depth-first walk through the tree, building CUDA arrays at the leafs, and computing the operation at each node. This is my current code as drafted in [this PR](https://github.com/SymbolicML/DynamicExpressions.jl/pull/62):~~
~~[removed; see below]~~
~~It uses an internal [`tree_mapreduce`](https://symbolicml.org/DynamicExpressions.jl/dev/utils/#Internal-utilities) method which as the name suggests is a generalization of a regular mapreduce to a tree-like structure, so you can have different a “map” for leaf nodes or branch nodes, and aggregate on `(parents, children...)`~~
~~With this, I am seeing the following performance:~~
~~[removed; see below]~~
~~So it is sitting at about a 5.5x speedup. However, this `10,000` array size is actually larger than typical – you might only evaluate over a small batch of ~100 elements at a time. At that size, the CPU is ~microseconds while the GPU is stil at about 1.5ms.~~
~~Therefore I am wondering what my options are here to improve this. Is there something in the CUDA array operations I’m not doing correctly, which is perhaps causing the CUDA threads to sync before the next operation? Or is this fundamentally a limitation with this approach? It doesn’t seem possible to make a single CUDA kernel for this entire scheme, unfortunately.~~
Edit: see latest version [here](https://discourse.julialang.org/t/improving-gpu-performance-for-symbolic-regression/108800/6). It’s all done in a single CUDA kernel.
---
### Author:  [@Raf](https://discourse.julialang.org/u/Raf)
#### Post date: [January 14, 2024, 11:10pm UTC](https://discourse.julialang.org/t/improving-gpu-performance-for-symbolic-regression/108800/2 "2024-01-14T23:10:28Z")
GPUs are almost never worth using for small arrays because kernel launch time will dwarf the operations, and 100 operations can only use a small fraction of your individually quite slow GPU cores anyway. Your few but super fast CPU cores will always eat problems like that for breakfast.
Try `CUDA.@profile` on your function for a breakdown of the timings.
There ares some things you can do. I’ve recently rewritten a bunch of my simulation code to handle this problem - I need replicates of every run (its stochastic) so I can stack 100 replicates of the array on another dimension to make it large enough to be worth launching a GPU kernel for.
But generally you need to have some single thing you need to do quite a few thousand times for it to be worth using a GPU.
---
### Author:  [@MilesCranmer](https://discourse.julialang.org/u/MilesCranmer)
#### Post date: [January 14, 2024, 11:41pm UTC](https://discourse.julialang.org/t/improving-gpu-performance-for-symbolic-regression/108800/3 "2024-01-14T23:41:47Z")
Thanks. I was hoping there could be some way for it to lazily dispatch the calls to the GPU or something, since it does technically end up being 300x evaluations of the same array shape (just with this complex dependency structure which gets configured on the CPU) – and closer to the leafs of the tree, you can actually have many operations being executed at once, since they don’t depend on eachother. But maybe I’m being overly optimistic about what can be done automatically here.
I can also pre-compute the entire allocation needed (= 150 x 100 elements for batch size of 100 and 150 leaf nodes) in advance; would that maybe help at all? It would still have this graph of calculation dependencies but it at least wouldn’t need to create allocations repeatedly.
---
### Author:  [@MilesCranmer](https://discourse.julialang.org/u/MilesCranmer)
#### Post date: [January 15, 2024, 12:49am UTC](https://discourse.julialang.org/t/improving-gpu-performance-for-symbolic-regression/108800/4 "2024-01-15T00:49:15Z")
I wonder if it’s also possible to do this entire thing in a single CUDA kernel on a preallocated storage array? The tricky thing is you would need to have a particular dependency structure for the threads so that the leaf nodes execute first, the (depth-1) nodes second, and so on. But you could totally generate all of that information on the CPU before making the kernel call, and pass it as an array to the GPU.
```julia
julia> using DynamicExpressions
julia> operators = OperatorEnum(binary_operators=(+, -, *, /), unary_operators=(cos, sin));
julia> x1, x2, x3 = [Node{Float64}(feature=i) for i in 1:3];
julia> tree = cos(x1) * x2 - 0.9 / x3; for _=1:2; tree=cos(tree) - 0.9*tree; end; tree
cos(cos((cos(x1) * x2) - (0.9 / x3)) - (0.9 * ((cos(x1) * x2) - (0.9 / x3)))) - (0.9 * (cos((cos(x1) * x2) - (0.9 / x3)) - (0.9 * ((cos(x1) * x2) - (0.9 / x3)))))
julia> local_depth = count_depth.(collect(nodes));
5-element Vector{Int64}:
8
7
6
5
4
```
This `local_depth` would be the depth at each node position, and basically tells you the execution order so that all dependent parts of the computation are executed beforehand. i.e., leaf nodes (like constants or features) would have a depth of 1, so they are executed first.
---
### Author:  [@MilesCranmer](https://discourse.julialang.org/u/MilesCranmer)
#### Post date: [January 15, 2024, 1:50am UTC](https://discourse.julialang.org/t/improving-gpu-performance-for-symbolic-regression/108800/5 "2024-01-15T01:50:38Z")
Edit: out-of-date; see next post!
---
### Author:  [@MilesCranmer](https://discourse.julialang.org/u/MilesCranmer)
#### Post date: [January 15, 2024, 5:49am UTC](https://discourse.julialang.org/t/improving-gpu-performance-for-symbolic-regression/108800/6 "2024-01-15T05:49:25Z")
Quick update – got it working with a single CUDA kernel! Here’s the [code](https://github.com/SymbolicML/DynamicExpressions.jl/blob/20a8f389cdff8628dafb13222fd66711ea45292d/ext/DynamicExpressionsCUDAExt.jl).
> **(click to expand)**
>
> ```julia
> module DynamicExpressionsCUDAExt
> #! format: off
>
> using CUDA
> using DynamicExpressions: OperatorEnum, AbstractExpressionNode
> using DynamicExpressions.EvaluateEquationModule: get_nbin, get_nuna
> using DynamicExpressions.AsArrayModule: as_array
>
> import DynamicExpressions.EvaluateEquationModule: eval_tree_array
>
> function eval_tree_array(
> tree::AbstractExpressionNode{T}, gcX::CuArray{T,2}, operators::OperatorEnum; _...
> ) where {T<:Number}
> (outs, is_good) = eval_tree_array((tree,), gcX, operators)
> return (only(outs), only(is_good))
> end
>
> function eval_tree_array(
> trees::NTuple{M,N}, gcX::CuArray{T,2}, operators::OperatorEnum; _...
> ) where {T<:Number,N<:AbstractExpressionNode{T},M}
> (; val, execution_order, roots, buffer) = as_array(Int32, trees...)
> num_launches = maximum(execution_order)
> num_elem = size(gcX, 2)
> num_nodes = size(buffer, 2)
>
> ## Floating point arrays:
> gworkspace = CuArray{T}(undef, num_elem, num_nodes)
> gval = CuArray(val)
>
> ## Index arrays (much faster to have `@view` here)
> gbuffer = CuArray(buffer)
> gdegree = @view gbuffer[1, :]
> gfeature = @view gbuffer[2, :]
> gop = @view gbuffer[3, :]
> gexecution_order = @view gbuffer[4, :]
> gidx_self = @view gbuffer[5, :]
> gidx_l = @view gbuffer[6, :]
> gidx_r = @view gbuffer[7, :]
> gconstant = @view gbuffer[8, :]
>
> num_threads = 256
> num_blocks = nextpow(2, ceil(Int, num_elem * num_nodes / num_threads))
>
> _launch_gpu_kernel!(
> num_threads, num_blocks, num_launches, gworkspace,
> # Thread info:
> num_elem, num_nodes, gexecution_order,
> # Input data and tree
> operators, gcX, gidx_self, gidx_l, gidx_r,
> gdegree, gconstant, gval, gfeature, gop,
> )
>
> out = ntuple(
> i -> @view(gworkspace[:, roots[i]]),
> Val(M)
> )
> is_good = ntuple(
> i -> true, # Up to user to find NaNs
> Val(M)
> )
>
> return (out, is_good)
> end
>
> function _launch_gpu_kernel!(
> num_threads, num_blocks, num_launches::Integer, buffer::AbstractArray{T,2},
> # Thread info:
> num_elem::Integer, num_nodes::Integer, execution_order::AbstractArray{I},
> # Input data and tree
> operators::OperatorEnum, cX::AbstractArray{T,2}, idx_self::AbstractArray, idx_l::AbstractArray, idx_r::AbstractArray,
> degree::AbstractArray, constant::AbstractArray, val::AbstractArray{T,1}, feature::AbstractArray, op::AbstractArray,
> ) where {I,T}
> nuna = get_nuna(typeof(operators))
> nbin = get_nbin(typeof(operators))
> (nuna > 10 || nbin > 10) && error("Too many operators. Kernels are only compiled up to 10.")
> gpu_kernel! = create_gpu_kernel(operators, Val(nuna), Val(nbin))
> for launch in one(I):I(num_launches)
> @cuda threads=num_threads blocks=num_blocks gpu_kernel!(
> buffer,
> launch, num_elem, num_nodes, execution_order,
> cX, idx_self, idx_l, idx_r,
> degree, constant, val, feature, op
> )
> end
> return nothing
> end
>
> # Need to pre-compute the GPU kernels with an `@eval` for each number of operators
> # 1. We need to use an `@nif` over operators, as GPU kernels
> # can't index into arrays of operators.
> # 2. `@nif` is evaluated at parse time and needs to know the number of
> # ifs to generate at that time, so we can't simply use specialization.
> # 3. We can't use `@generated` because we can't create closures in those.
> for nuna in 0:10, nbin in 0:10
> @eval function create_gpu_kernel(operators::OperatorEnum, ::Val{$nuna}, ::Val{$nbin})
> function (
> # Storage:
> buffer,
> # Thread info:
> launch::Integer, num_elem::Integer, num_nodes::Integer, execution_order::AbstractArray,
> # Input data and tree
> cX::AbstractArray, idx_self::AbstractArray, idx_l::AbstractArray, idx_r::AbstractArray,
> degree::AbstractArray, constant::AbstractArray, val::AbstractArray, feature::AbstractArray, op::AbstractArray,
> )
> i = (blockIdx().x - 1) * blockDim().x + threadIdx().x
> if i > num_elem * num_nodes
> return nothing
> end
>
> node = (i - 1) % num_nodes + 1
> elem = (i - node) ÷ num_nodes + 1
>
> if execution_order[node] != launch
> return nothing
> end
>
> cur_degree = degree[node]
> cur_idx = idx_self[node]
> if cur_degree == 0
> if constant[node] == 1
> cur_val = val[node]
> buffer[elem, cur_idx] = cur_val
> else
> cur_feature = feature[node]
> buffer[elem, cur_idx] = cX[cur_feature, elem]
> end
> else
> if cur_degree == 1
> cur_op = op[node]
> l_idx = idx_l[node]
> Base.Cartesian.@nif(
> $nuna,
> i -> i == cur_op,
> i -> let op = operators.unaops[i]
> buffer[elem, cur_idx] = op(buffer[elem, l_idx])
> end
> )
> else
> cur_op = op[node]
> l_idx = idx_l[node]
> r_idx = idx_r[node]
> Base.Cartesian.@nif(
> $nbin,
> i -> i == cur_op,
> i -> let op = operators.binops[i]
> buffer[elem, cur_idx] = op(buffer[elem, l_idx], buffer[elem, r_idx])
> end
> )
> end
> end
> return nothing
> end
> end
> end
>
> #! format: on
> end
>
> ```
The key function is `_launch_gpu_kernel` which actually launches the kernels. Basically, each `launch`, it goes one up the tree to the next node that can be computed, and spawns all the threads to all (array elements) x (available nodes).
In gif form, the CUDA launches are like this:
And the threads which don’t compute anything simply quit early. This is about 30% faster than the initial naive version I tried. I’m assuming there’s remaining slowness due to me misusing some part of CUDA.jl, so any other tips much appreciated 🙂
Also I want to note that this new attempt allows batching over expressions! So you can do, e.g.,
```julia
julia> eval_tree_array((tree1, tree2, tree3), X, operators)
```
and it will put them in one big array and work on them at the same time. So maybe this is a way to get your suggestion working @Raf – I need to find a way to batch evaluate over expressions.
Finally, here’s the profiling results:
```julia
julia> CUDA.@profile eval_tree_array(tree, X, operators)[1]
Profiler ran for 638.25 µs, capturing 376 events.
Host-side activity: calling CUDA APIs took 377.89 µs (59.21% of the trace)
┌──────────┬────────────┬───────┬───────────────────────────────────────┬─────────────────────────┐
│ Time (%) │ Total time │ Calls │ Time distribution │ Name │
├──────────┼────────────┼───────┼───────────────────────────────────────┼─────────────────────────┤
│ 21.93% │ 139.95 µs │ 27 │ 5.18 µs ± 15.46 ( 1.43 ‥ 82.49) │ cuMemAllocFromPoolAsync │
│ 19.24% │ 122.79 µs │ 31 │ 3.96 µs ± 3.28 ( 2.62 ‥ 20.98) │ cuLaunchKernel │
│ 5.53% │ 35.29 µs │ 9 │ 3.92 µs ± 4.71 ( 2.15 ‥ 16.45) │ cuMemcpyHtoDAsync │
│ 2.69% │ 17.17 µs │ 1 │ │ cuMemcpyDtoHAsync │
│ 1.57% │ 10.01 µs │ 11 │ 910.33 ns ± 381.74 (476.84 ‥ 1668.93) │ cuStreamSynchronize │
└──────────┴────────────┴───────┴───────────────────────────────────────┴─────────────────────────┘
Device-side activity: GPU was busy for 188.59 µs (29.55% of the trace)
┌──────────┬────────────┬───────┬─────────────────────────────────────┬──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
│ Time (%) │ Total time │ Calls │ Time distribution │ Name ⋯
├──────────┼────────────┼───────┼─────────────────────────────────────┼──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
│ 21.52% │ 137.33 µs │ 14 │ 9.81 µs ± 1.72 ( 6.91 ‥ 12.87) │ _Z3_5913CuDeviceArrayI7Float32Li2ELi1EE5Int64S1_S_I4BoolLi1ELi1EES_IS0_Li2ELi1EES_I5Int32Li1ELi1EES_IS3_Li1ELi1EES_IS3_Li1ELi1EES_I5UInt8Li1ELi1EES_IS2_Li1ELi1EES_IS0_Li1ELi1EES_I6UInt16L ⋯
│ 4.41% │ 28.13 µs │ 14 │ 2.01 µs ± 0.12 ( 1.91 ‥ 2.15) │ _Z16broadcast_kernel15CuKernelContext13CuDeviceArrayI4BoolLi1ELi1EE11BroadcastedI12CuArrayStyleILi1EE5TupleI5OneToI5Int64EE2__S4_I8ExtrudedIS0_I5Int32Li1ELi1EES4_IS1_ES4_IS6_EES9_EES6_ ⋯
│ 1.91% │ 12.16 µs │ 9 │ 1.35 µs ± 0.12 ( 1.19 ‥ 1.43) │ [copy pageable to device memory] ⋯
│ 0.67% │ 4.29 µs │ 1 │ │ _Z22partial_mapreduce_grid8identity7add_sum7Float3216CartesianIndicesILi1E5TupleI5OneToI5Int64EEES2_ILi1ES3_IS4_IS5_EEE3ValILitrueEE13CuDeviceArrayIS1_Li2ELi1EE11BroadcastedI12CuArrayStyl ⋯
│ 0.41% │ 2.62 µs │ 1 │ │ _Z16broadcast_kernel15CuKernelContext13CuDeviceArrayI7Float32Li1ELi1EE11BroadcastedI12CuArrayStyleILi1EE5TupleI5OneToI5Int64EE1_S4_I8ExtrudedIS0_IS1_Li1ELi1EES4_I4BoolES4_IS6_EES1_EES6_ ⋯
│ 0.37% │ 2.38 µs │ 1 │ │ _Z15getindex_kernel15CuKernelContext13CuDeviceArrayI7Float32Li1ELi1EES0_IS1_Li2ELi1EE5TupleI5Int64S3_E5SliceI5OneToIS3_EES3_ ⋯
│ 0.26% │ 1.67 µs │ 1 │ │ [copy device to pageable memory] ⋯
└──────────┴────────────┴───────┴─────────────────────────────────────┴──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
1 column omitted
```
* * *
Edit: Improved some of the redundant allocations. Now I’m seeing the GPU have a 3x improvement in perf on even 1,000 element arrays.
---
### Author:  [@Zentrik](https://discourse.julialang.org/u/Zentrik)
#### Post date: [January 15, 2024, 7:21am UTC](https://discourse.julialang.org/t/improving-gpu-performance-for-symbolic-regression/108800/7 "2024-01-15T07:21:34Z")
I haven’t looked through your code thoroughly so some general tips I would have are use 32 bit integers and floats if possible, don’t use recursion but instead a for/while loop and inline aggressively. Check out [Performance Tips · CUDA.jl](https://cuda.juliagpu.org/stable/tutorials/performance/)
EDIT: you don’t seem to be using recursion in the kernel so ignore that. One other thing would be using a grid stride loop so that you can use the optimal thread/ block sizes according to the occupancy calculator. I’m not sure that would be a big benefit though.
---
### Author:  [@MilesCranmer](https://discourse.julialang.org/u/MilesCranmer)
#### Post date: [January 15, 2024, 7:36am UTC](https://discourse.julialang.org/t/improving-gpu-performance-for-symbolic-regression/108800/8 "2024-01-15T07:36:51Z")
The CPU evaluation is recursion based (I tried a stack-based one a while ago but it didn’t make any difference and had a much greater complexity). But the new GPU one does everything in parallel for expression tree nodes which are independent from one another. Let me update the first post because that’s an older version.
There’s no recursion and the only for loop right now is in the `_launch_gpu_kernel!` part is to avoid data races (see gif above). So it looks like this:
```julia
function _launch_gpu_kernel!(
num_threads, num_blocks, num_launches::Integer, buffer::AbstractArray{T,2},
# Thread info:
num_elem::Integer, num_nodes::Integer, execution_order::AbstractArray{I},
# Input data and tree
operators::OperatorEnum, cX::AbstractArray{T,2}, idx_self::AbstractArray, idx_l::AbstractArray, idx_r::AbstractArray,
degree::AbstractArray, constant::AbstractArray, val::AbstractArray{T,1}, feature::AbstractArray, op::AbstractArray,
) where {I,T}
nuna = get_nuna(typeof(operators))
nbin = get_nbin(typeof(operators))
(nuna > 10 || nbin > 10) && error("Too many operators. Kernels are only compiled up to 10.")
gpu_kernel! = create_gpu_kernel(operators, Val(nuna), Val(nbin))
for launch in one(I):I(num_launches)
@cuda threads=num_threads blocks=num_blocks gpu_kernel!(
buffer,
launch, num_elem, num_nodes, execution_order,
cX, idx_self, idx_l, idx_r,
degree, constant, val, feature, op
)
end
return nothing
end
```
Every GPU thread checks the `launch` integer here to see if it is its turn to execute; otherwise it exits. But if there’s a better way to force the threads to run in a particular order (without needing to dispatch again and again) I’d love to hear it.
---
### Author:  [@MilesCranmer](https://discourse.julialang.org/u/MilesCranmer)
#### Post date: [January 15, 2024, 7:48am UTC](https://discourse.julialang.org/t/improving-gpu-performance-for-symbolic-regression/108800/9 "2024-01-15T07:48:46Z")
I also made some improvements to the memory allocation. It seems like each new call to `CuArray` is the bottleneck right now, so the new version (updated the code) uses views of a large array of integers (even for `Bool`). This helps a _lot_.
Surprisingly those parts are still the bottleneck though (more typical array size used for this) –
I’m wondering if there’s any way I can just have a single `CuArray` call at the top and use `@view` and interpret the bits as `Float32/Int32` as needed. Is that possible here?
---
### Author:  [@MilesCranmer](https://discourse.julialang.org/u/MilesCranmer)
#### Post date: [January 15, 2024, 7:56am UTC](https://discourse.julialang.org/t/improving-gpu-performance-for-symbolic-regression/108800/10 "2024-01-15T07:56:07Z")
Alright so that big red bar in that flame graph turned out to be this call in CUDA.jl in `execution.jl`:
```julia
function (kernel::HostKernel)(args...; threads::CuDim=1, blocks::CuDim=1, kwargs...)
call(kernel, map(cudaconvert, args)...; threads, blocks, kwargs...)
end
```
If I simply apply this change to it:
```diff
- function (kernel::HostKernel)(args...; threads::CuDim=1, blocks::CuDim=1, kwargs...)
+ function (kernel::HostKernel)(args::Vararg{Any,M}; threads::CuDim=1, blocks::CuDim=1, kwargs...) where {M}
```
this makes it specialize and gets another 20% improvement!
This is because Vararg doesn’t automatically specialize [Performance Tips · The Julia Language](https://docs.julialang.org/en/v1/manual/performance-tips/#Be-aware-of-when-Julia-avoids-specializing) – so that `map(cudaconvert, args)...` call causes a bit of damage.
@maleadt are you okay if we upstream this change?
---
### Author:  [@Zentrik](https://discourse.julialang.org/u/Zentrik)
#### Post date: [January 15, 2024, 8:37am UTC](https://discourse.julialang.org/t/improving-gpu-performance-for-symbolic-regression/108800/11 "2024-01-15T08:37:59Z")
Does the VSCode profiler (I assume that’s what you’re using) profile GPU code? I don’t think it does.
---
### Author:  [@maleadt](https://discourse.julialang.org/u/maleadt)
#### Post date: [January 15, 2024, 1:22pm UTC](https://discourse.julialang.org/t/improving-gpu-performance-for-symbolic-regression/108800/12 "2024-01-15T13:22:46Z")
> [@MilesCranmer](#):
>
> @maleadt are you okay if we upstream this change?
Yeah, I think that would be fine!
---
### Author:  [@MilesCranmer](https://discourse.julialang.org/u/MilesCranmer)
#### Post date: [January 15, 2024, 4:24pm UTC](https://discourse.julialang.org/t/improving-gpu-performance-for-symbolic-regression/108800/13 "2024-01-15T16:24:11Z")
Cool. Done! [Fix dynamic dispatch issue by MilesCranmer · Pull Request #2235 · JuliaGPU/CUDA.jl · GitHub](https://github.com/JuliaGPU/CUDA.jl/pull/2235)
---
### Author:  [@Vijay\_Eranti](https://discourse.julialang.org/u/Vijay_Eranti)
#### Post date: [September 12, 2024, 2:51am UTC](https://discourse.julialang.org/t/improving-gpu-performance-for-symbolic-regression/108800/14 "2024-09-12T02:51:22Z")
Have you tried openai triton instead of cuda
---
### Author:  [@gdalle](https://discourse.julialang.org/u/gdalle)
#### Post date: [September 12, 2024, 5:58am UTC](https://discourse.julialang.org/t/improving-gpu-performance-for-symbolic-regression/108800/15 "2024-09-12T05:58:56Z")
Hi, welcome to the communiy! In general we avoid responding to year-old posts because otherwise everyone gets pinged and they have long moved on. Please open a new topic if you have specific unanswered questions.