GPU computing in the Julia language

This first tutorial introduces GPU programming in Julia, from high-level array operations down to writing your own kernels. Everything we do in the second tutorial sits on top of the mechanisms shown here: building optimal control models with ExaModels.jl, and solving them with MadNLP.jl.

Julia

Both tutorials are written in Julia. You do not need to know the language to follow them. Familiarity with arrays and functions is enough, and every construct we use is explained where it appears. If you would like more, the language is documented in depth at docs.julialang.org, and this introduction from the JuMP documentation is a shorter on-ramp.

Julia is JIT compiled, so the first call is slow

Julia compiles specialized machine code the first time you call a function with a given argument type (on the CPU or the GPU). That first call therefore pays for the compiler as well as the computation, and every later call is fast.

Practical consequences today, all of them normal and none of them a fault of your session:

• A cell can take seconds on its first run and milliseconds afterwards.
• Re-run a cell before concluding anything is slow.
• Never time a first call. Every benchmark in this workshop uses the @btime macro (from BenchmarkTools.jl, introduced below), which warms up and repeats automatically, so compilation never lands in a reported number.
• The effect is largest for GPU code, where the first call compiles a GPU kernel as well. In tutorial 2 the first optimization solve of the session takes a couple of minutes for this reason; it is paid once.

One function, timed twice:

square(x) = x * x

t_first = @elapsed square(2.0)
t_second = @elapsed square(3.0)
(first_call_s = t_first, second_call_s = t_second,
 first_call_slower_by = round(Int, t_first / t_second))

(first_call_s = 0.005407312, second_call_s = 5.892e-6, first_call_slower_by = 918)

The first call included compiling square for Float64, which costs hundreds to thousands of times the multiplication itself. The second ran the already-compiled code. Every "why is this cell slow the first time?" moment today is the same effect at a larger scale, up to and including the first optimization solve in tutorial 2.

To measure a running time reliably we use BenchmarkTools.jl throughout this workshop. Its @btime macro warms up first, so compilation is excluded, then runs the expression many times and reports the minimum:

using BenchmarkTools

@btime square(2.0);

  11.163 ns (0 allocations: 0 bytes)

Nanoseconds, against the milliseconds of the raw first call above. When the time is needed as a value rather than a printout, for example to compute a speedup, we use the sibling macro @belapsed. It returns the minimum time in seconds.

A short ✏️ exercise appears later on, as a blank cell marked "your turn".

Arrays on the GPU

Using a GPU does not always mean writing GPU kernels. Much of what CUDA.jl offers is available from a high level. Standard array operations (arithmetic, broadcasting, map, reduce, linear algebra) are already extended to GPU arrays, so most GPU code in this workshop looks like ordinary Julia array code. We write one kernel by hand near the end, for the one pattern broadcasting does not cover, not because you will usually need to.

Julia's GPU support is organized under JuliaGPU. We use NVIDIA GPUs here, programmed through CUDA.jl:

using CUDA

A first check. This should show the GPU you were allocated:

CUDA.name(CUDA.device())

┌ Warning: Your Quadro GV100 GPU (compute capability 7.0) is not fully supported by CUDA 12.9.0.
│ Some functionality may be broken. Ensure you are using the latest version of CUDA.jl in combination with an up-to-date NVIDIA driver.
│ If that does not help, please file an issue to add support for the latest CUDA toolkit.
└ @ CUDACore ~/.julia/packages/CUDACore/OIFhX/lib/cudadrv/state.jl:231

"Quadro GV100"

In Julia, a vector on the CPU is allocated as

x_cpu = zeros(8)

8-element Vector{Float64}:
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0

The GPU counterpart is explicit. CUDA.zeros allocates in GPU memory and returns a CuArray:

x_gpu = CUDA.zeros(Float64, 8)

8-element CUDACore.CuArray{Float64, 1, CUDACore.DeviceMemory}:
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0

Info

The dimension-only constructors CUDA.zeros(n), CUDA.rand(n) and CUDA.randn(n) default to Float32, a heritage of machine learning where single precision is the norm. Their Base counterparts zeros(n) and rand(n) give Float64. Constructors that see your data keep its type: CuArray(randn(8)) and CUDA.fill(1.0, 8) are Float64. Optimal control is not machine learning. An interior-point method depends on residuals near 1e-8, which Float32 cannot represent against numbers of size one, so in this workshop we write Float64 explicitly whenever a constructor does not see a value.

