Skip to content

Commit

Permalink
evaluator: add BridgeDeviceEvaluator
Browse files Browse the repository at this point in the history
* Allow to deport the computation on the GPU, while using
  a CPU compatible solver.
* add CUSOLVERRF in the dependencies
* fix transpose in kernels
  • Loading branch information
frapac committed Jul 6, 2021
1 parent 0193b15 commit db1ca61
Show file tree
Hide file tree
Showing 20 changed files with 279 additions and 106 deletions.
2 changes: 2 additions & 0 deletions .github/workflows/action.yml
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ jobs:
- uses: julia-actions/setup-julia@latest
with:
version: ${{ matrix.julia-version }}
- run: julia --project deps/deps.jl
- uses: julia-actions/julia-buildpkg@latest
- uses: julia-actions/julia-runtest@latest

Expand All @@ -40,5 +41,6 @@ jobs:
- uses: julia-actions/setup-julia@latest
with:
version: ${{ matrix.julia-version }}
- run: julia --project deps/deps.jl
- uses: julia-actions/julia-buildpkg@latest
- uses: julia-actions/julia-runtest@latest
3 changes: 2 additions & 1 deletion benchmark/benchmarks.jl
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,8 @@ function run_benchmark(datafile, device, linsolver)
algo = linsolver(precond)
powerflow_solver = NewtonRaphson(tol=ntol)
nlp = ExaPF.ReducedSpaceEvaluator(polar;
linear_solver=algo, powerflow_solver=powerflow_solver)
powerflow_solver=powerflow_solver)
nlp.linear_solver = algo
convergence = ExaPF.update!(nlp, u0)
ExaPF.reset!(nlp)
convergence = ExaPF.update!(nlp, u0)
Expand Down
1 change: 1 addition & 0 deletions src/Evaluators/Evaluators.jl
Original file line number Diff line number Diff line change
Expand Up @@ -237,6 +237,7 @@ Reset evaluator `nlp` to default configuration.
function reset! end

include("common.jl")
include("bridge_evaluator.jl")

# Basic evaluators
include("reduced_evaluator.jl")
Expand Down
5 changes: 4 additions & 1 deletion src/Evaluators/MOI_wrapper.jl
Original file line number Diff line number Diff line change
Expand Up @@ -67,9 +67,12 @@ function MOI.eval_constraint(ev::MOIEvaluator, cons, x)
end

function MOI.eval_constraint_jacobian(ev::MOIEvaluator, jac, x)
n = length(x)
m = n_constraints(ev.nlp)
_update!(ev, x)
fill!(jac, 0)
jacobian!(ev.nlp, jac, x)
J = reshape(jac, m, n)
jacobian!(ev.nlp, J, x)
end

