Refactoring of ExaPF

Project.toml

@@ -5,6 +5,7 @@ version = "0.6.0"
 
 [deps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+CUDAKernels = "72cfdca4-0801-4ab0-bf6a-d52aa10adc57"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"

benchmark/benchmarks.jl

@@ -40,61 +40,51 @@ function run_benchmark(datafile, device, linsolver)
     ntol = 1e-6
     pf = PowerSystem.PowerNetwork(datafile)
     polar = PolarForm(pf, device)
-    cache = ExaPF.get(polar, ExaPF.PhysicalState())
-    jx = AutoDiff.Jacobian(polar, ExaPF.power_balance, State())
+    mapx = ExaPF.my_map(polar, State())
+    nx = length(mapx)
+    stack = ExaPF.NetworkStack(polar)
+
+    basis = ExaPF.PolarBasis(polar)
+    pflow = ExaPF.PowerFlowBalance(polar)
+    jx = ExaPF.MyJacobian(polar, pflow ∘ basis, mapx)
     J = jx.J
     npartitions = ceil(Int64,(size(jx.J,1)/64))
     if npartitions < 2
         npartitions = 2
     end
     precond = ExaPF.LinearSolvers.BlockJacobiPreconditioner(J, npartitions, device)
-    # Retrieve initial state of network
-    u0 = ExaPF.initial(polar, Control())
 
     algo = linsolver(J; P=precond)
     powerflow_solver = NewtonRaphson(tol=ntol)
+    VT = typeof(stack.input)
+    pf_buffer = ExaPF.NLBuffer{VT}(nx)
 
-    # Init variables
-    buffer = get(polar, ExaPF.PhysicalState())
-    jx = AutoDiff.Jacobian(polar, ExaPF.power_balance, State())
-
-    # Warmstart
-    ExaPF.init_buffer!(polar, buffer)
-    ExaPF.powerflow(polar, jx, buffer, powerflow_solver; linear_solver=algo)
+    # Warm-up
+    ExaPF.nlsolve!(
+        powerflow_solver, jx, stack; linear_solver=algo, nl_buffer=pf_buffer,
+    )
 
-    TimerOutputs.reset_timer!(ExaPF.TIMER)
-    ExaPF.init_buffer!(polar, buffer)
-    convergence = ExaPF.powerflow(polar, jx, buffer, powerflow_solver; linear_solver=algo)
+    ExaPF.init!(polar, stack)
+    res = @timed ExaPF.nlsolve!(
+        powerflow_solver, jx, stack; linear_solver=algo, nl_buffer=pf_buffer,
+    )
+    convergence = res.value
 
-    # Make sure we are converged
-    @assert(convergence.has_converged)
-
-    # Output
-    prettytime = TimerOutputs.prettytime
-    timers = ExaPF.TIMER.inner_timers
-    inner_timer = timers["Newton"]
-    return convergence.has_converged, timers, inner_timer
+    return convergence.has_converged, res.time
 end
 
 function main()
     linsolver = eval(Meta.parse("LinearSolvers.$(ARGS[1])"))
     device = eval(Meta.parse("$(ARGS[2])()"))
     datafile = joinpath(dirname(@__FILE__), ARGS[3])
 
-    has_converged, timers, inner_timer = run_benchmark(datafile, device, linsolver)
-
-    if ARGS[1] == "DirectSolver"
-        println("$(ARGS[1]), $(ARGS[2]), $(ARGS[3]),",
-                printtimer(timers, "Newton"),
-                ", $(has_converged)")
-    else
-        println("$(ARGS[1]), $(ARGS[2]), $(ARGS[3]),",
-                printtimer(timers, "Newton"),",",
-                printtimer(inner_timer, "Jacobian"),",",
-                printtimer(inner_timer, "Linear Solver"),
-                ", $(has_converged)")
-    end
+    has_converged, timer = run_benchmark(datafile, device, linsolver)
     @test has_converged
+
+    println("$(ARGS[1]), $(ARGS[2]), $(ARGS[3]),",
+            timer,
+            ", $(has_converged)")
+
 end
 
 main()

src/ExaPF.jl

@@ -6,24 +6,16 @@ using LinearAlgebra
 using SparseArrays
 
 import CUDA
-import CUDA.CUBLAS
 import CUDA.CUSPARSE
-import CUDA.CUSOLVER
 
 import ForwardDiff
 using KernelAbstractions
 const KA = KernelAbstractions
-using TimerOutputs: @timeit, TimerOutput
 
 import Base: show, get
 
-const VERBOSE_LEVEL_HIGH = 3
-const VERBOSE_LEVEL_MEDIUM = 2
-const VERBOSE_LEVEL_LOW = 1
-const VERBOSE_LEVEL_NONE = 0
-const TIMER = TimerOutput()
+export run_pf
 
-include("utils.jl")
 include("architectures.jl")
 
 # Templates
@@ -43,4 +35,9 @@ const LS = LinearSolvers
 # Polar formulation
 include("Polar/polar.jl")
 
+# CUDA extension
+if CUDA.has_cuda()
+    include("cuda_wrapper.jl")
+end
+
 end

src/LinearSolvers/LinearSolvers.jl

@@ -8,14 +8,15 @@ import Base: show
 
 using CUDA
 using KernelAbstractions
+using CUDAKernels
 import CUDA.CUBLAS
 import CUDA.CUSOLVER
 import CUDA.CUSPARSE
 import Krylov
 import LightGraphs
 import Metis
 
-import ..ExaPF: xnorm, csclsvqr!
+import ..ExaPF: xnorm
 
 const KA = KernelAbstractions
 
@@ -102,9 +103,12 @@ exa_factorize(J::Adjoint{T, SparseMatrixCSC{T, Int}}) where T = lu(J.parent)'
 DirectSolver(J; options...) = DirectSolver(exa_factorize(J))
 DirectSolver() = DirectSolver(nothing)
 
+function update!(s::DirectSolver, J::AbstractMatrix)
+    lu!(s.factorization, J) # Update factorization inplace
+end
+
 # Reuse factorization in update
 function ldiv!(s::DirectSolver{<:LinearAlgebra.Factorization}, y::AbstractVector, J::AbstractMatrix, x::AbstractVector)
-    lu!(s.factorization, J) # Update factorization inplace
     LinearAlgebra.ldiv!(y, s.factorization, x) # Forward-backward solve
     return 0
 end
@@ -122,10 +126,6 @@ function rdiv!(s::DirectSolver{<:LinearAlgebra.Factorization}, y::AbstractArray,
     LinearAlgebra.ldiv!(y, s.factorization', x) # Forward-backward solve
     return 0
 end
-function rdiv!(s::DirectSolver{<:LinearAlgebra.Factorization}, y::Array, J::SparseMatrixCSC, x::Array)
-    LinearAlgebra.ldiv!(y, s.factorization', x) # Forward-backward solve
-    return 0
-end
 
 function ldiv!(::DirectSolver{Nothing}, y::Vector, J::AbstractMatrix, x::Vector)
     F = lu(J)
@@ -143,18 +143,6 @@ function batch_ldiv!(s::DirectSolver{<:LinearAlgebra.Factorization}, Y, Js::Vect
     end
 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{Nothing},
-    y::CUDA.CuVector, J::CUSPARSE.CuSparseMatrixCSC, x::CUDA.CuVector,
-)
-    csclsvqr!(J, x, y, 1e-8, one(Cint), 'O')
-    return 0
-end
 get_transpose(::DirectSolver, M::CUSPARSE.CuSparseMatrixCSR) = CUSPARSE.CuSparseMatrixCSC(M)
 
 function rdiv!(s::DirectSolver{<:LinearAlgebra.Factorization}, y::CUDA.CuVector, J::CUSPARSE.CuSparseMatrixCSR, x::CUDA.CuVector)
@@ -163,8 +151,11 @@ function rdiv!(s::DirectSolver{<:LinearAlgebra.Factorization}, y::CUDA.CuVector,
     return 0
 end
 
-function update_preconditioner!(solver::AbstractIterativeLinearSolver, J, device)
-    update(solver.precond, J, device)
+function update!(solver::AbstractIterativeLinearSolver, J::SparseMatrixCSC)
+    update(solver.precond, J, CPU())
+end
+function update!(solver::AbstractIterativeLinearSolver, J::CUSPARSE.CuSparseMatrixCSR)
+    update(solver.precond, J, CUDADevice())
 end
 
 """