From 9f1fb26cfea4c38859df2b718929bf88e91a4e3d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Fri, 7 Aug 2026 18:51:12 +0200 Subject: [PATCH 1/2] Add GPU support for LBFGSData --- src/lbfgs.jl | 151 ++++++++++++++++++++++++++++------------------- src/utilities.jl | 8 +-- 2 files changed, 94 insertions(+), 65 deletions(-) diff --git a/src/lbfgs.jl b/src/lbfgs.jl index e856e070..d59dd039 100644 --- a/src/lbfgs.jl +++ b/src/lbfgs.jl @@ -1,7 +1,14 @@ export LBFGSOperator, InverseLBFGSOperator, diag, diag! -"A data type to hold information relative to LBFGS operators." -mutable struct LBFGSData{T, I <: Integer} +"A data type to hold information relative to LBFGS operators. + +`V<:AbstractVector{T}` parameterises the *n-sized* working buffers, so the +same operator can live on the CPU (`V = Vector{T}`) or on the GPU +(`V = CuVector{T}`, etc.). The small `mem`-sized bookkeeping arrays +(`ys`, `α`, `norm_b`) stay on the CPU on purpose — they are indexed by +scalar `k` inside `lbfgs_multiply` and would force scalar-getindex on a +GPU array." +mutable struct LBFGSData{T, I <: Integer, V <: AbstractVector{T}} const mem::I const scaling::Bool scaling_factor::T @@ -9,22 +16,23 @@ mutable struct LBFGSData{T, I <: Integer} σ₂::T σ₃::T opnorm_upper_bound::T # Upper bound for the operator norm ‖Bₖ‖₂ ≤ ‖B₀‖₂ + ∑ᵢ ‖bᵢ‖₂² - const s::Vector{Vector{T}} - const y::Vector{Vector{T}} + const s::Vector{V} + const y::Vector{V} const ys::Vector{T} const α::Vector{T} - const a::Vector{Vector{T}} - const b::Vector{Vector{T}} + const a::Vector{V} + const b::Vector{V} const norm_b::Vector{T} insert::I - const Ax::Vector{T} + const Ax::V const shifted_p::Matrix{T} # Temporary matrix used in the computation solve_shifted_system! const shifted_v::Vector{T} const shifted_u::Vector{T} end function LBFGSData( - T::Type, + ::Type{T}, + ::Type{V}, n::I; mem::I = 5, scaling::Bool = true, @@ -32,9 +40,10 @@ function LBFGSData( inverse::Bool = true, σ₂::Float64 = 0.99, σ₃::Float64 = 10.0, -) where {I <: Integer} +) where {T, I <: Integer, V <: AbstractVector{T}} maxmem = max(mem, 1) - LBFGSData{T, I}( + _zeros(::Type{V}, n) where {V} = fill!(V(undef, n), zero(eltype(V))) + LBFGSData{T, I, V}( maxmem, scaling, convert(T, 1), @@ -42,25 +51,30 @@ function LBFGSData( convert(T, σ₂), convert(T, σ₃), convert(T, 1), - [zeros(T, n) for _ = 1:maxmem], - [zeros(T, n) for _ = 1:maxmem], + [_zeros(V, n) for _ = 1:maxmem], + [_zeros(V, n) for _ = 1:maxmem], zeros(T, maxmem), inverse ? zeros(T, maxmem) : zeros(T, 0), - inverse ? Vector{Vector{T}}(undef, 0) : [zeros(T, n) for _ = 1:maxmem], - inverse ? Vector{Vector{T}}(undef, 0) : [zeros(T, n) for _ = 1:maxmem], + inverse ? V[] : [_zeros(V, n) for _ = 1:maxmem], + inverse ? V[] : [_zeros(V, n) for _ = 1:maxmem], inverse ? Vector{T}(undef, 0) : zeros(T, maxmem), 1, - Vector{T}(undef, n), + V(undef, n), Array{T}(undef, (n, 2 * maxmem)), Vector{T}(undef, 2 * maxmem), Vector{T}(undef, n), ) end +# Backwards-compatible: default to CPU `Vector{T}`. +LBFGSData(T::Type, n::I; kwargs...) where {I <: Integer} = + LBFGSData(T, Vector{T}, n; kwargs...) + LBFGSData(n::I; kwargs...) where {I <: Integer} = LBFGSData(Float64, n; kwargs...) "A type for limited-memory BFGS approximations." -mutable struct LBFGSOperator{T, I <: Integer, F, Ft, Fct} <: AbstractQuasiNewtonOperator{T} +mutable struct LBFGSOperator{T, I <: Integer, F, Ft, Fct, V <: AbstractVector{T}} <: + AbstractQuasiNewtonOperator{T} const nrow::I const ncol::I const symmetric::Bool @@ -69,7 +83,7 @@ mutable struct LBFGSOperator{T, I <: Integer, F, Ft, Fct} <: AbstractQuasiNewton const tprod!::Ft # apply the transpose operator to a vector const ctprod!::Fct # apply the transpose conjugate operator to a vector const inverse::Bool - const data::LBFGSData{T, I} + const data::LBFGSData{T, I, V} nprod::I ntprod::I nctprod::I @@ -84,25 +98,26 @@ LBFGSOperator{T}( tprod!::Ft, ctprod!::Fct, inverse::Bool, - data::LBFGSData{T, I}, -) where {T, I <: Integer, F, Ft, Fct} = LBFGSOperator{T, I, F, Ft, Fct}( - nrow, - ncol, - symmetric, - hermitian, - prod!, - tprod!, - ctprod!, - inverse, - data, - 0, - 0, - 0, -) + data::LBFGSData{T, I, V}, +) where {T, I <: Integer, F, Ft, Fct, V <: AbstractVector{T}} = + LBFGSOperator{T, I, F, Ft, Fct, V}( + nrow, + ncol, + symmetric, + hermitian, + prod!, + tprod!, + ctprod!, + inverse, + data, + 0, + 0, + 0, + ) has_args5(op::LBFGSOperator) = true isallocated5(op::LBFGSOperator) = true -storage_type(op::LBFGSOperator{T}) where {T} = Vector{T} +storage_type(op::LBFGSOperator{T, I, F, Ft, Fct, V}) where {T, I, F, Ft, Fct, V} = V """ InverseLBFGSOperator(T, n, [mem=5; scaling=true]) @@ -110,10 +125,15 @@ storage_type(op::LBFGSOperator{T}) where {T} = Vector{T} Construct a limited-memory BFGS approximation in inverse form. If the type `T` is omitted, then `Float64` is used. """ -function InverseLBFGSOperator(T::Type, n::I; kwargs...) where {I <: Integer} +function InverseLBFGSOperator( + ::Type{T}, + ::Type{V}, + n::I; + kwargs..., +) where {T, V <: AbstractVector{T}, I <: Integer} kwargs = Dict(kwargs) delete!(kwargs, :inverse) - lbfgs_data = LBFGSData(T, n; inverse = true, kwargs...) + lbfgs_data = LBFGSData(T, V, n; inverse = true, kwargs...) function lbfgs_multiply( res::AbstractVector, @@ -158,6 +178,8 @@ function InverseLBFGSOperator(T::Type, n::I; kwargs...) where {I <: Integer} return LBFGSOperator{T}(n, n, true, true, prod!, prod!, prod!, true, lbfgs_data) end +InverseLBFGSOperator(T::Type, n::Integer; kwargs...) = + InverseLBFGSOperator(T, Vector{T}, n; kwargs...) InverseLBFGSOperator(n::Integer; kwargs...) = InverseLBFGSOperator(Float64, n; kwargs...) """ @@ -166,10 +188,15 @@ InverseLBFGSOperator(n::Integer; kwargs...) = InverseLBFGSOperator(Float64, n; k Construct a limited-memory BFGS approximation in forward form. If the type `T` is omitted, then `Float64` is used. """ -function LBFGSOperator(T::Type, n::I; kwargs...) where {I <: Integer} +function LBFGSOperator( + ::Type{T}, + ::Type{V}, + n::I; + kwargs..., +) where {T, V <: AbstractVector{T}, I <: Integer} kwargs = Dict(kwargs) delete!(kwargs, :inverse) - lbfgs_data = LBFGSData(T, n; inverse = false, kwargs...) + lbfgs_data = LBFGSData(T, V, n; inverse = false, kwargs...) function lbfgs_multiply( res::AbstractVector, @@ -206,14 +233,16 @@ function LBFGSOperator(T::Type, n::I; kwargs...) where {I <: Integer} return LBFGSOperator{T}(n, n, true, true, prod!, prod!, prod!, false, lbfgs_data) end +LBFGSOperator(T::Type, n::Integer; kwargs...) = + LBFGSOperator(T, Vector{T}, n; kwargs...) LBFGSOperator(n::I; kwargs...) where {I <: Integer} = LBFGSOperator(Float64, n; kwargs...) function push_common!( - op::LBFGSOperator{T, I, F1, F2, F3}, - s::Vector{T}, - y::Vector{T}, + op::LBFGSOperator{T, I, F1, F2, F3, V}, + s::AbstractVector{T}, + y::AbstractVector{T}, ys::T, -) where {T, I, F1, F2, F3} +) where {T, I, F1, F2, F3, V} # op.counters.updates += 1 data = op.data insert = data.insert @@ -268,10 +297,10 @@ The third and fourth calling sequences are used in inverse LBFGS updating in con where α is the most recent steplength and g the gradient used when solving `d=-Hg`. """ function push!( - op::LBFGSOperator{T, I, F1, F2, F3}, - s::Vector{T}, - y::Vector{T}, -) where {T, I, F1, F2, F3} + op::LBFGSOperator{T, I, F1, F2, F3, V}, + s::AbstractVector{T}, + y::AbstractVector{T}, +) where {T, I, F1, F2, F3, V} if op.data.damped return push!(op, s, y, similar(s)) end @@ -288,11 +317,11 @@ function push!( end function push!( - op::LBFGSOperator{T, I, F1, F2, F3}, - s::Vector{T}, - y::Vector{T}, - Bs::Vector{T}, -) where {T, I, F1, F2, F3} + op::LBFGSOperator{T, I, F1, F2, F3, V}, + s::AbstractVector{T}, + y::AbstractVector{T}, + Bs::AbstractVector{T}, +) where {T, I, F1, F2, F3, V} if !op.data.damped error("This push! should be used for damped operators") elseif op.inverse @@ -322,13 +351,13 @@ function push!( end function push!( - op::LBFGSOperator{T, I, F1, F2, F3}, - s::Vector{T}, - y::Vector{T}, + op::LBFGSOperator{T, I, F1, F2, F3, V}, + s::AbstractVector{T}, + y::AbstractVector{T}, α::T, - g::Vector{T}, - Bs::Vector{T}, -) where {T, I, F1, F2, F3} + g::AbstractVector{T}, + Bs::AbstractVector{T}, +) where {T, I, F1, F2, F3, V} if !op.data.damped error("This push! should be used for damped operators") elseif !op.inverse @@ -358,12 +387,12 @@ function push!( end function push!( - op::LBFGSOperator{T, I, F1, F2, F3}, - s::Vector{T}, - y::Vector{T}, + op::LBFGSOperator{T, I, F1, F2, F3, V}, + s::AbstractVector{T}, + y::AbstractVector{T}, α::T, - g::Vector{T}, -) where {T, I, F1, F2, F3} + g::AbstractVector{T}, +) where {T, I, F1, F2, F3, V} push!(op, s, y, α, g, similar(g)) end diff --git a/src/utilities.jl b/src/utilities.jl index 343c413a..da5ab324 100644 --- a/src/utilities.jl +++ b/src/utilities.jl @@ -206,10 +206,10 @@ Erway, J. B., Jain, V., & Marcia, R. F. Shifted L-BFGS Systems. Optimization Met """ function solve_shifted_system!( x::AbstractVector{T}, - B::LBFGSOperator{T, I, F1, F2, F3}, + B::LBFGSOperator{T, I, F1, F2, F3, V}, b::AbstractVector{T}, σ::T, -) where {T, I, F1, F2, F3} +) where {T, I, F1, F2, F3, V} if σ < 0 throw(ArgumentError("σ must be nonnegative")) end @@ -280,9 +280,9 @@ ldiv!(x, B, b) function ldiv!( x::AbstractVector{T}, - B::LBFGSOperator{T, I, F1, F2, F3}, + B::LBFGSOperator{T, I, F1, F2, F3, V}, b::AbstractVector{T}, -) where {T, I, F1, F2, F3} +) where {T, I, F1, F2, F3, V} # Call solve_shifted_system! with σ = 0 solve_shifted_system!(x, B, b, T(0.0)) return x From 86feac4492355ca2b315f3174841bed21e318ae6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Sat, 19 Sep 2026 10:38:55 +0200 Subject: [PATCH 2/2] Add tests --- .buildkite/pipeline.yml | 2 + src/lbfgs.jl | 54 +++++++++++-------- src/utilities.jl | 21 ++++---- test/gpu/amdgpu.jl | 3 ++ test/gpu/jlarrays.jl | 3 ++ test/gpu/metal.jl | 3 ++ test/gpu/nvidia.jl | 3 ++ test/gpu/test_lbfgs.jl | 115 ++++++++++++++++++++++++++++++++++++++++ test/runtests.jl | 1 + 9 files changed, 175 insertions(+), 30 deletions(-) create mode 100644 test/gpu/test_lbfgs.jl diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 74fc8542..7b404974 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -11,6 +11,7 @@ steps: Pkg.add("CUDA") Pkg.instantiate() include("test/gpu/test_S_kwarg.jl") + include("test/gpu/test_lbfgs.jl") include("test/gpu/nvidia.jl")' timeout_in_minutes: 30 @@ -27,5 +28,6 @@ steps: Pkg.add("AMDGPU") Pkg.instantiate() include("test/gpu/test_S_kwarg.jl") + include("test/gpu/test_lbfgs.jl") include("test/gpu/amdgpu.jl")' timeout_in_minutes: 30 diff --git a/src/lbfgs.jl b/src/lbfgs.jl index d59dd039..a3c78855 100644 --- a/src/lbfgs.jl +++ b/src/lbfgs.jl @@ -4,11 +4,12 @@ export LBFGSOperator, InverseLBFGSOperator, diag, diag! `V<:AbstractVector{T}` parameterises the *n-sized* working buffers, so the same operator can live on the CPU (`V = Vector{T}`) or on the GPU -(`V = CuVector{T}`, etc.). The small `mem`-sized bookkeeping arrays +(`V = CuVector{T}`, etc.). `M<:AbstractMatrix{T}` stores the shifted-solve +workspace on the same backend. The small `mem`-sized bookkeeping arrays (`ys`, `α`, `norm_b`) stay on the CPU on purpose — they are indexed by scalar `k` inside `lbfgs_multiply` and would force scalar-getindex on a GPU array." -mutable struct LBFGSData{T, I <: Integer, V <: AbstractVector{T}} +mutable struct LBFGSData{T, I <: Integer, V <: AbstractVector{T}, M <: AbstractMatrix{T}} const mem::I const scaling::Bool scaling_factor::T @@ -25,9 +26,9 @@ mutable struct LBFGSData{T, I <: Integer, V <: AbstractVector{T}} const norm_b::Vector{T} insert::I const Ax::V - const shifted_p::Matrix{T} # Temporary matrix used in the computation solve_shifted_system! + const shifted_p::M # Temporary matrix used in the computation solve_shifted_system! const shifted_v::Vector{T} - const shifted_u::Vector{T} + const shifted_u::V end function LBFGSData( @@ -43,7 +44,9 @@ function LBFGSData( ) where {T, I <: Integer, V <: AbstractVector{T}} maxmem = max(mem, 1) _zeros(::Type{V}, n) where {V} = fill!(V(undef, n), zero(eltype(V))) - LBFGSData{T, I, V}( + Ax = V(undef, n) + shifted_p = similar(Ax, n, 2 * maxmem) + LBFGSData{T, I, V, typeof(shifted_p)}( maxmem, scaling, convert(T, 1), @@ -59,22 +62,28 @@ function LBFGSData( inverse ? V[] : [_zeros(V, n) for _ = 1:maxmem], inverse ? Vector{T}(undef, 0) : zeros(T, maxmem), 1, - V(undef, n), - Array{T}(undef, (n, 2 * maxmem)), + Ax, + shifted_p, Vector{T}(undef, 2 * maxmem), - Vector{T}(undef, n), + V(undef, n), ) end # Backwards-compatible: default to CPU `Vector{T}`. -LBFGSData(T::Type, n::I; kwargs...) where {I <: Integer} = - LBFGSData(T, Vector{T}, n; kwargs...) +LBFGSData(T::Type, n::I; kwargs...) where {I <: Integer} = LBFGSData(T, Vector{T}, n; kwargs...) LBFGSData(n::I; kwargs...) where {I <: Integer} = LBFGSData(Float64, n; kwargs...) "A type for limited-memory BFGS approximations." -mutable struct LBFGSOperator{T, I <: Integer, F, Ft, Fct, V <: AbstractVector{T}} <: - AbstractQuasiNewtonOperator{T} +mutable struct LBFGSOperator{ + T, + I <: Integer, + F, + Ft, + Fct, + V <: AbstractVector{T}, + M <: AbstractMatrix{T}, +} <: AbstractQuasiNewtonOperator{T} const nrow::I const ncol::I const symmetric::Bool @@ -83,7 +92,7 @@ mutable struct LBFGSOperator{T, I <: Integer, F, Ft, Fct, V <: AbstractVector{T} const tprod!::Ft # apply the transpose operator to a vector const ctprod!::Fct # apply the transpose conjugate operator to a vector const inverse::Bool - const data::LBFGSData{T, I, V} + const data::LBFGSData{T, I, V, M} nprod::I ntprod::I nctprod::I @@ -98,9 +107,9 @@ LBFGSOperator{T}( tprod!::Ft, ctprod!::Fct, inverse::Bool, - data::LBFGSData{T, I, V}, -) where {T, I <: Integer, F, Ft, Fct, V <: AbstractVector{T}} = - LBFGSOperator{T, I, F, Ft, Fct, V}( + data::LBFGSData{T, I, V, M}, +) where {T, I <: Integer, F, Ft, Fct, V <: AbstractVector{T}, M <: AbstractMatrix{T}} = + LBFGSOperator{T, I, F, Ft, Fct, V, M}( nrow, ncol, symmetric, @@ -120,10 +129,12 @@ isallocated5(op::LBFGSOperator) = true storage_type(op::LBFGSOperator{T, I, F, Ft, Fct, V}) where {T, I, F, Ft, Fct, V} = V """ + InverseLBFGSOperator(T, V, n; mem=5, scaling=true) InverseLBFGSOperator(T, n, [mem=5; scaling=true]) InverseLBFGSOperator(n, [mem=5; scaling=true]) Construct a limited-memory BFGS approximation in inverse form. If the type `T` -is omitted, then `Float64` is used. +is omitted, then `Float64` is used. Pass `V <: AbstractVector{T}` to select +the storage backend (default: `Vector{T}`). """ function InverseLBFGSOperator( ::Type{T}, @@ -183,10 +194,12 @@ InverseLBFGSOperator(T::Type, n::Integer; kwargs...) = InverseLBFGSOperator(n::Integer; kwargs...) = InverseLBFGSOperator(Float64, n; kwargs...) """ + LBFGSOperator(T, V, n; mem=5, scaling=true) LBFGSOperator(T, n; [mem=5, scaling=true]) LBFGSOperator(n; [mem=5, scaling=true]) Construct a limited-memory BFGS approximation in forward form. If the type `T` -is omitted, then `Float64` is used. +is omitted, then `Float64` is used. Pass `V <: AbstractVector{T}` to select +the storage backend (default: `Vector{T}`). """ function LBFGSOperator( ::Type{T}, @@ -233,8 +246,7 @@ function LBFGSOperator( return LBFGSOperator{T}(n, n, true, true, prod!, prod!, prod!, false, lbfgs_data) end -LBFGSOperator(T::Type, n::Integer; kwargs...) = - LBFGSOperator(T, Vector{T}, n; kwargs...) +LBFGSOperator(T::Type, n::Integer; kwargs...) = LBFGSOperator(T, Vector{T}, n; kwargs...) LBFGSOperator(n::I; kwargs...) where {I <: Integer} = LBFGSOperator(Float64, n; kwargs...) function push_common!( @@ -402,7 +414,7 @@ end Extract the diagonal of a L-BFGS operator in forward mode. """ function diag(op::LBFGSOperator{T}) where {T} - d = Vector{T}(undef, op.nrow) + d = storage_type(op)(undef, op.nrow) diag!(op, d) end diff --git a/src/utilities.jl b/src/utilities.jl index da5ab324..c78c9478 100644 --- a/src/utilities.jl +++ b/src/utilities.jl @@ -221,27 +221,30 @@ function solve_shifted_system!( @. x = x_0 * b max_i = 2 * data.mem - sign_i = 1 + # Add each positive rank-one term before its downdate. Subtracting + # a*a' first makes the initial system singular when σ == 0. + sign_i = -1 for i = 1:max_i j = (i + 1) ÷ 2 - k = mod(insert + j - 1, data.mem) + 1 + k = mod(insert + j - 2, data.mem) + 1 data.shifted_u .= ((sign_i == -1) ? data.b[k] : data.a[k]) - @. data.shifted_p[:, i] = x_0 * data.shifted_u + p_i = view(data.shifted_p, :, i) + @. p_i = x_0 * data.shifted_u - sign_t = 1 + sign_t = -1 for t = 1:(i - 1) - c0 = dot(view(data.shifted_p, :, t), data.shifted_u) + p_t = view(data.shifted_p, :, t) + c0 = dot(p_t, data.shifted_u) c1 = sign_t .* data.shifted_v[t] c2 = c1 * c0 - view(data.shifted_p, :, i) .+= c2 .* view(data.shifted_p, :, t) + p_i .+= c2 .* p_t sign_t = -sign_t end - data.shifted_v[i] = 1 / (1 - sign_i * dot(data.shifted_u, view(data.shifted_p, :, i))) - x .+= - sign_i * data.shifted_v[i] * (view(data.shifted_p, :, i)' * b) .* view(data.shifted_p, :, i) + data.shifted_v[i] = 1 / (1 - sign_i * dot(data.shifted_u, p_i)) + x .+= sign_i * data.shifted_v[i] * dot(p_i, b) .* p_i sign_i = -sign_i end return x diff --git a/test/gpu/amdgpu.jl b/test/gpu/amdgpu.jl index d12f3b92..30362506 100644 --- a/test/gpu/amdgpu.jl +++ b/test/gpu/amdgpu.jl @@ -18,3 +18,6 @@ using LinearOperators, AMDGPU @testset "AMDGPU S kwarg" test_S_kwarg(arrayType = ROCArray) end + +AMDGPU.allowscalar(false) +test_lbfgs_backend(arrayType = ROCArray) diff --git a/test/gpu/jlarrays.jl b/test/gpu/jlarrays.jl index 4ace84cc..5fe4a095 100644 --- a/test/gpu/jlarrays.jl +++ b/test/gpu/jlarrays.jl @@ -1 +1,4 @@ test_S_kwarg(arrayType = JLArray) + +JLArrays.allowscalar(false) +test_lbfgs_backend(arrayType = JLArray) diff --git a/test/gpu/metal.jl b/test/gpu/metal.jl index 22e3d224..6b8264c0 100644 --- a/test/gpu/metal.jl +++ b/test/gpu/metal.jl @@ -1,2 +1,5 @@ using Metal test_S_kwarg(arrayType = MtlArray, notMetal = false) + +Metal.allowscalar(false) +test_lbfgs_backend(arrayType = MtlArray) diff --git a/test/gpu/nvidia.jl b/test/gpu/nvidia.jl index 63bf7326..2cc98347 100644 --- a/test/gpu/nvidia.jl +++ b/test/gpu/nvidia.jl @@ -21,3 +21,6 @@ using LinearOperators, CUDA, CUDA.CUSPARSE, CUDA.CUSOLVER @testset "Nvidia S kwarg" test_S_kwarg(arrayType = CuArray) end + +CUDA.allowscalar(false) +test_lbfgs_backend(arrayType = CuArray) diff --git a/test/gpu/test_lbfgs.jl b/test/gpu/test_lbfgs.jl new file mode 100644 index 00000000..618d6979 --- /dev/null +++ b/test/gpu/test_lbfgs.jl @@ -0,0 +1,115 @@ +using Test, LinearAlgebra, LinearOperators + +function test_lbfgs_backend(; arrayType) + T = Float32 + n = 8 + x = T.(1:n) ./ n + v = arrayType(x) + V = typeof(v) + rtol = 5.0f-4 + atol = 5.0f-5 + + @testset "LBFGS with $arrayType" begin + @testset "mem=$mem, scaling=$scaling" for mem in (0, 3), scaling in (false, true) + B = LBFGSOperator(T, V, n; mem = mem, scaling = scaling) + H = InverseLBFGSOperator(T, V, n; mem = mem, scaling = scaling) + maxmem = max(mem, 1) + for op in (B, H) + @test LinearOperators.storage_type(op) == V + @test op.data.mem == maxmem + @test all(w -> w isa V, op.data.s) + @test all(w -> w isa V, op.data.y) + @test op.data.Ax isa V + @test op.data.shifted_p isa typeof(similar(v, n, 2 * maxmem)) + @test size(op.data.shifted_p) == (n, 2 * maxmem) + @test op * v isa V + @test Array(op * v) ≈ x + # Rejected curvature must not change the approximation. + push!(op, v, -v) + push!(op, v, zero(v)) + @test op.data.insert == 1 + @test Array(op * v) ≈ x + end + + pairs = Tuple{Vector{T}, Vector{T}}[] + # Check empty, partially filled, full, and wrapped history. + for k = 0:(maxmem + 2) + if k > 0 + s = T[sin(i + k) for i = 1:n] + y = (one(T) .+ T.(1:n) ./ n) .* s + push!(pairs, (s, y)) + length(pairs) > maxmem && popfirst!(pairs) + push!(B, arrayType(s), arrayType(y)) + push!(H, arrayType(s), arrayType(y)) + end + # Independent dense BFGS reference using the retained history. + γ = scaling && k > 0 ? dot(last(pairs)...) / dot(last(pairs)[2], last(pairs)[2]) : one(T) + dense = Matrix{T}(I, n, n) / γ + for (s, y) in pairs + Bs = dense * s + dense += y * y' / dot(s, y) - Bs * Bs' / dot(s, Bs) + end + for (op, expected) in ((B, dense * x), (H, dense \ x)) + @test op.data.insert == mod(k, maxmem) + 1 + @test Array(op * v) ≈ expected rtol = rtol atol = atol + out = copy(v) + mul!(out, op, v, T(2), T(0.5)) + @test Array(out) ≈ 2 .* expected .+ T(0.5) .* x rtol = rtol atol = atol + fill!(out, T(NaN)) + mul!(out, op, v, one(T), zero(T)) + @test Array(out) ≈ expected rtol = rtol atol = atol + end + @test diag(B) isa V + @test Array(diag(B)) ≈ diag(dense) rtol = rtol atol = atol + d = similar(v) + @test diag!(B, d) === d + @test Array(d) ≈ diag(dense) rtol = rtol atol = atol + sol = similar(v) + for σ in (zero(T), T(0.5)) + @test solve_shifted_system!(sol, B, v, σ) === sol + @test Array(sol) ≈ (dense + σ * I) \ x rtol = rtol atol = atol + @test Array(B * sol + σ * sol) ≈ x rtol = rtol atol = atol + end + @test ldiv!(sol, B, v) === sol + @test Array(sol) ≈ dense \ x rtol = rtol atol = atol + end + @test_throws LinearOperators.LinearOperatorException diag(H) + @test_throws ArgumentError solve_shifted_system!(similar(v), B, v, -one(T)) + for op in (B, H) + reset!(op) + @test op.data.insert == 1 + @test op.data.scaling_factor == one(T) + @test Array(op * v) ≈ x + end + @test Array(diag(B)) ≈ ones(T, n) + end + + @testset "damped updates" for constructor in (LBFGSOperator, InverseLBFGSOperator) + op = constructor(T, V, n; mem = 3, damped = true) + cpu = constructor(T, n; mem = 3, damped = true) + for k = 1:5 + g = T[cos(i + k) for i = 1:n] + s = -(cpu * g) + # Exercise both lower and upper curvature damping. + y = (isodd(k) ? T(0.001) : T(100)) .* (-g) + if op.inverse + if isodd(k) + push!(op, arrayType(s), arrayType(y), one(T), arrayType(g)) + push!(cpu, s, copy(y), one(T), g) + else + push!(op, arrayType(s), arrayType(y), one(T), arrayType(g), similar(v)) + push!(cpu, s, copy(y), one(T), g, similar(x)) + end + elseif isodd(k) + push!(op, arrayType(s), arrayType(y)) + push!(cpu, s, y) + else + push!(op, arrayType(s), arrayType(y), similar(v)) + push!(cpu, s, y, similar(x)) + end + @test Array(op * v) ≈ cpu * x rtol = rtol atol = atol + @test op.data.insert == cpu.data.insert + end + end + end +end diff --git a/test/runtests.jl b/test/runtests.jl index 24265704..31f9ff02 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -19,6 +19,7 @@ include("test_chainrules.jl") include("test_solve_shifted_system.jl") include("test_shifted_operator.jl") include("gpu/test_S_kwarg.jl") +include("gpu/test_lbfgs.jl") include("gpu/jlarrays.jl") if Sys.isapple() && occursin("arm64", Sys.MACHINE) include("gpu/metal.jl")