Add faster GMRES version

Author: avik-pal

.JuliaFormatter.toml

@@ -1,2 +1,3 @@
 style = "sciml"
-format_markdown = true
+format_markdown = true
+annotate_untyped_fields_with_any = false

Project.toml

@@ -5,6 +5,7 @@ version = "2.5.1"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
+ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 EnumX = "4e289a0a-7415-4d19-859d-a7e5c4648b56"
 FastLapackInterface = "29a986be-02c6-4525-aec4-84b980013641"
@@ -35,7 +36,6 @@ IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153"
 KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77"
 MKL_jll = "856f044c-d86e-5d09-b602-aeab76dc8ba7"
 Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
-NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 Pardiso = "46dd5b70-b6fb-5a00-ae2d-e8fea33afaf2"
 
 [extensions]
@@ -61,7 +61,6 @@ IterativeSolvers = "0.9.2"
 KLU = "0.3.0, 0.4"
 Krylov = "0.9"
 KrylovKit = "0.5, 0.6"
-NNlib = "0.9"
 PrecompileTools = "1"
 Preferences = "1"
 RecursiveFactorization = "0.2.8"
@@ -83,10 +82,9 @@ IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77"
 MKL_jll = "856f044c-d86e-5d09-b602-aeab76dc8ba7"
-Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
+Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
 MultiFloats = "bdf0d083-296b-4888-a5b6-7498122e68a5"
-NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"

ext/LinearSolveBlockDiagonalsExt.jl

@@ -2,8 +2,8 @@ module LinearSolveBlockDiagonalsExt
 
 using LinearSolve, BlockDiagonals
 
-function LinearSolve.init_cacheval(alg::SimpleGMRES{false}, A::BlockDiagonal, b, u, Pl, Pr,
-    maxiters::Int, abstol, reltol, verbose, assumptions; zeroinit = true)
+function LinearSolve.init_cacheval(alg::SimpleGMRES{false}, A::BlockDiagonal, b, args...;
+    kwargs...)
     @assert ndims(A) == 2 "ndims(A) == $(ndims(A)). `A` must have ndims == 2."
     # We need to perform this check even when `zeroinit == true`, since the type of the
     # cache is dependent on whether we are able to use the specialized dispatch.
@@ -17,8 +17,8 @@ function LinearSolve.init_cacheval(alg::SimpleGMRES{false}, A::BlockDiagonal, b,
         end
     end
     # Can't help but perform dynamic dispatch here
-    return LinearSolve._init_cacheval(Val(uniform_blocks), alg, A, b, u, Pl, Pr, maxiters,
-        abstol, reltol, verbose, assumptions; zeroinit, blocksize = usize)
+    return LinearSolve._init_cacheval(Val(uniform_blocks), alg, A, b, args...;
+        blocksize = usize, kwargs...)
 end
 
 end

ext/LinearSolveNNlibExt.jl

@@ -29,6 +29,7 @@ using LinearAlgebra.LAPACK: require_one_based_indexing, chkfinite, chkstride1,
 
 import GPUArraysCore
 import Preferences
+import ConcreteStructs: @concrete
 
 # wrap
 import Krylov

src/LinearSolve.jl

@@ -253,7 +253,7 @@ function SciMLBase.solve!(cache::LinearCache, alg::KrylovJL; kwargs...)
         Krylov.solve!(args...; M = M,
             kwargs...)
     elseif cache.cacheval isa Krylov.GmresSolver
-        Krylov.solve!(args...; M = M, N = N,
+        Krylov.solve!(args...; M = M, N = N, restart = alg.gmres_restart > 0,
             kwargs...)
     elseif cache.cacheval isa Krylov.BicgstabSolver
         Krylov.solve!(args...; M = M, N = N,