Handwritten adjoints and removal of Zygote

.github/workflows/action.yml

@@ -20,6 +20,8 @@ jobs:
       - uses: julia-actions/julia-runtest@latest
 
   test-moonshot:
+    env:
+      CUDA_VISIBLE_DEVICES: 1
     runs-on: self-hosted
     strategy:
       matrix:

Project.toml

@@ -18,7 +18,6 @@ Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
-Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
 CUDA = "^2.0"
@@ -32,7 +31,6 @@ MathOptInterface = "0.9"
 Metis = "1"
 SparseDiffTools = "1"
 TimerOutputs = "0.5"
-Zygote = "~0.6"
 julia = "^1.5"
 
 [extras]

src/ExaPF.jl

@@ -24,7 +24,6 @@ const KA = KernelAbstractions
 import MathOptInterface
 const MOI = MathOptInterface
 using TimerOutputs: @timeit, TimerOutput
-import Zygote
 
 import Base: show, get
 

src/Polar/Constraints/constraints.jl

@@ -14,6 +14,28 @@ include("active_power.jl")
 include("reactive_power.jl")
 include("line_flow.jl")
 
+# Generic functions
+function adjoint!(
+    polar::PolarForm,
+    func::Function,
+    ∂cons, cons,
+    stack, buffer,
+)
+    @assert is_constraint(func)
+    ∂pinj = similar(buffer.vmag) ; fill(∂pinj, 0)
+    ∂qinj = nothing # TODO
+    fill!(stack.∂vm, 0)
+    fill!(stack.∂va, 0)
+    adjoint!(
+        polar, func,
+        ∂cons, cons,
+        buffer.vmag, stack.∂vm,
+        buffer.vang, stack.∂va,
+        buffer.pinj, ∂pinj,
+        buffer.qinj, ∂qinj,
+    )
+end
+
 function jacobian_sparsity(polar::PolarForm, func::Function, xx::AbstractVariable)
     nbus = get(polar, PS.NumberOfBuses())
     Vre = Float64[i for i in 1:nbus]

src/Polar/Constraints/line_flow.jl

@@ -24,46 +24,19 @@ function flow_constraints(polar::PolarForm, cons, vm, va, pbus, qbus)
 end
 
 function flow_constraints_grad!(polar::PolarForm, cons_grad, buffer, weights)
-    T = typeof(buffer.vmag)
-    if isa(buffer.vmag, Array)
-        kernel! = accumulate_view!(KA.CPU())
-    else
-        kernel! = accumulate_view!(KA.CUDADevice())
-    end
-
     nlines = PS.get(polar.network, PS.NumberOfLines())
     nbus = PS.get(polar.network, PS.NumberOfBuses())
-    f = polar.topology.f_buses
-    t = polar.topology.t_buses
-    # Statements references below (1)
-    fr_vmag = buffer.vmag[f]
-    to_vmag = buffer.vmag[t]
-    fr_vang = buffer.vang[f]
-    to_vang = buffer.vang[t]
-    function lumping(fr_vmag, to_vmag, fr_vang, to_vang)
-        Δθ = fr_vang .- to_vang
-        cosθ = cos.(Δθ)
-        sinθ = sin.(Δθ)
-        cons = branch_flow_kernel_zygote(
-            polar.topology.yff_re, polar.topology.yft_re, polar.topology.ytf_re, polar.topology.ytt_re,
-            polar.topology.yff_im, polar.topology.yft_im, polar.topology.ytf_im, polar.topology.ytt_im,
-            fr_vmag, to_vmag,
-            cosθ, sinθ
-        )
-        return dot(weights, cons)
-    end
-    grad = Zygote.gradient(lumping, fr_vmag, to_vmag, fr_vang, to_vang)
-    gvmag = @view cons_grad[1:nbus]
-    gvang = @view cons_grad[nbus+1:2*nbus]
-    # This is basically the adjoint of the statements above (1). = becomes +=.
-    ev_vmag = kernel!(gvmag, grad[1], f, ndrange = nbus)
-    ev_vang = kernel!(gvang, grad[3], f, ndrange = nbus)
-    wait(ev_vmag)
-    wait(ev_vang)
-    ev_vmag = kernel!(gvmag, grad[2], t, ndrange = nbus)
-    ev_vang = kernel!(gvang, grad[4], t, ndrange = nbus)
-    wait(ev_vmag)
-    wait(ev_vang)
+    PT = polar.topology
+    fill!(cons_grad, 0)
+    adj_vmag = @view cons_grad[1:nbus]
+    adj_vang = @view cons_grad[nbus+1:2*nbus]
+    ev = adj_branch_flow_kernel!(polar.device)(weights, buffer.vmag, adj_vmag, 
+            buffer.vang, adj_vang,
+            PT.yff_re, PT.yft_re, PT.ytf_re, PT.ytt_re,
+            PT.yff_im, PT.yft_im, PT.ytf_im, PT.ytt_im,
+            PT.f_buses, PT.t_buses, nlines, ndrange = nlines
+    )
+    wait(ev)
     return cons_grad
 end
 