We are sharing a GPU today, so keep your arrays small

Several of us are on the same device, so every array in this notebook is deliberately modest and you should size yours the same way. The arithmetic is one multiplication: a Float64 costs 8 bytes, so an n-element vector is 8n bytes and an n × n matrix is 8n². The largest object here is the 2048 × 2048 matrix in the linear-algebra section, at 2048² × 8 = 34 MB. Running the whole notebook leaves about 145 MB of arrays live, and the memory pool never holds more than about 300 MB, so several of us fit on one device comfortably. CUDA.memory_info() returns the free and total bytes on the device at any time, so you can check before you allocate something big.

Data moves between the two worlds by conversion. CuArray(a) takes an array that lives in CPU memory and copies it to the device:

y_gpu = CuArray(randn(8))

8-element CUDACore.CuArray{Float64, 1, CUDACore.DeviceMemory}:
  0.09536474393523404
  1.9666541644356506
  0.37609245655694584
  1.5056621262875372
  1.085749145767373
 -0.19997564359563946
 -0.7320842787223049
 -2.481177522771381

and Array(a) copies a GPU array back to the host:

y_cpu = Array(y_gpu)

8-element Vector{Float64}:
  0.09536474393523404
  1.9666541644356506
  0.37609245655694584
  1.5056621262875372
  1.085749145767373
 -0.19997564359563946
 -0.7320842787223049
 -2.481177522771381

Transfers over the PCIe bus are slow compared to GPU memory bandwidth. The recipe for performance is: move data to the GPU once, keep the whole computation resident there, and bring back only the small result. That is what "GPU-resident interior-point methods" will mean later in the workshop.

Array programming with broadcasting

The easiest way to compute on a CuArray is Julia's broadcasting syntax, the same dot syntax you would use on the CPU. If you know MATLAB, you already know it: the dots are MATLAB's elementwise notation (.*, .^), generalized so that any Julia function broadcasts with a dot. The dotted assignment .= writes in place: it fills the array you already have, allocating nothing.

x_gpu .= 1.0

8-element CUDACore.CuArray{Float64, 1, CUDACore.DeviceMemory}:
 1.0
 1.0
 1.0
 1.0
 1.0
 1.0
 1.0
 1.0

Allocate the output once, then fill it in place with a whole dotted expression:

z_gpu = CUDA.zeros(Float64, 8)
z_gpu .= 2.0 .* x_gpu .+ sin.(y_gpu)

8-element CUDACore.CuArray{Float64, 1, CUDACore.DeviceMemory}:
 2.095220261582409
 2.922666121844938
 2.3672888494706656
 2.9978795177936486
 2.8846529519344326
 1.8013545401617352
 1.3315786634825684
 1.3865552507447

Each broadcast expression compiles to a single GPU kernel: Julia fuses the chain of dotted operations, so 2.0 .* x .+ sin.(y) reads x and y once and writes z once, rather than materializing intermediates.

This lets you write CPU/GPU-compatible functions: keep the body in array operations, take the array as an argument, and one definition serves both devices. The mechanism behind it is Julia's multiple dispatch: a function is compiled per argument type, so the same definition becomes a CPU loop when handed an Array and a GPU kernel when handed a CuArray. Define one function:

f(v) = 2.0 .* v .+ 1.0

f (generic function with 1 method)

Hand it a CPU array:

typeof(f(randn(4)))

Vector{Float64} (alias for Array{Float64, 1})

Hand the same function a GPU array:

typeof(f(CuArray(randn(4))))

CUDACore.CuArray{Float64, 1, CUDACore.DeviceMemory}

Same function, two compiled methods, selected by the type of the input. This mechanism is what will let one optimization model run on either device unchanged in the next tutorial.

What not to do: scalar indexing

What you must not do on the GPU is access elements one at a time:

try
    z_gpu[1]  # scalar indexing, an error in non-interactive code
catch err
    @error "Scalar indexing failed" typeof(err)
end

┌ Error: Scalar indexing failed
│   typeof(err) = ErrorException
└ @ Main.var"##277" ~/git/teaching/ifac2026-workshop/notebook/notebooks/1-gpu-computing.ipynb:4

