Use KA.jl for faster sym_givens

Author: Avik Pal

Project.toml

@@ -33,6 +33,7 @@ BlockDiagonals = "0a1fb500-61f7-11e9-3c65-f5ef3456f9f0"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 HYPRE = "b5ffcf37-a2bd-41ab-a3da-4bd9bc8ad771"
 IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153"
+KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77"
 MKL_jll = "856f044c-d86e-5d09-b602-aeab76dc8ba7"
 Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
@@ -43,6 +44,7 @@ LinearSolveBlockDiagonalsExt = "BlockDiagonals"
 LinearSolveCUDAExt = "CUDA"
 LinearSolveHYPREExt = "HYPRE"
 LinearSolveIterativeSolversExt = "IterativeSolvers"
+LinearSolveKernelAbstractionsExt = "KernelAbstractions"
 LinearSolveKrylovKitExt = "KrylovKit"
 LinearSolveMKLExt = "MKL_jll"
 LinearSolveMetalExt = "Metal"
@@ -58,6 +60,7 @@ GPUArraysCore = "0.1"
 HYPRE = "1.4.0"
 IterativeSolvers = "0.9.2"
 KLU = "0.3.0, 0.4"
+KernelAbstractions = "0.9"
 Krylov = "0.9"
 KrylovKit = "0.5, 0.6"
 PrecompileTools = "1"
@@ -79,6 +82,7 @@ HYPRE = "b5ffcf37-a2bd-41ab-a3da-4bd9bc8ad771"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
+KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77"
 MKL_jll = "856f044c-d86e-5d09-b602-aeab76dc8ba7"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"

ext/LinearSolveKernelAbstractionsExt.jl

@@ -0,0 +1,24 @@
+module LinearSolveKernelAbstractionsExt
+
+using LinearSolve, KernelAbstractions
+
+LinearSolve.__is_extension_loaded(::Val{:KernelAbstractions}) = true
+
+using GPUArraysCore
+
+function LinearSolve._fast_sym_givens!(c, s, R, nr::Int, inner_iter::Int, bsize::Int, Hbis)
+    backend = get_backend(Hbis)
+    kernel! = __fast_sym_givens_kernel!(backend)
+    kernel!(c[inner_iter], s[inner_iter], R[nr + inner_iter], Hbis; ndrange=bsize)
+    return c, s, R
+end
+
+@kernel function __fast_sym_givens_kernel!(c, s, R, @Const(Hbis))
+    idx = @index(Global)
+    @inbounds _c, _s, _ρ = LinearSolve._sym_givens(R[idx], Hbis[idx])
+    @inbounds c[idx] = _c
+    @inbounds s[idx] = _s
+    @inbounds R[idx] = _ρ
+end
+
+end

src/LinearSolve.jl

@@ -56,6 +56,11 @@ _isidentity_struct(λ::Number) = isone(λ)
 _isidentity_struct(A::UniformScaling) = isone(A.λ)
 _isidentity_struct(::SciMLOperators.IdentityOperator) = true
 
+# Dispatch Friendly way to check if an extension is loaded
+__is_extension_loaded(::Val) = false
+
+function _fast_sym_givens! end
+
 # Code
 
 const INCLUDE_SPARSE = Preferences.@load_preference("include_sparse", Base.USE_GPL_LIBS)

src/simplegmres.jl

@@ -105,6 +105,19 @@ function _sym_givens(a::T, b::T) where {T <: AbstractFloat}
     return (c, s, ρ)
 end
 
+function _sym_givens!(c, s, R, nr::Int, inner_iter::Int, bsize::Int, Hbis)
+    if __is_extension_loaded(Val(:KernelAbstractions))
+        return _fast_sym_givens!(c, s, R, nr, inner_iter, bsize, Hbis)
+    end
+    __res = _sym_givens.(R[nr + inner_iter], Hbis)
+    GPUArraysCore.@allowscalar foreach(1:bsize) do i
+        c[inner_iter][i] = __res[i][1]
+        s[inner_iter][i] = __res[i][2]
+        R[nr + inner_iter][i] = __res[i][3]
+    end
+    return c, s, R
+end
+
 _no_preconditioner(::Nothing) = true
 _no_preconditioner(::IdentityOperator) = true
 _no_preconditioner(::UniformScaling) = true
@@ -221,15 +234,7 @@ function SciMLBase.solve!(cache::SimpleGMRESCache{false}, lincache::LinearCache)
     while !(solved || tired || breakdown)
         # Initialize workspace.
         nr = 0  # Number of coefficients stored in Rₖ.
-        #=  TODO: Check that not zeroing out doesn't lead to incorrect results.
-        foreach(V) do v
-            v .= zero(T)  # Orthogonal basis of Kₖ(MAN, Mr₀).
-        end
-        s .= zero(T)  # Givens sines used for the factorization QₖRₖ = Hₖ₊₁.ₖ.
-        c .= zero(T)  # Givens cosines used for the factorization QₖRₖ = Hₖ₊₁.ₖ.
-        R .= zero(T)  # Upper triangular matrix Rₖ.
-        z .= zero(T)  # Right-hand of the least squares problem min ‖Hₖ₊₁.ₖyₖ - βe₁‖₂.
-        =#
+        # TODO: Check that not zeroing out doesn't lead to incorrect results.
 
         if restart
             xr .= zero(T)  # xr === Δx when restart is set to true
@@ -517,13 +522,7 @@ function SciMLBase.solve!(cache::SimpleGMRESCache{true}, lincache::LinearCache)
             # Compute and apply current Givens reflection Ωₖ.
             # [cₖ  sₖ] [ r̄ₖ.ₖ ] = [rₖ.ₖ]
             # [s̄ₖ -cₖ] [hₖ₊₁.ₖ]   [ 0  ]
-            # FIXME: Write inplace kernel
-            __res = _sym_givens.(R[nr + inner_iter], Hbis)
-            foreach(1:bsize) do i
-                c[inner_iter][i] = __res[i][1]
-                s[inner_iter][i] = __res[i][2]
-                R[nr + inner_iter][i] = __res[i][3]
-            end
+            _sym_givens!(c, s, R, nr, inner_iter, bsize, Hbis)
 
             # Update zₖ = (Qₖ)ᴴβe₁
             ζₖ₊₁ = conj.(s[inner_iter]) .* z[inner_iter]
@@ -567,15 +566,8 @@ function SciMLBase.solve!(cache::SimpleGMRESCache{true}, lincache::LinearCache)
                 pos = pos - j + 1            # position of rᵢ.ⱼ₋₁
             end
             # Rₖ can be singular if the system is inconsistent
-            # FIXME: Write with broadcasting
-            GPUArraysCore.@allowscalar foreach(1:bsize) do B
-                if abs(R[pos][B]) ≤ btol
-                    y[i][B] = zero(T)
-                    inconsistent = true
-                else
-                    y[i][B] /= R[pos][B]
-                end
-            end
+            y[i] .= ifelse.(abs.(R[pos]) .≤ btol, zero(T), y[i] ./ R[pos])  # yᵢ ← yᵢ / rᵢᵢ
+            inconsistent = any(abs.(R[pos]) .≤ btol)
         end
 
         # Form xₖ = NVₖyₖ