function MOI.eval_hessian_lagrangian(ev::MOIEvaluator, hess, x, σ, μ)
Expand Down
6 changes: 3 additions & 3 deletions src/Evaluators/auglag.jl
Original file line number Diff line number Diff line change
Expand Up @@ -88,11 +88,11 @@ function AugLagEvaluator(
return AugLagEvaluator(nlp, cons_type, cx, c₀, λ, λc, scaler, NLPCounter())
end
function AugLagEvaluator(
datafile::String; device=CPU(), options...
datafile::String; device=CPU(), scale=false, c₀=0.1, options...
)
nlp = ReducedSpaceEvaluator(datafile; device=device)
nlp = ReducedSpaceEvaluator(datafile; device=device, options...)
u0 = initial(nlp)
return AugLagEvaluator(nlp, u0; options...)
return AugLagEvaluator(nlp, u0; scale=scale, c₀=c₀)
end

has_hessian(nlp::AugLagEvaluator) = has_hessian(nlp.inner)
Expand Down
138 changes: 138 additions & 0 deletions src/Evaluators/bridge_evaluator.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,138 @@

#=
BridgeEvaluator
=#
struct BridgeDeviceArrays{VT, MT}
u::VT
g::VT
cons::VT
v::VT
y::VT
w::VT
jv::VT
J::MT
H::MT
end

function BridgeDeviceArrays(n::Int, m::Int, VT, MT)
BridgeDeviceArrays{VT, MT}(
VT(undef, n),
VT(undef, n),
VT(undef, m),
VT(undef, m),
VT(undef, m),
VT(undef, m),
VT(undef, n),
MT(undef, m, n),
MT(undef, n, n),
)
end

struct BridgeDeviceEvaluator{Evaluator, VT, MT, DVT, DMT} <: AbstractNLPEvaluator
inner::Evaluator
bridge::BridgeDeviceArrays{DVT, DMT}
end
function BridgeDeviceEvaluator(nlp::AbstractNLPEvaluator, device)
n, m = n_variables(nlp), n_constraints(nlp)
# Deporting device
VT = Array{Float64, 1}
MT = Array{Float64, 2}
if isa(nlp.model.device, CPU)
VTD = Array{Float64, 1}
MTD = Array{Float64, 2}
else
VTD = CUDA.CuArray{Float64, 1}
MTD = CUDA.CuArray{Float64, 2}
end
bridge = BridgeDeviceArrays(n, m, VTD, MTD)
return BridgeDeviceEvaluator{typeof(nlp), VT, MT, VTD, MTD}(nlp, bridge)
end
function BridgeDeviceEvaluator(case::String; device=KA.CPU())
nlp = ReducedSpaceEvaluator(case)
return BridgeDeviceEvaluator(nlp, device)
end

n_variables(nlp::BridgeDeviceEvaluator) = n_variables(nlp.inner)
n_constraints(nlp::BridgeDeviceEvaluator) = n_constraints(nlp.inner)
constraints_type(nlp::BridgeDeviceEvaluator) = constraints_type(nlp.inner)
has_hessian(nlp::BridgeDeviceEvaluator) = has_hessian(nlp.inner)
reset!(nlp::BridgeDeviceEvaluator) = reset!(nlp.inner)

# Getters
get(nlp::BridgeDeviceEvaluator, attr::AbstractNLPAttribute) = get(nlp.inner, attr)
get(nlp::BridgeDeviceEvaluator, attr::AbstractVariable) = get(nlp.inner, attr)
get(nlp::BridgeDeviceEvaluator, attr::PS.AbstractNetworkAttribute) = get(nlp.inner, attr)

function setvalues!(nlp::BridgeDeviceEvaluator, attr::PS.AbstractNetworkValues, values)
setvalues!(nlp.inner, attr, values)
end

function bounds(nlp::BridgeDeviceEvaluator{Ev, VT, MT, DVT, DMT}, attr::AbstractNLPAttribute) where {Ev, VT, MT, DVT, DMT}
b♭, b♯ = bounds(nlp.inner, attr)
return b♭ |> VT, b♯ |> VT
end

function initial(nlp::BridgeDeviceEvaluator{Ev, VT, MT, DVT, DMT}) where {Ev, VT, MT, DVT, DMT}
return initial(nlp.inner) |> VT
end

function update!(nlp::BridgeDeviceEvaluator, u)
copyto!(nlp.bridge.u, u)
return update!(nlp.inner, nlp.bridge.u)
end

objective(nlp::BridgeDeviceEvaluator, u) = objective(nlp.inner, nlp.bridge.u)

function constraint!(nlp::BridgeDeviceEvaluator, cons, u)
constraint!(nlp.inner, nlp.bridge.cons, nlp.bridge.u)
copyto!(cons, nlp.bridge.cons)
return
end

function gradient!(nlp::BridgeDeviceEvaluator, grad, u)
gradient!(nlp.inner, nlp.bridge.g, nlp.bridge.u)
copyto!(grad, nlp.bridge.g)
return
end

function jtprod!(nlp::BridgeDeviceEvaluator, jv, u, v)
copyto!(nlp.bridge.v, v)
jtprod!(nlp.inner, nlp.bridge.jv, nlp.bridge.u, nlp.bridge.v)
copyto!(jv, nlp.bridge.jv)
return
end

function jacobian!(nlp::BridgeDeviceEvaluator, jac, w)
jacobian!(nlp.inner, nlp.bridge.J, nlp.bridge.u)
copyto!(jac, nlp.bridge.J)
return
end

function ojtprod!(nlp::BridgeDeviceEvaluator, jv, u, σ, v)
copyto!(nlp.bridge.v, v)
ojtprod!(nlp.inner, nlp.bridge.jv, nlp.bridge.u, σ, nlp.bridge.v)
copyto!(jv, nlp.bridge.jv)
return
end

function hessprod!(nlp::BridgeDeviceEvaluator, hv, u, v)
copyto!(nlp.bridge.v, v)
hessprod!(nlp.inner, nlp.bridge.jv, nlp.bridge.u, nlp.bridge.v)
copyto!(hv, nlp.bridge.jv)
return
end

function hessian!(nlp::BridgeDeviceEvaluator, H, u)
hessian!(nlp.inner, nlp.bridge.H, nlp.bridge.u)
copyto!(H, nlp.bridge.H)
return
end

function hessian_lagrangian_penalty!(nlp::BridgeDeviceEvaluator, H, u, y, σ, w)
copyto!(nlp.bridge.w, w)
copyto!(nlp.bridge.y, y)
hessian_lagrangian_penalty!(nlp.inner, nlp.bridge.H, nlp.bridge.u, nlp.bridge.y, σ, nlp.bridge.w)
copyto!(H, nlp.bridge.H)
return
end

30 changes: 3 additions & 27 deletions src/Evaluators/common.jl
Original file line number Diff line number Diff line change
Expand Up @@ -111,7 +111,9 @@ function _inf_pr(nlp::AbstractNLPEvaluator, cons)
return max(err_inf, err_sup)
end

# Scaler utils
#=
SCALER
=#
abstract type AbstractScaler end

scale_factor(h, tol, η) = max(tol, η / max(1.0, h))
Expand Down Expand Up @@ -159,29 +161,3 @@ function MaxScaler(nlp::AbstractNLPEvaluator, u0::AbstractVector;
return MaxScaler{typeof(s_obj), typeof(s_cons)}(s_obj, s_cons, s_cons .* g♭, s_cons .* g♯)
end

struct BridgeDevice{VT, MT}
u::VT
g::VT
cons::VT
v::VT
y::VT
w::VT
jv::VT
J::MT
H::MT
end

function BridgeDevice(n::Int, m::Int, VT, MT)
BridgeDevice{VT, MT}(
VT(undef, n),
VT(undef, n),
VT(undef, m),
VT(undef, m),
VT(undef, m),
VT(undef, m),
VT(undef, n),
MT(undef, m, n),
MT(undef, n, n),
)
end

29 changes: 22 additions & 7 deletions src/Evaluators/reduced_evaluator.jl
Original file line number Diff line number Diff line change
Expand Up @@ -92,6 +92,7 @@ end
function ReducedSpaceEvaluator(
model::PolarForm{T, VI, VT, MT};
constraints=Function[voltage_magnitude_constraints, active_power_constraints, reactive_power_constraints],
linear_solver=nothing,
powerflow_solver=NewtonRaphson(tol=1e-12),
want_jacobian=true,
nbatch_hessian=1,
Expand All @@ -118,9 +119,10 @@ function ReducedSpaceEvaluator(
shift += m
end

SpMT = isa(model.device, CPU) ? SparseMatrixCSC : CUSPARSE.CuSparseMatrixCSR
# Build Linear Algebra
J = powerflow_jacobian(model)
linear_solver = DirectSolver(J)
J = powerflow_jacobian(model) |> SpMT
_linear_solver = isnothing(linear_solver) ? DirectSolver(J) : linear_solver

obj_ad = pullback_objective(model)
state_ad = FullSpaceJacobian(model, power_balance)
Expand Down Expand Up @@ -151,7 +153,7 @@ function ReducedSpaceEvaluator(
constraints, g_min, g_max,
buffer,
state_ad, obj_ad, cons_ad, cons_jac, hess_ad,
linear_solver, powerflow_solver, want_jacobian, false, want_hessian,
_linear_solver, powerflow_solver, want_jacobian, false, want_hessian,
)
end
function ReducedSpaceEvaluator(datafile::String; device=KA.CPU(), options...)
Expand Down Expand Up @@ -395,7 +397,6 @@ function full_jtprod!(nlp::ReducedSpaceEvaluator, jvx, jvu, u, v)
end

function jtprod!(nlp::ReducedSpaceEvaluator, jv, u, v)
@assert !isnothing(nlp.hesslag)
∂obj = nlp.obj_stack
μ = nlp.buffer.balance
jvx = ∂obj.stack.jvₓ ; fill!(jvx, 0)
Expand Down Expand Up @@ -569,14 +570,16 @@ end

# Batch Hessian
macro define_batch_hessian(function_name, target_function, args...)
fname_dispatch = Symbol("_" * String(function_name))
fname = Symbol(function_name)
argstup = Tuple(args)
quote
function $(esc(fname))(nlp::ReducedSpaceEvaluator, dest, $(map(esc, argstup)...))
function $(esc(fname_dispatch))(nlp::ReducedSpaceEvaluator, hesslag::BatchHessianLagrangian, dest, $(map(esc, argstup)...))
@assert has_hessian(nlp)
@assert n_batches(hesslag) > 1
n = ExaPF.n_variables(nlp)
∇²f = nlp.hesslag.hess
nbatch = size(nlp.hesslag.tmp_hv, 2)
∇²f = hesslag.hess
nbatch = size(hesslag.tmp_hv, 2)

# Allocate memory
v_cpu = zeros(n, nbatch)
Expand Down Expand Up @@ -609,6 +612,18 @@ macro define_batch_hessian(function_name, target_function, args...)
$target_function(nlp, hm, $(map(esc, argstup)...), v)
end
end
function $(esc(fname_dispatch))(nlp::ReducedSpaceEvaluator, hesslag::HessianLagrangian, dest, $(map(esc, argstup)...))
@assert has_hessian(nlp)
n = n_variables(nlp)
v = similar(x)
@inbounds for i in 1:n
hv = @view dest[:, i]
fill!(v, 0)
v[i] = 1.0
$target_function(nlp, hv, $(map(esc, argstup)...), v)
end
end
$(esc(fname))(nlp::ReducedSpaceEvaluator, dest, $(map(esc, argstup)...)) = $(esc(fname_dispatch))(nlp, nlp.hesslag, dest, $(map(esc, argstup)...))
end
end

Expand Down
Loading

0 comments on commit db1ca61

Please sign in to comment.