A single-element read forces a synchronization and a PCIe round-trip; a loop of them is catastrophically slow, so CUDA.jl disallows it outside the REPL. The message to take away: on the GPU you operate on whole arrays, never on scalars. If a computation seems to need element-by-element logic, it either fits map/mapreduce, or it needs a custom kernel (below).

Timing a loop against a broadcast

CUDA.@allowscalar unblocks scalar indexing, so we can do the thing you should never do and time it. The same update written two ways:

function axpy_loop!(y, x, a)
    for i in eachindex(y)
        y[i] += a * x[i]
    end
    return y
end

function axpy_bcast!(y, x, a)
    y .+= a .* x
    return y
end

axpy_bcast! (generic function with 1 method)

The same data on both devices:

# Ten thousand elements: the scalar loop makes its point at this size in
# about half a second; every element access is its own round-trip to the
# device, so the cost grows with the array and a million elements would
# take over half a minute.
n_loop = 10_000

y_cs = fill(2.0, n_loop)
x_cs = fill(1.0, n_loop)

y_gs = CUDA.fill(2.0, n_loop)
x_gs = CUDA.fill(1.0, n_loop);

@belapsed takes care of the warm-up and the repeats. It also brings one piece of notation with it: variables that live in the global scope are interpolated into the benchmarked expression with a $, so that the timing measures the operation itself and not the global-variable lookup around it.

First, loop against broadcast on the CPU:

t_cpu_loop = @belapsed axpy_loop!($y_cs, $x_cs, 3.0)
t_cpu_bcast_small = @belapsed axpy_bcast!($y_cs, $x_cs, 3.0)

(cpu_loop_ms = 1e3 * t_cpu_loop, cpu_broadcast_ms = 1e3 * t_cpu_bcast_small)

(cpu_loop_ms = 0.0023505555555555554, cpu_broadcast_ms = 0.0023517777777777777)

The same speed: Julia compiles the loop to the same machine code, so loops are not slow here the way they are in Python or MATLAB.

Now the same pair on the GPU. One new rule applies: GPU operations launch asynchronously, so the benchmarked expression must contain a CUDA.@sync or we time the launch rather than the computation. The scalar loop needs no @sync, because every single-element access already waits for the device.

t_gpu_bcast_small = @belapsed CUDA.@sync axpy_bcast!($y_gs, $x_gs, 3.0)

# `@elapsed`, not `@belapsed`: this one is slow enough that repeating it
# would cost minutes, and far too slow for the noise to matter.
t_gpu_loop = @elapsed CUDA.@allowscalar axpy_loop!(y_gs, x_gs, 3.0)

(gpu_broadcast_ms = 1e3 * t_gpu_bcast_small, gpu_loop_ms = 1e3 * t_gpu_loop)

(gpu_broadcast_ms = 0.022199, gpu_loop_ms = 388.434381)

The loop is not just slower than the broadcast. It is far slower than the same loop on the CPU:

(gpu_loop_vs_gpu_broadcast = round(Int, t_gpu_loop / t_gpu_bcast_small),
 gpu_loop_vs_cpu_loop = round(Int, t_gpu_loop / t_cpu_loop))

(gpu_loop_vs_gpu_broadcast = 17498, gpu_loop_vs_cpu_loop = 165252)

Each y[i] is a separate round-trip to the device, with the GPU's 5120 cores idle while it happens; the broadcast issues one kernel for all of them at once. Putting data on a GPU and then touching it element by element is worse than not using the GPU at all, and that is why CUDA.jl makes you write @allowscalar to do it.

At this size even the GPU broadcast trails the CPU one: launching a kernel costs tens of microseconds whatever it computes, and ten thousand elements is not enough arithmetic to repay it. The batched simulation at the end of this tutorial shows the crossover once the arrays are large enough.

map and reduce

Broadcasting is not the only high-level construct already extended for GPU arrays. First, map itself: map(f, a) applies the function f to every element of a and collects the results. On an ordinary CPU array:

map(x -> 2.0 * x^2 + 1.0, [0.0, 1.0, 2.0])

3-element Vector{Float64}:
 1.0
 3.0
 9.0

The x -> 2.0 * x^2 + 1.0 is an anonymous function, defined inline where it is used, without a name. Handed a CuArray, the same map runs on the device, and the function you pass is compiled into a GPU kernel for you:

