Add a FFTOperator (#26)

Author: amontoison

.github/workflows/action.yml

@@ -34,7 +34,7 @@ jobs:
     runs-on: cuda
     strategy:
       matrix:
-        julia-version: ['lts', '1']
+        julia-version: ['1']
         julia-arch: [x64]
 
     steps:
@@ -54,7 +54,7 @@ jobs:
     runs-on: amdgpu
     strategy:
       matrix:
-        julia-version: ['lts', '1']
+        julia-version: ['1']
         julia-arch: [x64]
 
     steps:

src/CompressedSensingIPM.jl

@@ -7,7 +7,7 @@ using FFTW
 using Krylov
 using NLPModels
 
-export FFTNLPModel, FFTKKTSystem, FFTParameters
+export FFTNLPModel, FFTKKTSystem, FFTParameters, FFTOperator
 
 include("fft_utils.jl")
 include("fft_model.jl")

src/fft_model.jl

@@ -13,19 +13,13 @@ mutable struct FFTParameters{AT,N,IM}
     end
 end
 
-mutable struct FFTNLPModel{T,VT,FFT,R,C,P} <: AbstractNLPModel{T,VT}
+mutable struct FFTNLPModel{T,VT,FFT,P} <: AbstractNLPModel{T,VT}
     meta::NLPModelMeta{T,VT}
+    counters::Counters
     parameters::P
     nβ::Int
-    counters::Counters
-    op::FFT
-    buffer_real::R
-    buffer_complex1::C
-    buffer_complex2::C
+    op_fft::FFT
     M_perpt_z0::VT
-    rdft::Bool
-    fft_timer::Base.RefValue{Float64}
-    mapping_timer::Base.RefValue{Float64}
     krylov_solver::Symbol
     preconditioner::Bool
 end
@@ -37,6 +31,7 @@ function FFTNLPModel{VT}(parameters::FFTParameters;
     T = eltype(VT)
     DFTdim = parameters.DFTdim   # problem size (1, 2, 3)
     DFTsize = parameters.DFTsize  # problem dimension
+    index_missing = parameters.index_missing
     nβ = prod(DFTsize)
     nvar = 2 * nβ
     ncon = 2 * nβ
@@ -69,31 +64,10 @@ function FFTNLPModel{VT}(parameters::FFTParameters;
     )
 
     # FFT operator
-    A_vec = VT(undef, nβ)
-    A = reshape(A_vec, DFTsize)
-    buffer_real = A
-    if rdft == true
-        op = plan_rfft(A)
-        M1 = (DFTsize[1] ÷ 2)
-        if DFTdim == 1
-            buffer_complex1 = Complex{T}.(A[1:M1+1])
-        elseif DFTdim == 2
-            buffer_complex1 = Complex{T}.(A[1:M1+1,:])
-        else
-            buffer_complex1 = Complex{T}.(A[1:M1+1,:,:])
-        end
-        buffer_complex2 = buffer_complex1
-    else
-        op = plan_fft(A)
-        buffer_complex1 = Complex{T}.(A)
-        buffer_complex2 = copy(buffer_complex1)
-    end
-    fft_timer = Ref{Float64}(0.0)
-    mapping_timer = Ref{Float64}(0.0)
-    tmp = M_perpt_z(buffer_real, buffer_complex1, buffer_complex2, op, DFTdim, DFTsize, parameters.z0, fft_timer, mapping_timer; rdft=rdft)
+    op_fft = FFTOperator{VT}(nβ, DFTdim, DFTsize, index_missing, rdft)
+    tmp = M_perpt_z(op_fft, parameters.z0)
     M_perpt_z0 = copy(tmp)
-    return FFTNLPModel(meta, parameters, nβ, Counters(), op, buffer_real, buffer_complex1,
-                       buffer_complex2, M_perpt_z0, rdft, fft_timer, mapping_timer, krylov_solver, preconditioner)
+    return FFTNLPModel(meta, Counters(), parameters, nβ, op_fft, M_perpt_z0, krylov_solver, preconditioner)
 end
 
 function NLPModels.obj(nlp::FFTNLPModel, x::AbstractVector)
@@ -103,7 +77,7 @@ function NLPModels.obj(nlp::FFTNLPModel, x::AbstractVector)
     lambda = nlp.parameters.lambda
     index_missing = nlp.parameters.index_missing
 
-    fft_val = M_perp_beta(nlp.buffer_real, nlp.buffer_complex1, nlp.buffer_complex2, nlp.op, DFTdim, DFTsize, x, index_missing, nlp.fft_timer, nlp.mapping_timer; rdft=nlp.rdft)
+    fft_val = M_perp_beta(nlp.op_fft, x)
     nβ = nlp.nβ
     beta = view(x, 1:nβ)
     c = view(x, nβ+1:2*nβ)
@@ -122,7 +96,7 @@ function NLPModels.grad!(nlp::FFTNLPModel, x::AbstractVector, g::AbstractVector)
     g_b = view(g, 1:nβ)
     g_c = view(g, nβ+1:2*nβ)
     beta = view(x, 1:nβ)
-    res = M_perpt_M_perp_vec(nlp.buffer_real, nlp.buffer_complex1, nlp.buffer_complex2, nlp.op, DFTdim, DFTsize, beta, index_missing, nlp.fft_timer, nlp.mapping_timer; rdft=nlp.rdft)
+    res = M_perpt_M_perp_vec(nlp.op_fft, beta)
     g_b .= res .- nlp.M_perpt_z0
     fill!(g_c, lambda)
     return g

src/fft_utils.jl

@@ -1,5 +1,62 @@
-function M_perpt_z(buffer_real, buffer_complex1, buffer_complex2, op, dim, _size, z, fft_timer, mapping_timer; rdft::Bool=false)
-    N = prod(_size)
+struct FFTOperator{R,C,OP,N,IM}
+    buffer_real::R
+    buffer_complex1::C
+    buffer_complex2::C
+    op::OP
+    DFTdim::Int64
+    DFTsize::NTuple{N,Int64}
+    index_missing::IM
+    fft_timer::Base.RefValue{Float64}
+    mapping_timer::Base.RefValue{Float64}
+    rdft::Bool
+end
+
+function FFTOperator{VT}(nβ, DFTdim, DFTsize, index_missing, rdft) where VT
+    T = eltype(VT)
+    A_vec = VT(undef, nβ)
+    A = reshape(A_vec, DFTsize)
+    buffer_real = A
+    if rdft == true
+        op = plan_rfft(A)
+        M1 = (DFTsize[1] ÷ 2)
+        if DFTdim == 1
+            buffer_complex1 = Complex{T}.(A[1:M1+1])
+        elseif DFTdim == 2
+            buffer_complex1 = Complex{T}.(A[1:M1+1,:])
+        else
+            buffer_complex1 = Complex{T}.(A[1:M1+1,:,:])
+        end
+        buffer_complex2 = buffer_complex1
+    else
+        op = plan_fft(A)
+        buffer_complex1 = Complex{T}.(A)
+        buffer_complex2 = copy(buffer_complex1)
+    end
+    fft_timer = Ref{Float64}(0.0)
+    mapping_timer = Ref{Float64}(0.0)
+
+    return FFTOperator(buffer_real, buffer_complex1, buffer_complex2, op, DFTdim, DFTsize, index_missing, fft_timer, mapping_timer, rdft)
+end
+
+function M_perpt_z(op_fft::FFTOperator, z)
+    return M_perpt_z(op_fft.buffer_real, op_fft.buffer_complex1, op_fft.buffer_complex2, op_fft.op,
+                     op_fft.DFTdim, op_fft.DFTsize, z, op_fft.fft_timer, op_fft.mapping_timer; rdft=op_fft.rdft)
+end
+
+function M_perp_beta(op_fft::FFTOperator, beta)
+    return M_perp_beta(op_fft.buffer_real, op_fft.buffer_complex1, op_fft.buffer_complex2, op_fft.op,
+                       op_fft.DFTdim, op_fft.DFTsize, beta, op_fft.index_missing, op_fft.fft_timer,
+                       op_fft.mapping_timer; rdft=op_fft.rdft)
+end
+
+function M_perpt_M_perp_vec(op_fft::FFTOperator, vec)
+    tmp = M_perp_beta(op_fft, vec)
+    tmp = M_perpt_z(op_fft, tmp)
+    return tmp
+end
+
+function M_perpt_z(buffer_real, buffer_complex1, buffer_complex2, op, DFTdim, DFTsize, z, fft_timer, mapping_timer; rdft::Bool=false)
+    N = prod(DFTsize)
 
     t1 = time_ns()
     if rdft
@@ -14,18 +71,18 @@ function M_perpt_z(buffer_real, buffer_complex1, buffer_complex2, op, dim, _size
 
     t3 = time_ns()
     beta = vec(buffer_real)
-    DFT_to_beta!(beta, dim, _size, temp; rdft)
+    DFT_to_beta!(beta, DFTdim, DFTsize, temp; rdft)
     t4 = time_ns()
     mapping_timer[] = mapping_timer[] + (t4 - t3) / 1e9
     return beta
 end
 
-function M_perp_beta(buffer_real, buffer_complex1, buffer_complex2, op, dim, _size, beta, idx_missing, fft_timer, mapping_timer; rdft::Bool=false)
-    N = prod(_size)
+function M_perp_beta(buffer_real, buffer_complex1, buffer_complex2, op, DFTdim, DFTsize, beta, index_missing, fft_timer, mapping_timer; rdft::Bool=false)
+    N = prod(DFTsize)
 
     t3 = time_ns()
     v = buffer_complex2
-    beta_to_DFT!(v, dim, _size, beta; rdft)
+    beta_to_DFT!(v, DFTdim, DFTsize, beta; rdft)
     t4 = time_ns()
     mapping_timer[] = mapping_timer[] + (t4 - t3) / 1e9
 
@@ -40,111 +97,111 @@ function M_perp_beta(buffer_real, buffer_complex1, buffer_complex2, op, dim, _si
     t2 = time_ns()
     fft_timer[] = fft_timer[] + (t2 - t1) / 1e9
 
-    buffer_real[idx_missing] .= 0
+    buffer_real[index_missing] .= 0
     return buffer_real
 end
 
-function M_perpt_M_perp_vec(buffer_real, buffer_complex1, buffer_complex2, op, dim, _size, vec, idx_missing, fft_timer, mapping_timer; rdft::Bool=false)
-    temp = M_perp_beta(buffer_real, buffer_complex1, buffer_complex2, op, dim, _size, vec, idx_missing, fft_timer, mapping_timer; rdft)
-    temp = M_perpt_z(buffer_real, buffer_complex1, buffer_complex2, op, dim, _size, temp, fft_timer, mapping_timer; rdft)
-    return temp
+function M_perpt_M_perp_vec(buffer_real, buffer_complex1, buffer_complex2, op, DFTdim, DFTsize, vec, index_missing, fft_timer, mapping_timer; rdft::Bool=false)
+    tmp = M_perp_beta(buffer_real, buffer_complex1, buffer_complex2, op, DFTdim, DFTsize, vec, index_missing, fft_timer, mapping_timer; rdft)
+    tmp = M_perpt_z(buffer_real, buffer_complex1, buffer_complex2, op, DFTdim, DFTsize, tmp, fft_timer, mapping_timer; rdft)
+    return tmp
 end
 
 # mapping between DFT and real vector beta
 
 # mapping DFT to beta
-# @param dim The dimension of the problem (dim = 1, 2, 3)
-# @param size The size of each dimension of the problem
-#(we only consider the cases when the sizes are even for all the dimenstions)
+# @param DFTdim The DFTdimension of the problem (DFTdim = 1, 2, 3)
+# @param size The size of each DFTdimension of the problem
+#(we only consider the cases when the sizes are even for all the DFTdimenstions)
 #(size is a tuple, e.g. size = (10, 20, 30))
 # @param v DFT
 
 # @details This fucnction maps DFT to beta
 
-# @return A 1-dimensional real vector beta whose length is the product of size
+# @return A 1-DFTdimensional real vector beta whose length is the product of size
 # @example
-# >dim = 2;
+# >DFTdim = 2;
 # >size1 = (6, 8)
 # >x = randn(6, 8)
 # >v = fft(x)/sqrt(prod(size1))
-# >beta = DFT_to_beta(dim, size1, v)
+# >beta = DFT_to_beta(DFTdim, size1, v)
 
-function DFT_to_beta!(beta, dim::Int, size, v; rdft::Bool=false)
-    if (dim == 1)
+function DFT_to_beta!(beta, DFTdim::Int, size, v; rdft::Bool=false)
+    if (DFTdim == 1)
         DFT_to_beta_1d!(beta, v, size; rdft)
-    elseif (dim == 2)
+    elseif (DFTdim == 2)
         DFT_to_beta_2d!(beta, v, size; rdft)
     else
         DFT_to_beta_3d!(beta, v, size; rdft)
     end
     return beta
 end
 
-function DFT_to_beta(dim::Int, size, v::Array{ComplexF64}; rdft::Bool=false)
+function DFT_to_beta(DFTdim::Int, size, v::Array{ComplexF64}; rdft::Bool=false)
     N = prod(size)
     beta = Vector{Float64}(undef, N)
-    DFT_to_beta!(beta, dim, size, v; rdft)
+    DFT_to_beta!(beta, DFTdim, size, v; rdft)
     return beta
 end
 
-function DFT_to_beta(dim::Int, size, v::CuArray{ComplexF64}; rdft::Bool=false)
+function DFT_to_beta(DFTdim::Int, size, v::CuArray{ComplexF64}; rdft::Bool=false)
     N = prod(size)
     beta = CuVector{Float64}(undef, N)
-    DFT_to_beta!(beta, dim, size, v; rdft)
+    DFT_to_beta!(beta, DFTdim, size, v; rdft)
     return beta
 end
 
-function DFT_to_beta(dim::Int, size, v::ROCArray{ComplexF64}; rdft::Bool=false)
+function DFT_to_beta(DFTdim::Int, size, v::ROCArray{ComplexF64}; rdft::Bool=false)
     N = prod(size)
     beta = ROCVector{Float64}(undef, N)
-    DFT_to_beta!(beta, dim, size, v; rdft)
+    DFT_to_beta!(beta, DFTdim, size, v; rdft)
     return beta
 end
 
 # mapping beta to DFT
-# @param dim The dimension of the problem (dim = 1, 2, 3)
-# @param size The size of each dimension of the problem
-#(we only consider the cases when the sizes are even for all the dimenstions)
+# @param DFTdim The DFTdimension of the problem (DFTdim = 1, 2, 3)
+# @param size The size of each DFTdimension of the problem
+#(we only consider the cases when the sizes are even for all the DFTdimenstions)
 #(size is a tuple, e.g. size = (10, 20, 30))
-# @param beta A 1-dimensional real vector with length equal to the product of size
+# @param beta A 1-DFTdimensional real vector with length equal to the product of size
 
 # @details This fucnction maps beta to DFT
 
 # @return DFT DFT shares the same size as param sizes
 
 # @example
-# >dim = 2;
+# >DFTdim = 2;
 # >size1 = (6, 8)
 # >x = randn(6, 8)
 # >v = fft(x)/sqrt(prod(size1))
-# >beta = DFT_to_beta(dim, size1, v)
-# >w = beta_to_DFT(dim, size1, beta) (w should be equal to v)
+# >beta = DFT_to_beta(DFTdim, size1, v)
+# >w = beta_to_DFT(DFTdim, size1, beta) (w should be equal to v)
 
-function beta_to_DFT!(v, dim::Int, size, beta; rdft::Bool=false)
-    if (dim == 1)
+function beta_to_DFT!(v, DFTdim::Int, size, beta; rdft::Bool=false)
+    if (DFTdim == 1)
         v = beta_to_DFT_1d!(v, beta, size; rdft)
-    elseif (dim == 2)
+    elseif (DFTdim == 2)
         v = beta_to_DFT_2d!(v, beta, size; rdft)
     else
         v = beta_to_DFT_3d!(v, beta, size; rdft)
     end
     return v
 end
 
-function beta_to_DFT(dim::Int, size, beta::StridedVector{Float64}; rdft::Bool=false)
+function beta_to_DFT(DFTdim::Int, size, beta::StridedVector{Float64}; rdft::Bool=false)
     v = Array{ComplexF64}(undef, size)
-    beta_to_DFT!(v, dim, size, beta; rdft)
+    beta_to_DFT!(v, DFTdim, size, beta; rdft)
     return v
 end
 
-function beta_to_DFT(dim::Int, size, beta::StridedCuVector{Float64}; rdft::Bool=false)
+function beta_to_DFT(DFTdim::Int, size, beta::StridedCuVector{Float64}; rdft::Bool=false)
     v = CuArray{ComplexF64}(undef, size)
-    beta_to_DFT!(v, dim, size, beta; rdft)
+    beta_to_DFT!(v, DFTdim, size, beta; rdft)
     return v
 end
 
-function beta_to_DFT(dim::Int, size, beta::AMDGPU.StridedROCVector{Float64}; rdft::Bool=false)
+function beta_to_DFT(DFTdim::Int, size, beta::AMDGPU.StridedROCVector{Float64}; rdft::Bool=false)
     v = ROCArray{ComplexF64}(undef, size)
-    beta_to_DFT!(v, dim, size, beta; rdft)
+    beta_to_DFT!(v, DFTdim, size, beta; rdft)
     return v
 end

src/kkt.jl

@@ -43,7 +43,7 @@ function LinearAlgebra.mul!(y::AbstractVector, K::CondensedFFTKKTSystem, x::Abst
     xz = view(x, nβ+1:2*nβ)
 
     # Evaluate Mᵀ M xβ
-    Mβ .= M_perpt_M_perp_vec(nlp.buffer_real, nlp.buffer_complex1, nlp.buffer_complex2, nlp.op, DFTdim, DFTsize, xβ, index_missing, nlp.fft_timer, nlp.mapping_timer; K.nlp.rdft)
+    Mβ .= M_perpt_M_perp_vec(nlp.op_fft, xβ)
 
     yβ .= beta .* yβ .+ alpha .* (Mβ .+ K.Λ1 .* xβ .+ K.Λ2 .* xz)
     yz .= beta .* yz .+ alpha .* (K.Λ2 .* xβ .+ K.Λ1 .* xz)
@@ -252,7 +252,7 @@ function MadNLP.mul!(y::VT, kkt::FFTKKTSystem, x::VT, alpha::Number, beta::Numbe
     xy2 = view(_x, 5*nβ+1:6*nβ)
 
     # Evaluate (MᵀM) * xβ
-    Mβ .= M_perpt_M_perp_vec(kkt.nlp.buffer_real, kkt.nlp.buffer_complex1, kkt.nlp.buffer_complex2, kkt.nlp.op, DFTdim, DFTsize, xβ, index_missing, kkt.nlp.fft_timer, kkt.nlp.mapping_timer; kkt.nlp.rdft)
+    Mβ .= M_perpt_M_perp_vec(kkt.nlp.op_fft, xβ)
     yβ .= beta .* yβ .+ alpha .* (Mβ .- xy1 .+ xy2)
     yz .= beta .* yz .- alpha .* (xy1 .+ xy2)
     ys1 .= beta .* ys1 .- alpha .* xy1