GPU Support
AbstractOperators.jl provides GPU compatibility through a lightweight extension model. Most operators work transparently with GPU arrays on any of the four supported backends: CUDA.jl, AMDGPU.jl, oneAPI.jl, and OpenCL.jl.
The main exceptions are:
NFFTOperators.jl, which currently supports GPU execution only with CUDA.jlWaveletOperators.jl, which currently works on CPU onlyDCTandIDCT(in FFTWOperators.jl) needs explicit loading ofAcceleratedDCTsto activate GPU support
Using GPU Arrays
Simply pass a GPU array to the mul! function or the * operator:
using AbstractOperators, CUDA # or AMDGPU, oneAPI, or OpenCL
x_gpu = CuArray(randn(Float32, 10))
y_gpu = CuArray(zeros(Float32, 9))
F = FiniteDiff(Float32, (10,)) # CPU operator, GPU arrays as input
mul!(y_gpu, F, x_gpu) # works transparentlyFor operators that carry internal arrays (like DiagOp), construct them from GPU arrays to get a GPU-typed operator:
d_gpu = CuArray(rand(Float32, 10))
D = DiagOp(d_gpu) # GPU DiagOp — threading disabled automatically
y_gpu = similar(x_gpu)
mul!(y_gpu, D, x_gpu)Threading and GPU
CPU threading (via Polyester.jl's @batch) is automatically disabled when GPU arrays are detected. This happens through the _should_thread mechanism:
- For operators with array fields (e.g.,
DiagOp), the constructor checks_should_thread(d)which returnsfalseforAbstractGPUArrays. - For GPU arrays passed to operators at construction time (e.g.,
Variation(gpu_x)), threading is disabled automatically. - The
GpuExtextension overrides_should_thread(::AbstractGPUArray) = false.
Storage Type Tracking
Every operator tracks its storage type via domain_array_type and codomain_array_type. This enables correct intermediate buffer allocation in composed operators:
using AbstractOperators, CUDA
# Create GPU operators
D = DiagOp(CuArray(rand(Float32, 5)))
domain_array_type(D) # CuArray{Float32}
codomain_array_type(D) # CuArray{Float32}
# Composition allocates GPU buffers automatically
F = FiniteDiff(CuArray{Float32}, (10,)) # array_type-aware FiniteDiff
domain_array_type(F) # CuArray{Float32}For operators without an array at construction time, pass a GPU array to the array-based constructor:
F = FiniteDiff(CuArray(ones(Float32, 5, 4)), 1) # deduces storage type from input
V = Variation(CuArray(ones(Float32, 4, 3))) # threaded=false automaticallyOr use the array_type keyword argument (for operators that support it):
Z = Zeros(Float32, (4,), Float32, (3,); array_type=CuArray)GpuExt Extension
GPU-specific overrides live in ext/GpuExt/. The extension is loaded automatically when GPUArrays.jl is loaded (directly or via any GPU backend like CUDA.jl).
The extension provides:
_should_thread(::AbstractGPUArray) = false— disables CPU threading for GPU arraysarray_type_display_string(::Type{<:AbstractGPUArray}) = "ᵍᵖᵘ"— shows ᵍᵖᵘ superscript in operator display- Variation adjoint GPU override — vectorized reshape-based stencil (no scalar indexing)
- BroadCast threading GPU override — falls back to non-threaded path
Testing with JLArrays
The test suite includes GPU tests using JLArrays.jl, a simple CPU-backed array type that mimics GPU behavior (no scalar indexing, GPUArrays interface). GPU tests are tagged :gpu:
julia --project=test -e '
using TestItemRunner
@run_package_tests filter=ti -> :gpu in ti.tags
'Backend-Specific Notes
Most operators in AbstractOperators.jl, DSPOperators.jl, and the GPU-compatible parts of FFTWOperators.jl follow the same generic GPU-array execution model and therefore work with CUDA.jl, AMDGPU.jl, oneAPI.jl, and OpenCL.jl.
NFFTOperators GPU Support (CUDA only)
NFFTOperators.jl supports GPU via the array_type keyword argument. This requires loading CUDA.jl (or any package that loads GPUArrays + Adapt):
using NFFTOperators, CUDA
image_size = (128, 128)
trajectory = rand(Float32, 2, 128, 50) .- 0.5f0
dcf = rand(Float32, 128, 50)
# Create GPU NFFT operator — ksp_buffer and dcf are CuArrays
op_gpu = NFFTOp(image_size, trajectory, dcf; array_type=CuArray, threaded=false)
x_gpu = CUDA.randn(ComplexF32, image_size)
y_gpu = op_gpu * x_gpu # forward: image → k-space (on GPU)
img_rec = op_gpu' * y_gpu # adjoint: k-space → image (on GPU)Note: GPU NFFT requires CUDA.jl because it relies on cuFFT plans. JLArrays (used in tests) does not support FFT plans and cannot be used with NFFTOperators GPU.
The trajectory k is kept on CPU (as required by NFFT.jl's GPU plan); only the computation buffers (ksp_buffer, dcf) and input/output arrays are on GPU.
FFTWOperators DCT/IDCT (AcceleratedDCTs)
DCT and IDCT are CPU-only unless AcceleratedDCTs is loaded. When AcceleratedDCTs is available and imported, the package activates GPU support for those operators through its extension:
using AbstractOperators, FFTWOperators, CUDA
import AcceleratedDCTs # explicitly import to activate GPU support for DCT/IDCT
x_gpu = CUDA.randn(Float32, 64)
dct_op = DCT(x_gpu)
y_gpu = dct_op * x_gpu
idct_op = IDCT(x_gpu)
x_rec_gpu = idct_op * y_gpuIf AcceleratedDCTs is not imported, DCT and IDCT continue to use the CPU FFTW implementation, and will not work with GPU arrays. In that case, wrap them with CpuOperatorWrapper to use in GPU pipelines:
using AbstractOperators, FFTWOperators, CUDA
x_gpu = CUDA.randn(Float32, 64)
dct_op = CpuOperatorWrapper(DCT(Float32, (64,)); array_type = CuArray{Float32}) # CPU operator wrapped for GPU use
y_gpu = dct_op * x_gpu # GPU in → CPU DCT → GPU outWaveletOperators CPU-only status
WaveletOperators.jl currently relies on CPU execution. Its operators do not yet support GPU arrays, so wavelet transforms should remain on CPU or be wrapped explicitly as CPU operators when building mixed CPU/GPU pipelines.
CpuOperatorWrapper
For operators that do not natively support GPU arrays (e.g., FFTWOperators DCT, custom CPU-only operators), use CpuOperatorWrapper. This wrapper preallocates CPU buffers for the operator's domain and codomain, allowing GPU arrays to be passed in and out while the computation happens on CPU.
AbstractOperators.OperatorWrapper — Type
OperatorWrapper{Op, DB, CB, DS, CS}A wrapper that adapts any operator's storage type independently of the wrapped operator. Useful for making CPU operators transparently usable with GPU arrays.
When mul!(y, wrapper, x) is called the wrapper:
- Copies
xto the preallocated domain buffer (same storage as the inner op) - Executes the inner operator on its native buffers
- Copies the result to
y
GPU mul! overrides are provided by the GpuExt extension (requires loading GPUArrays).
Construction
OperatorWrapper(op::AbstractOperator; array_type = Array)array_type sets the outer storage type reported by the wrapper (used by check). This is independent of the wrapped operator's storage type, which governs the internal CPU buffers.
Notes
- Thread safety: each wrapper has its own buffers; use
copy_operatorfor parallel use.
Example
julia> using AbstractOperators
julia> op = FiniteDiff(Float64, (8,), 1)
δx ℝ^8 -> ℝ^7
julia> w = OperatorWrapper(op)
CPU[δx] ℝ^8 -> ℝ^7
julia> size(w) == size(op)
trueAbstractOperators.OperatorWrapper — Method
OperatorWrapper(op::AbstractOperator; array_type = Array)Wrap op, preallocating internal CPU buffers from its domain/codomain. array_type sets the outer storage type used in check (default: Array).
Example:
using AbstractOperators, CUDA
# Any CPU operator
op = FiniteDiff(Float32, (64,)) # or any FFTWOperators, etc.
# Wrap it — preallocates CPU buffers for domain and codomain
wrapper = CpuOperatorWrapper(op; array_type = CuArray{Float32}) # specify GPU array type for buffers
x_gpu = CUDA.randn(Float32, 64)
y_gpu = similar(x_gpu, 63)
mul!(y_gpu, wrapper, x_gpu) # GPU in → CPU compute → GPU out
mul!(x_gpu, wrapper', y_gpu) # GPU in → CPU adjoint → GPU out