Add multiple timers

Author: amontoison

fft_model.jl

@@ -24,6 +24,8 @@ mutable struct FFTNLPModel{T,VT,FFT,R,C} <: AbstractNLPModel{T,VT}
     buffer_complex1::C
     buffer_complex2::C
     rdft::Bool
+    fft_timer::Ref{Float64}
+    mapping_timer::Ref{Float64}
 end
 
 function FFTNLPModel{T,VT}(parameters::FFTParameters; rdft::Bool=false) where {T,VT}
@@ -80,7 +82,9 @@ function FFTNLPModel{T,VT}(parameters::FFTParameters; rdft::Bool=false) where {T
         buffer_complex1 = Complex{T}.(A)
         buffer_complex2 = copy(buffer_complex1)
     end
-    return FFTNLPModel(meta, parameters, N, Counters(), op, buffer_real, buffer_complex1, buffer_complex2, rdft)
+    fft_timer = Ref{Float64}(0.0)
+    mapping_timer = Ref{Float64}(0.0)
+    return FFTNLPModel(meta, parameters, N, Counters(), op, buffer_real, buffer_complex1, buffer_complex2, rdft, fft_timer, mapping_timer)
 end
 
 include("kkt.jl")
@@ -186,7 +190,7 @@ function NLPModels.obj(nlp::FFTNLPModel, x::AbstractVector)
     index_missing = nlp.parameters.paramf[5]
     # Mt = nlp.parameters.paramf[6]
 
-    fft_val = M_perp_beta(nlp.buffer_real, nlp.buffer_complex1, nlp.buffer_complex2, nlp.op, DFTdim, DFTsize, x, index_missing; nlp.rdft)
+    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; nlp.rdft)
     N = nlp.N
     beta = view(x, 1:N)
     c = view(x, N+1:2*N)
@@ -207,7 +211,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.rdft)
+    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; nlp.rdft)
     g_b .= res .- M_perptz
     fill!(g_c, lambda)
     return g