u_gpu = map(x -> 2.0 * x^2 + 1.0, y_gpu)

8-element CUDACore.CuArray{Float64, 1, CUDACore.DeviceMemory}:
  1.0181888687716656
  8.735457204984174
  1.2828910717580764
  5.534036877073415
  3.35770241506916
  1.0799805160629805
  2.0718947823047147
 13.312483799011854

map allocates a fresh array for the result. map!(f, dest, a) writes into one you already own instead. This is the same in-place idea as .= above, and the form to use inside a loop:

dest_gpu = similar(y_gpu)
map!(x -> 2.0 * x^2 + 1.0, dest_gpu, y_gpu)

8-element CUDACore.CuArray{Float64, 1, CUDACore.DeviceMemory}:
  1.0181888687716656
  8.735457204984174
  1.2828910717580764
  5.534036877073415
  3.35770241506916
  1.0799805160629805
  2.0718947823047147
 13.312483799011854

The general parallel aggregation is a reduction: reduce(op, a) combines all elements with op in a parallel tree on the device, and mapreduce(f, op, a) applies f to each element on the way in. It fuses the two, so the mapped array is never built:

(reduce(+, u_gpu), mapreduce(abs2, +, u_gpu))

(36.392635535036035, 303.57179112253704)

Nearly every aggregation you already use is a special case of this, and so it is already available on the GPU. sum is reduce(+, ·), while maximum, minimum, extrema, count, any and all are the same tree with a different operator. There is an in-place family as well: sum!(dest, A) and maximum!(dest, A) reduce along a dimension into an array you already own.

(sum(abs2, z_gpu), maximum(abs, z_gpu), count(>(0.0), z_gpu), extrema(z_gpu))

(42.785001337235464, 2.9978795177936486, 8, (1.3315786634825684, 2.9978795177936486))

---

✏️ Exercise: reductions in an interior-point solve

Your turn. Every iteration of the solver you will use in the next tutorial computes a handful of reductions over vectors this shape. Here are three of them, on stand-in data:

c  = CUDA.randn(Float64, 10^5)          # constraint residual
xv = CUDA.rand(Float64, 10^5) .+ 0.1    # primal variables, positive
zv = CUDA.rand(Float64, 10^5) .+ 0.1    # their duals, positive
dx = CUDA.randn(Float64, 10^5)          # a step direction

Write each with mapreduce, in one pass and with no temporary array:

1. Primal infeasibility, the largest absolute residual, maximum |c|.
2. Complementarity, the largest product max |x_i z_i|. Two arrays at once: mapreduce takes several and walks them together.
3. The fraction-to-boundary step: the largest α keeping x + α dx positive, that is min(-x_i / dx_i) over the entries where dx_i < 0, and Inf if there are none. Return Inf for entries that do not constrain the step, since Inf is the identity for min.

# your code here

Solution: reductions-solution.

---

Dense linear algebra

CuArray plugs into Julia's standard linear algebra, backed by NVIDIA's cuBLAS and cuSOLVER libraries, with the same functions under the same names:

using LinearAlgebra

m = 2048
A_gpu = CUDA.randn(Float64, m, m)
v_gpu = CUDA.randn(Float64, m);

A matrix–vector product runs on cuBLAS:

w_gpu = A_gpu * v_gpu

2048-element CUDACore.CuArray{Float64, 1, CUDACore.DeviceMemory}:
 -64.88370268482025
  22.843138075297993
 -55.60042164255302
 -52.670906423581165
   8.375534701915473
 -34.0423149407661
  62.21900772546391
  13.296348194920718
  40.39426767545733
 -47.97813462800943
   ⋮
 -45.67855599941788
  45.38351038680202
  64.29313855701335
  67.46315071379313
   1.9967568135027562
 -59.06693017553814
 -31.94832520630457
  11.288628787339547
  39.26093096639351

Like .= for broadcasting, linear algebra has in-place variants that write into an output you already allocated. mul!(w, A, v) computes A * v and stores it in w, allocating nothing. In a hot loop, such as every iteration of an optimization solver, this is the form you want:

mul!(w_gpu, A_gpu, v_gpu);

Factorizations work too. Build a symmetric positive definite matrix, the kind that sits at the core of an interior-point iteration:

