linear solver: store factorization directly inside DirectSolver

* remove factorization attribute in ReducedSpaceEvaluator
* add a new rdiv! function for right solve

src/Evaluators/reduced_evaluator.jl

@@ -86,7 +86,6 @@ mutable struct ReducedSpaceEvaluator{T, VI, VT, MT, Jacx, Jacu, JacCons, Hess} <
     # Options
     linear_solver::LinearSolvers.AbstractLinearSolver
     powerflow_solver::AbstractNonLinearSolver
-    factorization::Union{Nothing, Factorization}
     has_jacobian::Bool
     update_jacobian::Bool
     has_hessian::Bool
@@ -95,7 +94,7 @@ end
 function ReducedSpaceEvaluator(
     model, x, u;
     constraints=Function[voltage_magnitude_constraints, active_power_constraints, reactive_power_constraints],
-    linear_solver=DirectSolver(),
+    linear_solver=direct_linear_solver(model),
     powerflow_solver=NewtonRaphson(tol=1e-12),
     want_jacobian=true,
     want_hessian=true,
@@ -143,7 +142,7 @@ function ReducedSpaceEvaluator(
         constraints, g_min, g_max,
         buffer,
         state_ad, obj_ad, cons_ad, cons_jac, hess_ad,
-        linear_solver, powerflow_solver, nothing, want_jacobian, false, want_hessian,
+        linear_solver, powerflow_solver, want_jacobian, false, want_hessian,
     )
 end
 function ReducedSpaceEvaluator(
@@ -222,18 +221,6 @@ end
 bounds(nlp::ReducedSpaceEvaluator, ::Variables) = (nlp.u_min, nlp.u_max)
 bounds(nlp::ReducedSpaceEvaluator, ::Constraints) = (nlp.g_min, nlp.g_max)
 
-function _update_factorization!(nlp::ReducedSpaceEvaluator)
-    ∇gₓ = nlp.state_jacobian.x.J
-    if !isa(∇gₓ, SparseMatrixCSC) || isa(nlp.linear_solver, LinearSolvers.AbstractIterativeLinearSolver)
-        return
-    end
-    if !isnothing(nlp.factorization)
-        lu!(nlp.factorization, ∇gₓ)
-    else
-        nlp.factorization = lu(∇gₓ)
-    end
-end
-
 ## Callbacks
 function update!(nlp::ReducedSpaceEvaluator, u)
     jac_x = nlp.state_jacobian.x
@@ -255,8 +242,6 @@ function update!(nlp::ReducedSpaceEvaluator, u)
 
     # Refresh values of active and reactive powers at generators
     update!(nlp.model, PS.Generators(), PS.ActivePower(), nlp.buffer)
-    # Update factorization if direct solver
-    _update_factorization!(nlp)
     # Evaluate Jacobian of power flow equation on current u
     AutoDiff.jacobian!(nlp.model, nlp.state_jacobian.u, nlp.buffer)
     # Specify that constraint's Jacobian is not up to date
@@ -289,8 +274,7 @@ function _forward_solve!(nlp::ReducedSpaceEvaluator, y, x)
         ∇gₓ = nlp.state_jacobian.x.J
         LinearSolvers.ldiv!(nlp.linear_solver, y, ∇gₓ, x)
     else
-        ∇gₓ = nlp.factorization
-        ldiv!(y, ∇gₓ, x)
+        LinearSolvers.ldiv!(nlp.linear_solver, y, x)
     end
 end
 
@@ -307,8 +291,7 @@ function _backward_solve!(nlp::ReducedSpaceEvaluator, y::VT, x::VT) where {VT <:
         ∇gT = LinearSolvers.get_transpose(nlp.linear_solver, ∇gₓ)
         LinearSolvers.ldiv!(nlp.linear_solver, y, ∇gT, x)
     else
-        ∇gf = nlp.factorization
-        LinearSolvers.ldiv!(nlp.linear_solver, y, ∇gf', x)
+        LinearSolvers.rdiv!(nlp.linear_solver, y, x)
     end
 end
 
@@ -394,11 +377,9 @@ function jacobian!(nlp::ReducedSpaceEvaluator, jac, u)
     m, nᵤ = size(Ju)
     ∇gᵤ = nlp.state_jacobian.u.J
     ∇gₓ = nlp.state_jacobian.x.J
-    # Compute factorization with UMFPACK
-    ∇gfac = nlp.factorization
-    # Compute adjoints all in once, using the same factorization
+    # Compute state sensitivities all in once
     μ = zeros(nₓ, nᵤ)
-    ldiv!(μ, ∇gfac, ∇gᵤ)
+    LinearSolvers.ldiv!(nlp.linear_solver, μ, ∇gᵤ)
     # Compute reduced Jacobian
     copy!(jac, Ju)
     mul!(jac, Jx, μ, -1.0, 1.0)
@@ -447,10 +428,11 @@ end
 
 function jtprod!(nlp::ReducedSpaceEvaluator, jv, u, v)
     ∂obj = nlp.obj_stack
+    μ = nlp.buffer.balance
     jvx = ∂obj.stack.jvₓ ; fill!(jvx, 0)
     jvu = ∂obj.stack.jvᵤ ; fill!(jvu, 0)
     full_jtprod!(nlp, jvx, jvu, u, v)
-    reduced_gradient!(nlp, jv, jvx, jvu, jvx, u)
+    reduced_gradient!(nlp, jv, jvx, jvu, μ, u)
 end
 
 function ojtprod!(nlp::ReducedSpaceEvaluator, jv, u, σ, v)

src/LinearSolvers/LinearSolvers.jl

@@ -54,18 +54,36 @@ get_transpose(::AbstractLinearSolver, M::AbstractMatrix) = transpose(M)
 
 """
     ldiv!(solver, x, A, y)
+    ldiv!(solver, x, y)
 
 * `solver::AbstractLinearSolver`: linear solver to solve the system
 * `x::AbstractVector`: Solution
 * `A::AbstractMatrix`: Input matrix
 * `y::AbstractVector`: RHS
 
 Solve the linear system ``A x = y`` using the algorithm
-specified in `solver`.
+specified in `solver`. If `A` is not specified, the function
+will used directly the factorization stored inplace.
 
 """
 function ldiv! end
 
+"""
+    rdiv!(solver, x, A, y)
+    rdiv!(solver, x, y)
+
+* `solver::AbstractLinearSolver`: linear solver to solve the system
+* `x::AbstractVector`: Solution
+* `A::AbstractMatrix`: Input matrix
+* `y::AbstractVector`: RHS
+
+Solve the linear system ``A^⊤ x = y`` using the algorithm
+specified in `solver`. If `A` is not specified, the function
+will used directly the factorization stored inplace.
+
+"""
+function rdiv! end
+
 """
     DirectSolver <: AbstractLinearSolver
 
@@ -75,27 +93,43 @@ Solve linear system ``A x = y`` with direct linear algebra.
 * On CUDA GPU, `DirectSolver` redirects the resolution to the method `CUSOLVER.csrlsvqr`
 
 """
-struct DirectSolver <: AbstractLinearSolver end
-DirectSolver(precond) = DirectSolver()
-function ldiv!(::DirectSolver,
-    y::Vector, J::AbstractMatrix, x::Vector,
-)
-    y .= J \ x
+struct DirectSolver{Fac<:Union{Nothing, LinearAlgebra.Factorization}} <: AbstractLinearSolver
+    factorization::Fac
+end
+
+DirectSolver(J::SparseMatrixCSC) = DirectSolver(lu(J))
+DirectSolver() = DirectSolver(nothing)
+DirectSolver(precond::AbstractPreconditioner) = DirectSolver(nothing)
+
+# Reuse factorization in update
+function ldiv!(s::DirectSolver{<:LinearAlgebra.Factorization}, y::Vector, J::AbstractMatrix, x::Vector)
+    lu!(s.factorization, J) # Update factorization inplace
+    LinearAlgebra.ldiv!(y, s.factorization, x) # Forward-backward solve
     return 0
 end
-function ldiv!(::DirectSolver,
-    y::Vector, J::Factorization, x::Vector,
-)
-    LinearAlgebra.ldiv!(y, J, x)
+# Solve system Ax = y
+function ldiv!(s::DirectSolver{<:LinearAlgebra.Factorization}, y::AbstractArray, x::AbstractArray)
+    LinearAlgebra.ldiv!(y, s.factorization, x) # Forward-backward solve
+    return 0
+end
+# Solve system A'x = y
+function rdiv!(s::DirectSolver{<:LinearAlgebra.Factorization}, y::Array, x::Array)
+    LinearAlgebra.ldiv!(y, s.factorization', x) # Forward-backward solve
     return 0
 end
-function ldiv!(::DirectSolver,
+
+function ldiv!(::DirectSolver{Nothing}, y::Vector, J::AbstractMatrix, x::Vector)
+    y .= J \ x
+    return 0
+end
+
+function ldiv!(::DirectSolver{Nothing},
     y::CUDA.CuVector, J::CUSPARSE.CuSparseMatrixCSR, x::CUDA.CuVector,
 )
     CUSOLVER.csrlsvqr!(J, x, y, 1e-8, one(Cint), 'O')
     return 0
 end
-function ldiv!(::DirectSolver,
+function ldiv!(::DirectSolver{Nothing},
     y::CUDA.CuVector, J::CUSPARSE.CuSparseMatrixCSC, x::CUDA.CuVector,
 )
     csclsvqr!(J, x, y, 1e-8, one(Cint), 'O')

src/Polar/polar.jl

@@ -268,6 +268,28 @@ function init_buffer!(form::PolarForm{T, IT, VT, MT}, buffer::PolarNetworkState)
     copyto!(buffer.qinj, qbus)
 end
 
+function direct_linear_solver(polar::PolarForm)
+    is_cpu = isa(polar.device, KA.CPU)
+    if is_cpu
+        jac = jacobian_sparsity(polar, power_balance, State())
+        return LinearSolvers.DirectSolver(jac)
+    else
+        # Factorization is not yet supported on the GPU
+        return LinearSolvers.DirectSolver(nothing)
+    end
+end
+
+function build_preconditioner(polar::PolarForm; nblocks=-1)
+    jac = jacobian_sparsity(polar, power_balance, State())
+    n = size(jac, 1)
+    npartitions = if nblocks > 0
+        nblocks
+    else
+        div(n, 32)
+    end
+    return LinearSolvers.BlockJacobiPreconditioner(jac, npartitions, polar.device)
+end
+
 function Base.show(io::IO, polar::PolarForm)
     # Network characteristics
     nbus = PS.get(polar.network, PS.NumberOfBuses())

test/Evaluators/powerflow.jl

@@ -4,6 +4,8 @@ using KernelAbstractions
 using Test
 
 import ExaPF: PowerSystem
+import ExaPF: LinearSolvers
+const LS = LinearSolvers
 
 @testset "Powerflow solver" begin
     datafile = joinpath(INSTANCES_DIR, "case14.raw")
@@ -19,13 +21,13 @@ import ExaPF: PowerSystem
 
     @testset "Deport computation on device $device" for device in DEVICES
         polar = PolarForm(pf, device)
-        jac = ExaPF.jacobian_sparsity(polar, ExaPF.power_balance, State())
-        precond = ExaPF.LinearSolvers.BlockJacobiPreconditioner(jac, npartitions, device)
+        precond = ExaPF.build_preconditioner(polar; nblocks=npartitions)
         # Retrieve initial state of network
         x0 = ExaPF.initial(polar, State())
         uk = ExaPF.initial(polar, Control())
 
         @testset "Powerflow solver $(LinSolver)" for LinSolver in ExaPF.list_solvers(device)
+            (LinSolver == LS.DirectSolver) && continue
             algo = LinSolver(precond)
             xk = copy(x0)
             nlp = ExaPF.ReducedSpaceEvaluator(