src/Polar/Constraints/power_balance.jl

@@ -11,8 +11,8 @@ function _power_balance!(
         kernel! = residual_kernel!(KA.CUDADevice())
     end
     ev = kernel!(F, v_m, v_a,
-                 ybus_re.nzval, ybus_re.colptr, ybus_re.rowval,
-                 ybus_im.nzval,
+                 ybus_re.colptr, ybus_re.rowval,
+                 ybus_re.nzval, ybus_im.nzval,
                  pinj, qinj, pv, pq, nbus,
                  ndrange=npv+npq)
     wait(ev)
@@ -49,6 +49,32 @@ function bounds(polar::PolarForm{T, IT, VT, MT}, ::typeof(power_balance)) where
     return xzeros(VT, m) , xzeros(VT, m)
 end
 
+# Adjoint
+function adjoint!(
+    polar::PolarForm,
+    ::typeof(power_balance),
+    cons, ∂cons,
+    vm, ∂vm,
+    va, ∂va,
+    pinj, ∂pinj,
+    qinj, ∂qinj,
+)
+    nbus = get(polar, PS.NumberOfBuses())
+    ref = polar.indexing.index_ref
+    pv = polar.indexing.index_pv
+    pq = polar.indexing.index_pq
+    ybus_re, ybus_im = get(polar.topology, PS.BusAdmittanceMatrix())
+
+    adj_residual_polar!(
+        cons, ∂cons,
+        vm, ∂vm,
+        va, ∂va,
+        ybus_re, ybus_im,
+        pinj, ∂pinj,
+        qinj, pv, pq, nbus,
+    )
+end
+
 function matpower_jacobian(polar::PolarForm, ::State, ::typeof(power_balance), V)
     nbus = get(polar, PS.NumberOfBuses())
     pf = polar.network

src/Polar/Constraints/reactive_power.jl

@@ -56,6 +56,36 @@ function bounds(polar::PolarForm{T, IT, VT, MT}, ::typeof(reactive_power_constra
     return convert(VT, q_min), convert(VT, q_max)
 end
 
+# Adjoint
+function adjoint!(
+    polar::PolarForm,
+    ::typeof(reactive_power_constraints),
+    cons, ∂cons,
+    vm, ∂vm,
+    va, ∂va,
+    pinj, ∂pinj,
+    qinj, ∂qinj,
+)
+    nbus = get(polar, PS.NumberOfBuses())
+    ref = polar.indexing.index_ref
+    pv = polar.indexing.index_pv
+    pq = polar.indexing.index_pq
+    pv_to_gen = polar.indexing.index_pv_to_gen
+    ref_to_gen = polar.indexing.index_ref_to_gen
+    ybus_re, ybus_im = get(polar.topology, PS.BusAdmittanceMatrix())
+
+    adj_reactive_power!(
+        cons, ∂cons,
+        vm, ∂vm,
+        va, ∂va,
+        ybus_re, ybus_im,
+        pinj, ∂pinj,
+        qinj, ∂qinj,
+        polar.reactive_load,
+        pv, pq, ref, pv_to_gen, ref_to_gen, nbus,
+    )
+end
+
 # Jacobian
 function jacobian(
     polar::PolarForm,

src/Polar/derivatives.jl

@@ -125,6 +125,87 @@ function AutoDiff.jacobian!(polar::PolarForm, jac::AutoDiff.Jacobian, buffer)
     return jac.J
 end
 
+#structure, F, vm, va, ybus_re, ybus_im, pinj, qinj, pv, pq, ref)
+function AutoDiff.Hessian(polar::PolarForm{T, VI, VT, MT}, func) where {T, VI, VT, MT}
+    @assert is_constraint(func)
+
+    if isa(polar.device, CPU)
+        A = Vector
+    elseif isa(polar.device, CUDADevice)
+        A = CUDA.CuVector
+    end
+
+    pf = polar.network
+    nbus = PS.get(pf, PS.NumberOfBuses())
+    n_cons = size_constraint(polar, func)
+
+    map = VI(polar.hessianstructure.map)
+    nmap = length(map)
+
+    x = VT(zeros(Float64, 4*nbus))
+
+    t1s{N} = ForwardDiff.Dual{Nothing,Float64, N} where N
+    t1sx = A{t1s{1}}(x)
+    t1sF = A{t1s{1}}(zeros(Float64, n_cons))
+    t1sseeds = A{ForwardDiff.Partials{1,Float64}}(undef, nmap)
+    varx = view(x, map)
+    t1svarx = view(t1sx, map)
+    VP = typeof(t1sseeds)
+    VD = typeof(t1sx)
+    return AutoDiff.Hessian{typeof(func), VI, VT, MT, Nothing, VP, VD, typeof(varx), typeof(t1svarx)}(
+        func, t1sseeds, t1sF, x, t1sx, map, varx, t1svarx
+    )
+end
+
+# λ' * H * v
+function AutoDiff.adj_hessian_prod!(
+    polar, H::AutoDiff.Hessian, buffer, λ, v,
+)
+    nbus = get(polar, PS.NumberOfBuses())
+    x = H.x
+    ntgt = length(v)
+    t1sx = H.t1sx
+    adj_t1sx = similar(t1sx)
+    t1sF = H.t1sF
+    adj_t1sF = similar(t1sF)
+    # Move data
+    x[1:nbus] .= buffer.vmag
+    x[nbus+1:2*nbus] .= buffer.vang
+    x[2*nbus+1:3*nbus] .= buffer.pinj
+    x[3*nbus+1:4*nbus] .= buffer.qinj
+    # Init dual variables
+    t1sx .= H.x
+    adj_t1sx .= 0.0
+    t1sF .= 0.0
+    adj_t1sF .= λ
+    # Seeding
+    nmap = length(H.map)
+
+    # Init seed
+    for i in 1:nmap
+        H.t1sseeds[i] = ForwardDiff.Partials{1, Float64}(NTuple{1, Float64}(v[i]))
+    end
+    AutoDiff.seed!(H.t1sseeds, H.varx, H.t1svarx, nbus)
+
+    adjoint!(
+        polar, H.func,
+        t1sF, adj_t1sF,
+        view(t1sx, 1:nbus), view(adj_t1sx, 1:nbus), # vmag
+        view(t1sx, nbus+1:2*nbus), view(adj_t1sx, nbus+1:2*nbus), # vang
+        view(t1sx, 2*nbus+1:3*nbus), view(adj_t1sx, 2*nbus+1:3*nbus), # pinj
+        view(t1sx, 3*nbus+1:4*nbus), view(adj_t1sx, 3*nbus+1:4*nbus), # qinj
+    )
+
+    # TODO, this is redundant
+    ps = ForwardDiff.partials.(adj_t1sx[H.map])
+    res = similar(v)
+    res .= 0.0
+    for i in 1:length(ps)
+        res[i] = ps[i].values[1]
+    end
+    return res
+end
+
 ## Utils
 # Expression of Jacobians from MATPOWER
 function _matpower_residual_jacobian(V, Ybus)