Compressed LBFGS (forward) operator

.buildkite/pipeline.yml

@@ -9,7 +9,8 @@ steps:
     command: |
       julia --color=yes --project -e '
       using Pkg
-      Pkg.add("CUDA")
+      # Pkg.add("CUDA")
+      Pkg.add(url="https://github.com/JuliaGPU/CUDA.jl", rev="master")
       Pkg.instantiate()
       include("test/gpu/nvidia.jl")'
     timeout_in_minutes: 30

Project.toml

@@ -7,12 +7,14 @@ FastClosures = "9aa1b823-49e4-5ca5-8b0f-3971ec8bab6a"
 LDLFactorizations = "40e66cde-538c-5869-a4ad-c39174c6795b"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
+Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 
 [compat]
 FastClosures = "0.2, 0.3"
 LDLFactorizations = "0.9, 0.10"
+Requires = "1.3"
 TimerOutputs = "^0.5"
 julia = "^1.6.0"
 

docs/make.jl

@@ -1,4 +1,5 @@
 using Documenter, LinearOperators
+using LinearOperators.ModCompressedLBFGSOperator
 
 makedocs(
   modules = [LinearOperators],

docs/src/reference.md

@@ -13,5 +13,5 @@ Pages = ["reference.md"]
 ```
 
 ```@autodocs
-Modules = [LinearOperators]
+Modules = [LinearOperators, ModCompressedLBFGSOperator]
 ```

src/compressed_lbfgs.jl

@@ -0,0 +1,221 @@
+module ModCompressedLBFGSOperator
+#=
+Compressed LBFGS implementation from:
+    REPRESENTATIONS OF QUASI-NEWTON MATRICES AND THEIR USE IN LIMITED MEMORY METHODS
+    Richard H. Byrd, Jorge Nocedal and Robert B. Schnabel (1994)
+    DOI: 10.1007/BF01582063
+
+Implemented by Paul Raynaud (supervised by Dominique Orban)
+=#
+
+using LinearAlgebra, LinearAlgebra.BLAS
+using Requires
+
+default_matrix_type(; T::DataType=Float64) = Matrix{T}
+default_vector_type(; T::DataType=Float64) = Vector{T}
+
+@init begin
+  @require CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" begin
+    default_matrix_type(; T::DataType=Float64) = CUDA.CuMatrix{T, CUDA.Mem.DeviceBuffer}
+    default_vector_type(; T::DataType=Float64) = CUDA.CuVector{T, CUDA.Mem.DeviceBuffer}
+  end
+  # this scheme may be extended to other GPU modules
+end
+
+export CompressedLBFGSOperator
+export default_matrix_type, default_vector_type
+
+"""
+    CompressedLBFGSOperator{T, M<:AbstractMatrix{T}, V<:AbstractVector{T}}
+
+A LBFGS limited-memory operator.
+It represents a linear application Rⁿˣⁿ, considering at most `mem` BFGS updates.
+This implementation considers the bloc matrices reoresentation of the BFGS (forward) update.
+It follows the algorithm described in [REPRESENTATIONS OF QUASI-NEWTON MATRICES AND THEIR USE IN LIMITED MEMORY METHODS](https://link.springer.com/article/10.1007/BF01582063) from Richard H. Byrd, Jorge Nocedal and Robert B. Schnabel (1994).
+This operator considers several fields directly related to the bloc representation of the operator:
+- `mem`: the maximal memory of the operator;
+- `n`: the dimension of the linear application;
+- `k`: the current memory's size of the operator;
+- `α`: scalar for `B₀ = α I`;
+- `Sₖ`: retain the `k`-th last vectors `s` from the updates parametrized by `(s,y)`;
+- `Yₖ`: retain the `k`-th last vectors `y` from the updates parametrized by `(s,y)`;;
+- `Dₖ`: a diagonal matrix mandatory to perform the linear application and to form the matrix;
+- `Lₖ`: a lower diagonal mandatory to perform the linear application and to form the matrix.
+In addition to this structures which are circurlarly update when `k` reaches `mem`, we consider other intermediate data structures renew at each update:
+- `chol_matrix`: a matrix required to store a Cholesky factorization of a Rᵏˣᵏ matrix;
+- `intermediate_1`: a R²ᵏˣ²ᵏ matrix;
+- `intermediate_2`: a R²ᵏˣ²ᵏ matrix;
+- `inverse_intermediate_1`: a R²ᵏˣ²ᵏ matrix;
+- `inverse_intermediate_2`: a R²ᵏˣ²ᵏ matrix;
+- `intermediary_vector`: a vector ∈ Rᵏ to store intermediate solutions;
+- `sol`: a vector ∈ Rᵏ to store intermediate solutions;
+This implementation is designed to work either on CPU or GPU.
+"""
+mutable struct CompressedLBFGSOperator{T, M<:AbstractMatrix{T}, V<:AbstractVector{T}}
+  mem::Int # memory of the operator
+  n::Int # vector size
+  k::Int # k ≤ mem, active memory of the operator
+  α::T # B₀ = αI
+  Sₖ::M # gather all sₖ₋ₘ 
+  Yₖ::M # gather all yₖ₋ₘ 
+  Dₖ::Diagonal{T,V} # mem * mem
+  Lₖ::LowerTriangular{T,M} # mem * mem
+
+  chol_matrix::M # 2m * 2m
+  intermediate_diagonal::Diagonal{T,V} # mem * mem
+  intermediate_1::UpperTriangular{T,M} # 2m * 2m
+  intermediate_2::LowerTriangular{T,M} # 2m * 2m
+  inverse_intermediate_1::UpperTriangular{T,M} # 2m * 2m
+  inverse_intermediate_2::LowerTriangular{T,M} # 2m * 2m
+  intermediary_vector::V # 2m
+  sol::V # mem
+end
+
+function columnshift!(A::AbstractMatrix{T}; direction::Int=-1, indicemax::Int=size(A)[1]) where T
+  map(i-> view(A,:,i+direction) .= view(A,:,i), 1-direction:indicemax)
+  return A
+end
+
+function vectorshift!(v::AbstractVector{T}; direction::Int=-1, indicemax::Int=length(v)) where T
+  view(v, 1:indicemax+direction) .= view(v,1-direction:indicemax)
+  return v
+end
+
+"""
+    CompressedLBFGSOperator(n::Int; [T=Float64, mem=5], gpu:Bool)
+
+A implementation of a LBFGS operator (forward), representing a `nxn` linear application.
+It considers at most `k` BFGS iterates, and fit the architecture depending if it is launched on a CPU or a GPU.
+"""
+function CompressedLBFGSOperator(n::Int; mem::Int=5, T=Float64, M=default_matrix_type(; T), V=default_vector_type(; T))
+  α = (T)(1)
+  k = 0  
+  Sₖ = M(undef, n, mem)
+  Yₖ = M(undef, n, mem)
+  Dₖ = Diagonal(V(undef, mem))
+  Lₖ = LowerTriangular(M(undef, mem, mem))
+  Lₖ.data .= (T)(0)
+
+  chol_matrix = M(undef, mem, mem)
+  intermediate_diagonal = Diagonal(V(undef, mem))
+  intermediate_1 = UpperTriangular(M(undef, 2*mem, 2*mem))
+  intermediate_2 = LowerTriangular(M(undef, 2*mem, 2*mem))
+  inverse_intermediate_1 = UpperTriangular(M(undef, 2*mem, 2*mem))
+  inverse_intermediate_2 = LowerTriangular(M(undef, 2*mem, 2*mem))
+  intermediary_vector = V(undef, 2*mem)
+  sol = V(undef, 2*mem)
+  return CompressedLBFGSOperator{T,M,V}(mem, n, k, α, Sₖ, Yₖ, Dₖ, Lₖ, chol_matrix, intermediate_diagonal, intermediate_1, intermediate_2, inverse_intermediate_1, inverse_intermediate_2, intermediary_vector, sol)
+end
+
+function Base.push!(op::CompressedLBFGSOperator{T,M,V}, s::V, y::V) where {T,M,V<:AbstractVector{T}}
+  if op.k < op.mem # still some place in the structures
+    op.k += 1
+    view(op.Sₖ, :, op.k) .= s
+    view(op.Yₖ, :, op.k) .= y
+    view(op.Dₖ.diag, op.k) .= dot(s, y)
+    mul!(view(op.Lₖ.data, op.k, 1:op.k-1), transpose(view(op.Yₖ, :, 1:op.k-1)), view(op.Sₖ, :, op.k) )
+  else # k == mem update circurlarly the intermediary structures
+    columnshift!(op.Sₖ; indicemax=op.k)
+    columnshift!(op.Yₖ; indicemax=op.k)
+    # op.Dₖ .= circshift(op.Dₖ, (-1, -1))
+    vectorshift!(op.Dₖ.diag; indicemax=op.k)
+    view(op.Sₖ, :, op.k) .= s
+    view(op.Yₖ, :, op.k) .= y
+    view(op.Dₖ.diag, op.k) .= dot(s, y)
+
+    map(i-> view(op.Lₖ, i:op.mem-1, i-1) .= view(op.Lₖ, i+1:op.mem, i), 2:op.mem)
+    mul!(view(op.Lₖ.data, op.k, 1:op.k-1), transpose(view(op.Yₖ, :, 1:op.k-1)), view(op.Sₖ, :, op.k) )
+  end
+
+  # step 4 and 6
+  precompile_iterated_structure!(op)
+
+  # secant equation fails if uncommented
+  # op.α = dot(y,s)/dot(s,s)
+  return op
+end
+
+# Algorithm 3.2 (p15)
+# Theorem 2.3 (p6)
+function Base.Matrix(op::CompressedLBFGSOperator{T,M,V}) where {T,M,V}
+  B₀ = M(zeros(T, op.n, op.n))
+  map(i -> B₀[i, i] = op.α, 1:op.n)
+
+  BSY = M(undef, op.n, 2*op.k)
+  (op.k > 0) && (BSY[:, 1:op.k] = B₀ * op.Sₖ[:, 1:op.k])
+  (op.k > 0) && (BSY[:, op.k+1:2*op.k] = op.Yₖ[:, 1:op.k])
+  _C = M(undef, 2*op.k, 2*op.k)
+  (op.k > 0) && (_C[1:op.k, 1:op.k] .= transpose(op.Sₖ[:, 1:op.k]) * op.Sₖ[:, 1:op.k])
+  (op.k > 0) && (_C[1:op.k, op.k+1:2*op.k] .= op.Lₖ[1:op.k, 1:op.k])
+  (op.k > 0) && (_C[op.k+1:2*op.k, 1:op.k] .= transpose(op.Lₖ[1:op.k, 1:op.k]))
+  (op.k > 0) && (_C[op.k+1:2*op.k, op.k+1:2*op.k] .= .- op.Dₖ[1:op.k, 1:op.k])
+  C = inv(_C)
+
+  Bₖ = B₀ .- BSY * C * transpose(BSY)
+  return Bₖ
+end
+
+# Algorithm 3.2 (p15)
+# step 4, Jₖ is computed only if needed
+function inverse_cholesky(op::CompressedLBFGSOperator{T,M,V}) where {T,M,V}
+  view(op.intermediate_diagonal.diag, 1:op.k) .= inv.(view(op.Dₖ.diag, 1:op.k))
+
+  mul!(view(op.inverse_intermediate_1, 1:op.k, 1:op.k), view(op.intermediate_diagonal, 1:op.k, 1:op.k), transpose(view(op.Lₖ, 1:op.k, 1:op.k)))
+  mul!(view(op.chol_matrix, 1:op.k, 1:op.k), view(op.Lₖ, 1:op.k, 1:op.k), view(op.inverse_intermediate_1, 1:op.k, 1:op.k))
+
+  mul!(view(op.chol_matrix, 1:op.k, 1:op.k), transpose(view(op.Sₖ, :, 1:op.k)), view(op.Sₖ, :, 1:op.k), op.α, (T)(1))
+
+  cholesky!(Symmetric(view(op.chol_matrix, 1:op.k, 1:op.k)))
+  Jₖ = transpose(UpperTriangular(view(op.chol_matrix, 1:op.k, 1:op.k)))
+  return Jₖ
+end
+
+# step 6, must be improve
+function precompile_iterated_structure!(op::CompressedLBFGSOperator)
+  Jₖ = inverse_cholesky(op)
+
+  # constant update
+  view(op.intermediate_1, op.k+1:2*op.k, 1:op.k) .= 0
+  view(op.intermediate_2, 1:op.k, op.k+1:2*op.k) .= 0
+  view(op.intermediate_1, op.k+1:2*op.k, op.k+1:2*op.k) .= transpose(Jₖ)
+  view(op.intermediate_2, op.k+1:2*op.k, op.k+1:2*op.k) .= Jₖ
+
+  # updates related to D^(1/2)
+  view(op.intermediate_diagonal.diag, 1:op.k) .= sqrt.(view(op.Dₖ.diag, 1:op.k))
+  view(op.intermediate_1, 1:op.k,1:op.k) .= .- view(op.intermediate_diagonal, 1:op.k, 1:op.k)
+  view(op.intermediate_2, 1:op.k, 1:op.k) .= view(op.intermediate_diagonal, 1:op.k, 1:op.k)
+
+  # updates related to D^(-1/2)
+  view(op.intermediate_diagonal.diag, 1:op.k) .= (x -> 1/sqrt(x)).(view(op.Dₖ.diag, 1:op.k))
+  mul!(view(op.intermediate_1, 1:op.k,op.k+1:2*op.k), view(op.intermediate_diagonal, 1:op.k, 1:op.k), transpose(view(op.Lₖ, 1:op.k, 1:op.k)))
+  mul!(view(op.intermediate_2, op.k+1:2*op.k, 1:op.k), view(op.Lₖ, 1:op.k, 1:op.k), view(op.intermediate_diagonal, 1:op.k, 1:op.k))
+  view(op.intermediate_2, op.k+1:2*op.k, 1:op.k) .= view(op.intermediate_2, op.k+1:2*op.k, 1:op.k) .* -1
+
+  view(op.inverse_intermediate_1, 1:2*op.k, 1:2*op.k) .= inv(op.intermediate_1[1:2*op.k, 1:2*op.k])
+  view(op.inverse_intermediate_2, 1:2*op.k, 1:2*op.k) .= inv(op.intermediate_2[1:2*op.k, 1:2*op.k])
+end
+
+# Algorithm 3.2 (p15)
+function LinearAlgebra.mul!(Bv::V, op::CompressedLBFGSOperator{T,M,V}, v::V) where {T,M,V<:AbstractVector{T}}
+  # step 1-4 and 6 mainly done by Base.push!
+  # step 5
+  mul!(view(op.sol, 1:op.k), transpose(view(op.Yₖ, :, 1:op.k)), v)
+  mul!(view(op.sol, op.k+1:2*op.k), transpose(view(op.Sₖ, :, 1:op.k)), v)
+  # scal!(op.α, view(op.sol, op.k+1:2*op.k)) # more allocation, slower
+  view(op.sol, op.k+1:2*op.k) .*= op.α
+
+  mul!(view(op.intermediary_vector, 1:2*op.k), view(op.inverse_intermediate_2, 1:2*op.k, 1:2*op.k), view(op.sol, 1:2*op.k))
+  mul!(view(op.sol, 1:2*op.k), view(op.inverse_intermediate_1, 1:2*op.k, 1:2*op.k), view(op.intermediary_vector, 1:2*op.k))
+
+  # step 7 
+  mul!(Bv, view(op.Yₖ, :, 1:op.k),  view(op.sol, 1:op.k))
+  mul!(Bv, view(op.Sₖ, :, 1:op.k), view(op.sol, op.k+1:2*op.k), - op.α, (T)(-1))
+  Bv .+= op.α .* v 
+  return Bv
+end
+
+end
+
+using ..ModCompressedLBFGSOperator
+export CompressedLBFGSOperator
+export default_matrix_type, default_vector_type

src/qn.jl

@@ -5,3 +5,5 @@ import LinearAlgebra.diag
 
 include("lbfgs.jl")
 include("lsr1.jl")
+
+include("compressed_lbfgs.jl")

test/gpu/nvidia.jl

@@ -14,3 +14,5 @@ using LinearOperators, CUDA, CUDA.CUSPARSE, CUDA.CUSOLVER
   y = M * v
   @test y isa CuVector{Float32}
 end
+
+include("../test_clbfgs.jl")

test/runtests.jl

@@ -7,6 +7,7 @@ include("test_linop_allocs.jl")
 include("test_adjtrans.jl")
 include("test_cat.jl")
 include("test_lbfgs.jl")
+include("test_clbfgs.jl")
 include("test_lsr1.jl")
 include("test_kron.jl")
 include("test_callable.jl")

test/test_clbfgs.jl

@@ -0,0 +1,21 @@
+@testset "CompressedLBFGSOperator operator" begin
+  iter=50
+  n=100
+  n=5
+  types = [Float32, Float64]
+  for T in types  
+    lbfgs = CompressedLBFGSOperator(n; T) # mem=5
+    V = LinearOperators.default_vector_type(;T)
+    Bv = V(rand(T, n))
+    s = V(rand(T, n))
+    mul!(Bv, lbfgs, s) # warm-up
+    for i in 1:iter
+      s = V(rand(T, n))
+      y = V(rand(T, n))
+      push!(lbfgs, s, y)
+      allocs = @allocated mul!(Bv, lbfgs, s)
+      @test allocs == 0
+      @test Bv ≈ y
+    end  
+  end
+end