[skip ci] clean CUDA deps

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"

src/ExaPF.jl

@@ -6,9 +6,6 @@ using LinearAlgebra
 using SparseArrays
 
 import CUDA
-import CUDA.CUBLAS
-import CUDA.CUSPARSE
-import CUDA.CUSOLVER
 
 import ForwardDiff
 using KernelAbstractions
@@ -37,4 +34,9 @@ const LS = LinearSolvers
 # Polar formulation
 include("Polar/polar.jl")
 
+# CUDA extension
+if CUDA.has_cuda()
+    include("cuda_wrapper.jl")
+end
+
 end

src/Polar/polar.jl

@@ -15,10 +15,6 @@ end
 function PolarForm(pf::PS.PowerNetwork, device::KA.CPU)
     return PolarForm{Float64, Vector{Int}, Vector{Float64}, Matrix{Float64}}(pf, device)
 end
-function PolarForm(pf::PS.PowerNetwork, device::KA.GPU)
-    return PolarForm{Float64, CuVector{Int}, CuVector{Float64}, CuMatrix{Float64}}(pf, device)
-end
-
 # Convenient constructor
 PolarForm(datafile::String, device) = PolarForm(PS.PowerNetwork(datafile), device)
 

src/architectures.jl

@@ -3,7 +3,6 @@ abstract type AbstractArchitecture end
 
 # norm
 xnorm(x::AbstractVector) = norm(x, 2)
-xnorm(x::CUDA.CuVector) = CUBLAS.nrm2(x)
 
 xnorm_inf(a) = maximum(abs.(a))
 
@@ -15,15 +14,3 @@ function get_jacobian_types(::CPU)
     return SMT, A
 end
 
-function get_jacobian_types(::GPU)
-    SMT = CUSPARSE.CuSparseMatrixCSR{Float64}
-    A = CUDA.CuVector
-    return SMT, A
-end
-
-function Base.unsafe_wrap(Atype::Type{CUDA.CuArray{T, 1, CUDA.Mem.DeviceBuffer}},
-                          p::CUDA.CuPtr{T}, dim::Integer;
-                          own::Bool=false, ctx::CUDA.CuContext=CUDA.context()) where {T}
-    unsafe_wrap(CUDA.CuArray{T, 1}, p, (dim,); own, ctx)
-end
-

src/autodiff.jl

@@ -4,7 +4,6 @@ module AutoDiff
 using SparseArrays
 
 using CUDA
-import CUDA.CUSPARSE
 import ForwardDiff
 import SparseDiffTools
 using KernelAbstractions

src/cuda_wrapper.jl

@@ -0,0 +1,50 @@
+
+import CUDA.CUBLAS
+import CUDA.CUSPARSE
+import CUDA.CUSOLVER
+
+using CUDAKernels
+
+function PolarForm(pf::PS.PowerNetwork, device::CUDADevice)
+    return PolarForm{Float64, CuVector{Int}, CuVector{Float64}, CuMatrix{Float64}}(pf, device)
+end
+
+default_sparse_matrix(::CUDADevice) = CuSparseMatrixCSR{Float64, Int}
+xnorm(x::CUDA.CuVector) = CUBLAS.nrm2(x)
+
+function get_jacobian_types(::CUDADevice)
+    SMT = default_sparse_matrix(CUDADevice())
+    A = CUDA.CuVector
+    return SMT, A
+end
+
+function Base.unsafe_wrap(Atype::Type{CUDA.CuArray{T, 1, CUDA.Mem.DeviceBuffer}},
+                          p::CUDA.CuPtr{T}, dim::Integer;
+                          own::Bool=false, ctx::CUDA.CuContext=CUDA.context()) where {T}
+    unsafe_wrap(CUDA.CuArray{T, 1}, p, (dim,); own, ctx)
+end
+
+# Differentiable LinearAlgebra.mul! for ForwardDiff
+@kernel function _spmm_kernel!(Y, X, colVal, rowPtr, nzVal, alpha, beta, n, m)
+    i, k = @index(Global, NTuple)
+    Y[i, k] *= beta
+    @inbounds for c in rowPtr[i]:rowPtr[i+1]-1
+        j = colVal[c]
+        Y[i, k] += alpha * nzVal[c] * X[j, k]
+    end
+end
+
+function LinearAlgebra.mul!(Y::AbstractArray{T, 2}, A::CUSPARSE.CuSparseMatrixCSR, X::AbstractArray{T, 2}, alpha::Number, beta::Number) where {T <: ForwardDiff.Dual}
+    n, m = size(A)
+    p = size(X, 2)
+    @assert size(Y, 1) == n
+    @assert size(X, 1) == m
+    @assert size(X, 2) == size(Y, 2)
+
+    ndrange = (n, p)
+    ev = _spmm_kernel!(CUDADevice())(
+        Y, X, A.colVal, A.rowPtr, A.nzVal, alpha, beta, n, m,
+        ndrange=ndrange,
+    )
+    wait(ev)
+end

test/runtests.jl

@@ -19,9 +19,6 @@ const BENCHMARK_DIR = joinpath(dirname(@__FILE__), "..", "benchmark")
 const CASES = ["case9.m", "case30.m"]
 
 ARCHS = Any[(CPU(), Array, SparseMatrixCSC)]
-if has_cuda_gpu()
-    include("gpu.jl")
-end
 
 # Load test modules
 @isdefined(TestLinearSolvers)    || include("TestLinearSolvers.jl")