# 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:** 1\
**Showing post:** 6

<div class="post-metadata">

**Author:** ![MilesCranmer](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/milescranmer/32/21070_2.png) [@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")

</div>

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:

![graphviz](https://global.discourse-cdn.com/julialang/original/3X/0/1/0131fd2d41de6620ca20730da3005f0f34a9e782.gif)

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.

---

_[View the full topic](https://discourse.julialang.org/t/improving-gpu-performance-for-symbolic-regression/108800)._