GPU Acceleration

RegularizedLeastSquares.jl supports generic GPU acceleration. This means that the user can use any GPU array type that supports the GPUArrays interface. This includes CUDA.jl, AMDGPU.jl, and Metal.jl. To perform a reconstruction on the GPU, one has to load a GPU backend package such as CUDA and specify the GPU array type:

using CUDA
gpu = CuArray

At first we will look at an example of dense GPU arrays.

using RegularizedLeastSquares
A = gpu(rand(Float32, 32, 16))
x = gpu(rand(Float32, 16))
b = A*x;

Solvers adapt their states based on the type of the given measurement vector. This means that the solver will automatically switch to GPU acceleration if a GPU array is passed as the measurement vector.

solver = createLinearSolver(CGNR, A; reg = L2Regularization(0.0001), iterations=32);
x_approx = solve!(solver, b)
16-element Vector{Float32}:
 0.14148785
 0.7033924
 0.87699914
 0.95843476
 0.44417214
 0.92356414
 0.02555638
 0.52980185
 0.97269547
 0.17114109
 0.44588813
 0.2555654
 0.7436719
 0.5728242
 0.651373
 0.878649

This adaption does not include the operator. So if we want to compare with CPU result, we need to construct a new solver with a CPU operator.

solver = createLinearSolver(CGNR, Array(A); reg = L2Regularization(0.0001), iterations=32);
x_cpu = solve!(solver, Array(b))
isapprox(Array(x_approx), x_cpu)
true

Moving a CGNR Solver with Adapt.jl

A CGNR solver can also be adapted to the target backend with Adapt.jl before solving:

using Adapt
cpu_solver = createLinearSolver(CGNR, Array(A); reg = L2Regularization(0.0001), iterations=32);
gpu_solver = Adapt.adapt(gpu, cpu_solver)
x_adapted = solve!(gpu_solver, b)
isapprox(Array(x_adapted), x_cpu)
true

Each operator of the solver must support Adapt.jl. FISTA, OptISTA, POGM, ADMM, SplitBregman, DirectSolver and PseudoInverse can be adapted in the same way, FISTA, OptISTA, POGM, ADMM and SplitBregman also after some iterations; PseudoInverse moves its stored SVD without recomputing it, and DirectSolver needs lu on the target array type to solve. Other solvers do not support Adapt.adapt yet.

The built-in regularization terms can be adapted on their own. This includes the internal cache of an already used TVRegularization. Masks and image geometry stay on the host:

tv = TVRegularization(0.01f0; shape = (4, 4))
gpu_tv = Adapt.adapt(gpu, tv)
prox!(gpu_tv, gpu(rand(Float32, 16)));

Projection and plug-and-play terms adapt their function, model and input_transform with the Adapt.jl rules of these objects. A closure moves the arrays it captures, while a callable type that holds arrays needs its own Adapt.adapt_structure. Global variables are not moved, and for a target of the same backend, such as Array, arrays are shared as Adapt.jl does:

struct Weighting{V}
  w::V
end
(f::Weighting)(x) = x .* f.w
Adapt.@adapt_structure Weighting
pnp = PlugAndPlayRegularization(0.5f0; model = Weighting(rand(Float32, 4, 4)), shape = [4, 4])
gpu_pnp = Adapt.adapt(gpu, pnp)
prox!(gpu_pnp, gpu(rand(Float32, 16)));

Matrix-Free Operators

A special case is the usage of matrix-free operators. Since these operators do not have a concrete matrix representation, their GPU support depends on their implementation. Since not all multiplications within a solver approximation are in-place, the operator also needs to support the * operation and construct an appropriate result vector. For matrix-free operators based on LinearOperators.jl, this can be achieved by implementing the LinearOperators.storage_type method.

In the following, we will take another look at the CS example and execute it on the GPU. Note that for the JLArray example we chose a small phantom, since the JLArray implementation is not optimized for performance:

using ImagePhantoms, ImageGeoms
N = 32
image = shepp_logan(N, SheppLoganToft())

using Random, LinearOperatorCollection
randomIndices = shuffle(eachindex(image))
sampledIndices = sort(randomIndices[1:div(end, 3)]);

To construct the operator, we need to convert the indices to a GPU array. We also need to specify the correct storage type. In both LinearOperators.jl and LinearOperatorCollection.jl this is done with the S keyword argument.

gpu_indices = gpu(sampledIndices)
A = SamplingOp(eltype(image), pattern = gpu_indices, shape = size(image), S = typeof(b));

Let's inspect the storage type of the operator:

using LinearOperatorCollection.LinearOperators
LinearOperators.storage_type(A)
Vector{Float32} (alias for Array{Float32, 1})

Afterwards we can use the operator as usual:

b = A*vec(gpu(image));

And use it in the solver:

using RegularizedLeastSquares
reg = TVRegularization(0.01; shape=size(image))
solver = createLinearSolver(FISTA, A; reg=reg, iterations=20)
img_approx = solve!(solver,b);

To visualize the reconstructed image, we need to reshape the result vector to the correct shape and convert it to a CPU array:

img_approx = reshape(Array(img_approx),size(image))
32×32 Matrix{Float32}:
 0.00119667  0.00123909  0.00136678  …  0.00148501  0.00139946  0.00135747
 0.00149177  0.00153366  0.0016509      0.00171522  0.00158308  0.00150764
 0.00203268  0.00207058  0.00217892     0.00213827  0.0019355   0.00181378
 0.0028053   0.00286482  0.00303267     0.00271994  0.00242909  0.0022681
 0.00387512  0.00404172  0.00440299     0.0034585   0.0030365   0.00283842
 0.00530106  0.00566458  0.00623544  …  0.0044751   0.00384565  0.00356077
 0.00707608  0.00765575  0.00813243     0.00588361  0.0050134   0.00455244
 0.00919181  0.0100342   0.0112879      0.00743416  0.00655525  0.00591873
 0.0115854   0.0123185   0.0129953      0.00850243  0.00825599  0.00776055
 0.0137998   0.0145059   0.0144599      0.00880777  0.00966183  0.00998095
 ⋮                                   ⋱              ⋮           
 0.0119713   0.0124483   0.0126928      0.00736927  0.0071556   0.00679527
 0.0113331   0.0117565   0.0114168      0.00671001  0.00605895  0.0053016
 0.00920377  0.00940486  0.00899623  …  0.00617894  0.00530353  0.0044771
 0.00664182  0.00680301  0.00677675     0.00612664  0.00485158  0.00401955
 0.00469124  0.00474067  0.0048384      0.00545809  0.00429975  0.00367721
 0.00347191  0.00340635  0.00340481     0.00465513  0.00381191  0.00339763
 0.00267492  0.00260641  0.00256515     0.00406781  0.0034345   0.00312379
 0.00219097  0.00216298  0.00214904  …  0.00371033  0.00316058  0.00288522
 0.00196785  0.00197082  0.00198506     0.00353228  0.00301165  0.00274956

We will again use CairoMakie for visualization:

using CairoMakie
function plot_image(figPos, img; title = "", width = 150, height = 150)
  ax = CairoMakie.Axis(figPos; yreversed=true, title, width, height)
  hidedecorations!(ax)
  heatmap!(ax, img)
end
fig = Figure()
plot_image(fig[1,1], image, title = "Image")
samplingMask = fill(false, size(image))
samplingMask[sampledIndices] .= true
plot_image(fig[1,2], image .* samplingMask, title = "Sampled Image")
plot_image(fig[1,3], img_approx, title = "Reconstructed Image")
resize_to_layout!(fig)
fig
Example block output

This page was generated using Literate.jl.