S_gpu = A_gpu * A_gpu' + m * I;

and factor it with Cholesky on the device, via cuSOLVER:

F = cholesky(S_gpu);

F now holds the factor; \ solves the linear system with it:

b_gpu = CUDA.randn(Float64, m)
x_sol = F \ b_gpu
norm(S_gpu * x_sol - b_gpu)

5.262402628704362e-14

And the in-place solve: ldiv!(x, F, b) writes the solution into x without allocating a fresh vector. Factor once, then solve many right-hand sides in place:

x_sol2 = similar(b_gpu)
ldiv!(x_sol2, F, b_gpu)
norm(S_gpu * x_sol2 - b_gpu)

5.262402628704362e-14

Sparse linear algebra

The matrices in this workshop are mostly sparse, and sparsity is a different data structure rather than a special case of a dense array. Julia holds one in a SparseMatrixCSC; the GPU counterpart lives in CUDA.jl's CUSPARSE submodule, and moving one across is a conversion like any other:

using SparseArrays
using CUDA.CUSPARSE

S_cpu = sprand(4096, 4096, 0.001) + 10I
S_sparse = CuSparseMatrixCSR(S_cpu)

4096×4096 cuSPARSE.CuSparseMatrixCSR{Float64, Int32} with 20798 stored entries:
⎡⣷⣯⣻⣿⣾⣏⣿⣿⣿⣶⣻⣛⣽⣶⡟⣿⣫⢿⣿⣿⣞⣶⣿⣿⣿⣾⣻⣿⣿⢮⣿⣿⣿⣿⣿⣿⣗⣿⡿⣻⎤
⎢⢹⣿⣿⣿⣿⢿⣯⣿⡽⣿⡿⣯⣯⣟⣿⣿⣿⢿⣿⣴⡾⣿⣿⣿⣿⣿⣿⠿⣿⣿⣾⣿⣟⣯⣿⣻⣜⣿⣿⣷⎥
⎢⣿⣿⣿⣿⣿⣿⣽⣿⡿⣾⣿⣿⣿⣿⣾⡿⣿⣿⣶⣯⣿⣽⡿⣿⣻⣻⡿⢿⣿⣾⣽⣧⣽⣿⣾⣿⣿⡿⣿⣿⎥
⎢⣿⣿⣟⣿⣶⣿⣿⣿⣿⢻⣷⣿⡟⣯⣿⣿⣿⣷⣷⣷⣯⣿⣿⣿⣿⣿⣷⡿⣷⡽⣿⢛⣿⣿⣯⡿⣿⣹⣿⣯⎥
⎢⣮⣻⣿⣯⣿⣿⣯⢻⣿⣿⣷⣿⣿⣽⣿⣿⣿⣿⣵⣟⣷⣾⣻⣟⣿⣛⣿⣿⣿⣟⡿⣿⣿⣿⣿⣿⣧⣽⣺⣿⎥
⎢⢽⣿⣿⣾⣿⣿⣿⣿⣿⣿⢿⣿⣷⢿⣟⣿⣷⣿⣿⣿⣿⣻⣯⣿⣟⡿⣿⣯⣿⣿⣿⣾⣿⣿⣞⣽⣿⣿⣿⡿⎥
⎢⣼⣿⡿⣟⣿⣯⢻⣿⣿⣷⡿⣿⣿⣿⣺⣿⣿⣿⣻⣿⣿⣿⣾⢿⣟⣿⣻⣿⣿⣷⣿⣿⣿⣯⣿⣿⣿⣾⣿⣿⎥
⎢⡷⡿⣿⣿⣾⣿⣿⣿⣿⡽⣿⣟⣿⣾⣿⣿⣿⡵⣿⣿⣽⣿⣿⣿⣿⣻⣿⣼⣯⣿⣼⣿⣿⣿⣿⣯⣷⠿⣞⠇⎥
⎢⣽⣿⣳⣶⣿⣿⣿⣿⣿⢿⣾⣯⣾⣿⡼⣿⣿⣾⣟⡿⣿⣿⣏⢿⡿⣾⢿⣿⣿⣿⣿⣿⡿⣿⣟⣿⣿⢿⣿⣏⎥
⎢⣽⣾⣿⣿⣽⡻⣶⣿⣿⣿⡻⣿⣿⣿⣿⣿⣿⣿⣿⣿⣟⣿⣯⢾⣿⣟⣿⣿⣿⣿⣾⣿⣿⣿⣿⣿⣟⣷⣿⣯⎥
⎢⣿⣏⣽⣿⡿⣿⡿⣿⣿⢿⣿⣿⢽⣟⣿⣿⣿⣿⢽⣿⣿⣷⢿⣯⣫⣿⣿⣿⡿⣿⣿⢻⣿⣻⣿⣽⣽⣿⣻⢟⎥
⎢⣯⣿⣾⣟⣟⣿⣿⣿⣏⣿⣾⣿⣿⣿⣿⣻⣿⣿⣿⣿⣿⣿⣿⣿⣟⣿⣿⣿⢷⣯⣿⢟⣿⣿⣻⣿⣿⣿⢿⢯⎥
⎢⢿⣻⣿⣿⡿⣿⣾⣿⢿⢷⢿⣿⣿⣿⣿⣿⣿⣾⣿⣽⣿⣿⡿⣿⣿⣿⣏⡿⣞⣿⣿⡿⣟⣿⣿⣿⣽⣛⣷⢿⎥
⎢⢿⣿⣿⣿⣿⣿⢷⣿⣿⡿⣿⣿⣾⣷⣿⣿⣿⣿⡿⢿⣿⣿⣿⣿⣿⣿⣿⣿⡿⣿⣿⣽⠿⣟⣿⣿⣿⢻⣾⣿⎥
⎢⣿⣿⣿⣻⣿⢾⣿⣿⢿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣾⣟⣿⡿⣿⣽⢿⣿⣿⣿⢿⡿⢯⢿⣿⣫⣺⣿⣻⣣⎥
⎢⣻⣿⡷⡯⣾⣷⣿⣟⣽⣿⣿⣿⣿⣯⣿⣿⣿⣿⣿⣿⣟⣯⢿⣿⣯⣷⢿⣿⣽⣿⣿⣽⣿⣾⣿⡻⣿⣿⣽⡿⎥
⎢⣾⣗⣿⣿⣻⣿⣽⡾⣵⢿⣽⡻⣿⣿⣿⢿⣿⢿⣿⣿⣿⣯⣿⣿⣿⣻⣽⣷⣧⣵⣿⢿⣻⣿⣿⣿⣿⣻⣿⡶⎥
⎢⣿⣿⣽⣿⣿⣿⣿⠺⢿⣻⣿⣯⣿⣿⣿⡿⣿⡿⣟⣿⣿⣿⣽⣯⣟⣯⢿⣗⣻⣫⣿⣿⣿⣿⣿⣿⣾⣻⢾⣽⎥
⎢⡹⣿⣿⣿⣿⣿⣷⣿⣿⣟⣷⣿⣿⣿⣿⣿⣯⣿⣷⣿⣻⣻⣷⣿⣿⣿⣿⣿⣻⣿⣿⣿⣿⣿⣿⣿⣿⣽⣿⣿⎥
⎣⢯⣿⣛⣷⣯⡿⡿⣻⣿⣿⡷⣿⣿⣿⡿⣿⢿⢿⢿⣽⢿⣿⣿⣿⢿⣾⠯⢿⣿⣷⢼⣶⣾⣿⣿⣿⣿⣿⣿⣷⎦