@@ -232,7 +236,7 @@ function NLPModels.hprod!(
     n = prod(DFTsize)
     hv_b = view(hv, 1:n)
     hv_c = view(hv, n+1:2*n)
-    hv_b .= M_perpt_M_perp_vec(nlp.buffer_real, nlp.buffer_complex1, nlp.buffer_complex2, nlp.op, DFTdim, DFTsize, v[1:n], index_missing; nlp.rdft)
+    hv_b .= M_perpt_M_perp_vec(nlp.buffer_real, nlp.buffer_complex1, nlp.buffer_complex2, nlp.op, DFTdim, DFTsize, v[1:n], index_missing, nlp.fft_timer, nlp.mapping_timer; nlp.rdft)
     fill!(hv_c, 0.0)
     return hv
 end
@@ -276,7 +280,7 @@ function NLPModels.hess_coord!(
         for i in 1:nβ
             fill!(v, 0.0)
             v[i] = 1.0
-            H[:, i] .= M_perpt_M_perp_vec(nlp.buffer_real, nlp.buffer_complex1, nlp.buffer_complex2, nlp.op, DFTdim, DFTsize, v, index_missing; nlp.rdft)
+            H[:, i] .= M_perpt_M_perp_vec(nlp.buffer_real, nlp.buffer_complex1, nlp.buffer_complex2, nlp.op, DFTdim, DFTsize, v, index_missing, nlp.fft_timer, nlp.mapping_timer; nlp.rdft)
         end
 
         cnt = 1

fft_utils.jl

@@ -25,39 +25,55 @@ include("mapping_gpu.jl")
 # >size1 = 4;
 # >M_perptz = M_perp_tz(z_zero, dim, size1)
 
-function M_perp_tz(buffer_real, buffer_complex1, buffer_complex2, op, dim, _size, z_zero; rdft::Bool=false)
+function M_perp_tz(buffer_real, buffer_complex1, buffer_complex2, op, dim, _size, z_zero, fft_timer, mapping_timer; rdft::Bool=false)
     N = prod(_size)
+
+    t1 = time_ns()
     if rdft
         temp = mul!(buffer_complex1, op, z_zero)  # op_rfft
     else
         buffer_complex2 .= z_zero  # z_zero should be store in a complex buffer for mul!
         temp = mul!(buffer_complex1, op, buffer_complex2)  # op_fft
     end
     temp ./= sqrt(N)
+    t2 = time_ns()
+    fft_timer[] = fft_timer[] + (t2 - t1) / 1e9
+
+    t3 = time_ns()
     beta = vec(buffer_real)
     DFT_to_beta!(beta, dim, _size, 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; rdft::Bool=false)
+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)
+
+    t3 = time_ns()
     v = buffer_complex2
     beta_to_DFT!(v, dim, _size, beta; rdft)
+    t4 = time_ns()
+    mapping_timer[] = mapping_timer[] + (t4 - t3) / 1e9
 
+    t1 = time_ns()
     if rdft
         ldiv!(buffer_real, op, v)  # op_rfft
         buffer_real .*= sqrt(N)
     else
         temp = ldiv!(buffer_complex1, op, v)  # op_fft
         buffer_real .= real.(temp) .* sqrt(N)
     end
+    t2 = time_ns()
+    fft_timer[] = fft_timer[] + (t2 - t1) / 1e9
+
     buffer_real[idx_missing] .= 0
     return buffer_real
 end
 
-function M_perpt_M_perp_vec(buffer_real, buffer_complex1, buffer_complex2, op, dim, _size, vec, idx_missing; rdft::Bool=false)
-    temp = M_perp_beta(buffer_real, buffer_complex1, buffer_complex2, op, dim, _size, vec, idx_missing; rdft)
-    temp = M_perp_tz(buffer_real, buffer_complex1, buffer_complex2, op, dim, _size, temp; rdft)
+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_perp_tz(buffer_real, buffer_complex1, buffer_complex2, op, dim, _size, temp, fft_timer, mapping_timer; rdft)
     return temp
 end
 

fft_vishwas.jl

@@ -129,6 +129,11 @@ N = length(results.solution) ÷ 2
 beta_MadNLP = results.solution[1:N]
 println("Timer: $(timer)")
 
+# solver.kkt.krylov_iterations
+# solver.kkt.krylov_timer
+# nlp.fft_timer[]
+# nlp.mapping_timer[]
+
 using DelimitedFiles
 open("sol_vishwas.txt", "w") do io
     writedlm(io, Vector(beta_MadNLP))

kkt.jl

@@ -24,14 +24,18 @@ struct CondensedFFTKKT{T, VT, FFT, R, C} <: AbstractMatrix{T}
     buffer_complex1::C  # Buffer for fft and ifft
     buffer_complex2::C  # Buffer for fft and ifft
     rdft::Bool
+    fft_timer::Ref{Float64}
+    mapping_timer::Ref{Float64}
 end
 
 function CondensedFFTKKT{T, VT}(nlp::FFTNLPModel{T, VT}) where {T, VT}
     nβ = nlp.N
     buf1 = VT(undef, nβ)
     Λ1 = VT(undef, nβ)
     Λ2 = VT(undef, nβ)
-    return CondensedFFTKKT{T, VT, typeof(nlp.op), typeof(nlp.buffer_real), typeof(nlp.buffer_complex1)}(nβ, nlp.parameters, buf1, Λ1, Λ2, nlp.op, nlp.buffer_real, nlp.buffer_complex1, nlp.buffer_complex2, nlp.rdft)
+    return CondensedFFTKKT{T, VT, typeof(nlp.op), typeof(nlp.buffer_real), typeof(nlp.buffer_complex1)}(
+                nβ, nlp.parameters, buf1, Λ1, Λ2, nlp.op, nlp.buffer_real,
+                nlp.buffer_complex1, nlp.buffer_complex2, nlp.rdft, nlp.fft_timer, nlp.mapping_timer)
 end
 
 Base.size(K::CondensedFFTKKT) = (2*K.nβ, 2*K.nβ)
@@ -54,7 +58,7 @@ function LinearAlgebra.mul!(y::AbstractVector, K::CondensedFFTKKT, x::AbstractVe
     xz  = view(x, nβ+1:2*nβ)
 
     # Evaluate Mᵀ M xβ
-    Mβ .= M_perpt_M_perp_vec(K.buffer_real, K.buffer_complex1, K.buffer_complex2, K.op, DFTdim, DFTsize, xβ, index_missing; K.rdft)
+    Mβ .= M_perpt_M_perp_vec(K.buffer_real, K.buffer_complex1, K.buffer_complex2, K.op, DFTdim, DFTsize, xβ, index_missing, K.fft_timer, K.mapping_timer; K.rdft)
 
     yβ .= beta .* yβ .+ alpha .* (Mβ .+ K.Λ1 .* xβ .+ K.Λ2 .* xz)
     yz .= beta .* yz .+ alpha .* (K.Λ2 .* xβ .+ K.Λ1 .* xz)
@@ -124,6 +128,8 @@ struct FFTKKTSystem{T, VI, VT, MT, LS} <: MadNLP.AbstractReducedKKTSystem{T, VT,
     z1::VT           # dimension nβ
     z2::VT           # dimension 2 * nβ
     linear_solver::LS
+    krylov_iterations::Vector{Int}
+    krylov_timer::Vector{Float64}
 end
 
 function MadNLP.create_kkt_system(
@@ -167,7 +173,7 @@ function MadNLP.create_kkt_system(
         reg, pr_diag, du_diag, l_diag, u_diag, l_lower, u_lower,
         ind_cons.ind_lb, ind_cons.ind_ub,
         z1, z2,
-        linear_solver,
+        linear_solver, Int[], Float64[],
     )
 end
 
@@ -259,7 +265,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.K.buffer_real, kkt.K.buffer_complex1, kkt.K.buffer_complex2, kkt.K.op, DFTdim, DFTsize, xβ, index_missing; kkt.K.rdft)
+    Mβ .= M_perpt_M_perp_vec(kkt.K.buffer_real, kkt.K.buffer_complex1, kkt.K.buffer_complex2, kkt.K.op, DFTdim, DFTsize, xβ, index_missing, kkt.K.fft_timer, kkt.K.mapping_timer; kkt.K.rdft)
     yβ .= beta .* yβ .+ alpha .* (Mβ .- xy1 .+ xy2)
     yz .= beta .* yz .- alpha .* (xy1 .+ xy2)
     ys1 .= beta .* ys1 .- alpha .* xy1
@@ -358,6 +364,8 @@ function MadNLP.solve!(kkt::FFTKKTSystem, w::MadNLP.AbstractKKTVector)
     # Solve with CG
     Krylov.solve!(kkt.linear_solver, kkt.K, b, M=kkt.P, atol=1e-12, rtol=0.0, verbose=0)
     x = Krylov.solution(kkt.linear_solver)
+    push!(kkt.krylov_iterations, kkt.linear_solver |> niterations)
+    push!(kkt.krylov_timer, kkt.linear_solver.stats.timer)
 
     # Unpack solution
     w1 .= x[1:nβ]                              # / x