Sparse matrix-vector products are available and fast. mul! dispatches to CUSPARSE, so the in-place form you already know works unchanged:

u_in = CUDA.randn(Float64, 4096)
u_out = similar(u_in)

mul!(u_out, S_sparse, u_in)

norm(u_out)

651.4104323060327

That covers the products. Sparse factorization is the harder problem, and a large part of what made GPU interior-point methods difficult historically: a factorization has to follow the index structure and the fill-in it creates as it proceeds, which parallelizes far less naturally than a dense block does.

We do not solve that here. The answer is cuDSS, NVIDIA's sparse direct solver, reached from Julia through CUDSS.jl. You will not call it directly today: in tutorial 2 MadNLPGPU does it for you, and the lecture covers what it is doing.

The rest of the standard library

Past broadcasting, reduce and linear algebra there is one more family: operations that are neither elementwise nor aggregations nor matrix algebra. Generating, sorting and searching all have native GPU implementations, behind the generic names you already know. A million uniform random numbers, generated on the device:

r_gpu = CUDA.rand(Float64, 10^6)

1000000-element CUDACore.CuArray{Float64, 1, CUDACore.DeviceMemory}:
 0.9867469156399651
 0.9599872603561677
 0.798389246599728
 0.29736306006965935
 0.7264618850191891
 0.993604398588835
 0.17179565180790518
 0.9841504690477176
 0.48209606599770066
 0.64994716046974
 ⋮
 0.20824577729007315
 0.7133344294870403
 0.506937781269057
 0.08883092875613868
 0.5242117539022915
 0.4228399878752925
 0.4953145575222588
 0.48036091615625404
 0.5604508300595523

Sorting them runs a parallel GPU sort, a different algorithm from the CPU one, and you get it by calling sort:

s_gpu = sort(r_gpu)

1000000-element CUDACore.CuArray{Float64, 1, CUDACore.DeviceMemory}:
 2.8170808324956553e-7
 4.1885748275349854e-7
 1.9094513840633986e-6
 2.0971991704921145e-6
 2.8260188916218176e-6
 3.572706507914969e-6
 3.849673913369767e-6
 3.8951056716141075e-6
 7.125311125066869e-6
 7.230111392531047e-6
 ⋮
 0.9999908783014664
 0.9999911131974988
 0.9999930606952037
 0.9999932013571005
 0.9999939318330999
 0.9999947632961126
 0.9999949054788455
 0.9999989770398894
 0.9999993262933122

And searching is parallel too. The indices of every element above 0.99:

findall(>(0.99), r_gpu)

9977-element CUDACore.CuArray{Int64, 1, CUDACore.DeviceMemory}:
      6
    338
    352
    658
    716
    778
    781
    788
    793
   1238
      ⋮
 999149
 999612
 999689
 999713
 999727
 999811
 999882
 999938
 999988

sortperm, cumsum, accumulate, reverse, logical-mask indexing, and the dense factorizations (lu, cholesky, qr) are native as well. Not everything is. unique, for one, falls back to slow element-at-a-time access. When in doubt, try your operation on a CuArray and watch for scalar-indexing complaints before relying on it in a hot loop.

---

✏️ Exercise: Monte-Carlo π

Your turn. Estimate π on the GPU: draw N = 10^6 uniform points in the unit square with CUDA.rand, and use mapreduce to count how many land inside the quarter disk x^2 + y^2 ≤ 1. Four times that fraction estimates π. Write it as a single fused reduction, with no temporary array of distances or of booleans.

# your code here

Solution: montecarlo-pi-solution.

---

Kernels with KernelAbstractions.jl

Broadcasting covers elementwise operations, but not everything is elementwise. A stencil is the classic exception: each output element reads several neighboring inputs. The discrete 1-D Laplacian is the second derivative on a grid with spacing h, and it is the building block of diffusion equations and PDE-constrained optimization:

$$ (Lx)_i = \frac{x_{i-1} - 2x_i + x_{i+1}}{h^2}, \qquad i = 1, \dots, n, $$

with Dirichlet boundary conditions. The solution is pinned to zero just outside the grid, so x_0 = x_{n+1} = 0.

Output i needs inputs i-1, i, and i+1, so it is not a broadcast. It is a job for a kernel: a function executed by many GPU threads at once, each handling one output element. We write it with KernelAbstractions.jl (KA), which expresses a kernel once and runs it on NVIDIA, AMD, and Intel GPUs, or on multithreaded CPUs. ExaModels and MadNLP are built on KA, and that is what makes them vendor-agnostic.

using KernelAbstractions
const KA = KernelAbstractions

@kernel function laplacian!(y, @Const(x), h2)
    i = @index(Global, Linear)
    n = length(x)
    if i == 1                                  # Dirichlet: x_0 = 0
        y[i] = (-2x[i] + x[i+1]) / h2
    elseif i == n                              # Dirichlet: x_{n+1} = 0
        y[i] = (x[i-1] - 2x[i]) / h2
    else
        y[i] = (x[i-1] - 2x[i] + x[i+1]) / h2
    end
end

laplacian! (generic function with 4 methods)

Three pieces of KA syntax: @kernel marks the function as a kernel, @index(Global, Linear) gives each thread its own index, which is the thread/block bookkeeping you would otherwise do by hand, and @Const declares an argument read-only. One thread handles one grid point, and the if is how a kernel expresses a boundary condition.

To run it: instantiating the kernel on a backend decides where it executes, and ndrange says how many threads to launch, here one per grid point. Our test function is a sine wave. It vanishes at both ends of the grid, so it satisfies the Dirichlet condition, and its second derivative is minus itself:

n_g = 2^20
h = 2pi / (n_g + 1)
x_wave = CuArray(sin.(h .* (1:n_g)))
y_wave = CUDA.zeros(Float64, n_g);

backend = KA.get_backend(x_wave)  # CUDABackend for a CuArray (ROCBackend for AMD, oneAPIBackend for Intel)
laplacian!(backend)(y_wave, x_wave, h^2; ndrange = n_g)
KA.synchronize(backend)

So the result should equal -x. The 1/h^2 amplifies floating-point roundoff, so expect agreement to a few digits rather than to machine precision. That is the arithmetic, not the kernel:

maximum(abs, y_wave .+ x_wave)

2.7257937598723636e-5

And the same kernel, unchanged, on the CPU, with one different backend:

x_wave_c = sin.(h .* (1:n_g))
y_wave_c = zeros(n_g);

laplacian!(CPU())(y_wave_c, x_wave_c, h^2; ndrange = n_g)
KA.synchronize(CPU())
maximum(abs, y_wave_c .+ x_wave_c)

2.7257937598723636e-5

Batched simulation

To close, a computation shaped like the ones in the rest of the workshop: simulate a batch of damped pendulums with different initial angles, all at once. The picture first, since the setting is easier seen than said:

Every instance obeys the same equations and differs only in its state, so a step for the whole batch is one operation on an array rather than a loop over pendulums. That is the shape a GPU is built for, and the same shape the optimization problems in the next tutorial take.

One pendulum obeys the second-order ODE

$$ \ddot{\theta} = -\frac{g}{L}\,\sin\theta - c\,\dot{\theta}, $$

or as a first-order system in the angle $\theta$ and the angular velocity $\omega = \dot{\theta}$:

$$ \dot{\theta} = \omega, \qquad \dot{\omega} = -\frac{g}{L}\,\sin\theta - c\,\omega, $$

with $g/L = 9.81$ and damping $c = 0.1$. The simplest way to simulate it is the semi-implicit Euler method, stepping forward with time step $\Delta t$ (the angle update uses the freshly updated velocity):

$$ \omega \leftarrow \omega - \Delta t\,(9.81\,\sin\theta + 0.1\,\omega), \qquad \theta \leftarrow \theta + \Delta t\,\omega. $$

The code below does this for the whole batch at once. The state is stored as arrays over the batch, so one Euler step is a broadcast: the same operation on every instance, which is the pattern GPUs are built for:

function simulate!(θ, ω, dt, nsteps)
    for _ in 1:nsteps
        ω .-= dt .* (9.81 .* sin.(θ) .+ 0.1 .* ω)
        θ .+= dt .* ω
    end
    return θ
end

batch = 100_000
θ0 = collect(range(-π, π; length = batch));

θ_c = copy(θ0)
ω_c = zeros(batch)
t_cpu_sim = @belapsed simulate!($θ_c, $ω_c, 1e-3, 1000)

θ_g = CuArray(θ0)
ω_g = CUDA.zeros(Float64, batch)
t_gpu_sim = @belapsed CUDA.@sync simulate!($θ_g, $ω_g, 1e-3, 1000)

(t_cpu = t_cpu_sim, t_gpu = t_gpu_sim, speedup = t_cpu_sim / t_gpu_sim)

(t_cpu = 2.134275244, t_gpu = 0.032262902, speedup = 66.15261218597136)

One code, both devices, and a substantial speedup on the GPU. Two caveats. The CPU number is a single thread: Julia broadcasts do not multithread on their own, so a multicore CPU closes part of the gap, though the memory-bandwidth arithmetic still favors the GPU. Each Euler step also launched two kernels, so kernel-launch overhead is why the speedup grows with batch size. Shrink the batch far enough and the launches dominate the arithmetic. That crossover is worth remembering when you size a problem for a GPU.

---

This notebook was generated using Literate.jl.