Allow vm and vn to be of different types

For problems `Ax=b`, this allows the vector space of `x` (`vn`) to be
different from the vector space of `b` (`vm`); also for some bipartite
problems

    [A B; C D] [x; y] = [b; c]

it allows `x` and `y` to be of different types.

The key change is to introduce `KrylovWorkspaceNext{T,FC,Sm,Sn}`, for
which `Sm` and `Sn` are the types of the length-`m` and length-`n`
vectors, respectively. `KrylovWorkspace{T,FC,S}` is now just a type
alias for `KrylovWorkspaceNext{T,FC,S,S}`. This is done for reasons of
backwards compatibility; in the next breaking release, we should rename
`KrylovWorkspaceNext` to `KrylovWorkspace`.

This PR exploits the new flexibility for CGLS, LSQR, TRICG, TRIMR, and
GPMR. Extending to other solvers is left as an exercise for the reader.

Fixes #1037

Author: timholy

src/cgls.jl

@@ -126,7 +126,7 @@ args_cgls = (:A, :b)
 kwargs_cgls = (:M, :ldiv, :radius, :λ, :atol, :rtol, :itmax, :timemax, :verbose, :history, :callback, :iostream)
 
 @eval begin
-  function cgls!(workspace :: CglsWorkspace{T,FC,S}, $(def_args_cgls...); $(def_kwargs_cgls...)) where {T <: AbstractFloat, FC <: FloatOrComplex{T}, S <: AbstractVector{FC}}
+  function cgls!(workspace :: CglsWorkspace{T,FC,Sm,Sn}, $(def_args_cgls...); $(def_kwargs_cgls...)) where {T <: AbstractFloat, FC <: FloatOrComplex{T}, Sm <: AbstractVector{FC}, Sn <: AbstractVector{FC}}
 
     # Timer
     start_time = time_ns()
@@ -142,13 +142,13 @@ kwargs_cgls = (:M, :ldiv, :radius, :λ, :atol, :rtol, :itmax, :timemax, :verbose
 
     # Check type consistency
     eltype(A) == FC || @warn "eltype(A) ≠ $FC. This could lead to errors or additional allocations in operator-vector products."
-    ktypeof(b) == S || error("ktypeof(b) must be equal to $S")
+    ktypeof(b) == Sm || error("ktypeof(b) must be equal to $Sm")
 
     # Compute the adjoint of A
     Aᴴ = A'
 
     # Set up workspace.
-    allocate_if(!MisI, workspace, :Mr, S, workspace.r)  # The length of Mr is m
+    allocate_if(!MisI, workspace, :Mr, Sm, workspace.r)  # The length of Mr is m
     x, p, s, r, q, stats = workspace.x, workspace.p, workspace.s, workspace.r, workspace.q, workspace.stats
     rNorms, ArNorms = stats.residuals, stats.Aresiduals
     reset!(stats)

src/gpmr.jl

@@ -157,7 +157,7 @@ optargs_gpmr = (:x0, :y0)
 kwargs_gpmr = (:C, :D, :E, :F, :ldiv, :gsp, :λ, :μ, :reorthogonalization, :atol, :rtol, :itmax, :timemax, :verbose, :history, :callback, :iostream)
 
 @eval begin
-  function gpmr!(workspace :: GpmrWorkspace{T,FC,S}, $(def_args_gpmr...); $(def_kwargs_gpmr...)) where {T <: AbstractFloat, FC <: FloatOrComplex{T}, S <: AbstractVector{FC}}
+  function gpmr!(workspace :: GpmrWorkspace{T,FC,Sm,Sn}, $(def_args_gpmr...); $(def_kwargs_gpmr...)) where {T <: AbstractFloat, FC <: FloatOrComplex{T}, Sm <: AbstractVector{FC}, Sn <: AbstractVector{FC}}
 
     # Timer
     start_time = time_ns()
@@ -181,8 +181,8 @@ kwargs_gpmr = (:C, :D, :E, :F, :ldiv, :gsp, :λ, :μ, :reorthogonalization, :ato
     # Check type consistency
     eltype(A) == FC || @warn "eltype(A) ≠ $FC. This could lead to errors or additional allocations in operator-vector products."
     eltype(B) == FC || @warn "eltype(B) ≠ $FC. This could lead to errors or additional allocations in operator-vector products."
-    ktypeof(b) == S || error("ktypeof(b) must be equal to $S")
-    ktypeof(c) == S || error("ktypeof(c) must be equal to $S")
+    ktypeof(b) == Sm || error("ktypeof(b) must be equal to $Sm")
+    ktypeof(c) == Sn || error("ktypeof(c) must be equal to $Sn")
 
     # Determine λ and μ associated to generalized saddle point systems.
     gsp && (λ = one(FC) ; μ = zero(FC))
@@ -192,10 +192,10 @@ kwargs_gpmr = (:C, :D, :E, :F, :ldiv, :gsp, :λ, :μ, :reorthogonalization, :ato
     warm_start && (μ ≠ 0) && !FisI && error("Warm-start with right preconditioners is not supported.")
 
     # Set up workspace.
-    allocate_if(!CisI, workspace, :q , S, workspace.x)  # The length of q is m
-    allocate_if(!DisI, workspace, :p , S, workspace.y)  # The length of p is n
-    allocate_if(!EisI, workspace, :wB, S, workspace.x)  # The length of wB is m
-    allocate_if(!FisI, workspace, :wA, S, workspace.y)  # The length of wA is n
+    allocate_if(!CisI, workspace, :q , Sm, workspace.x)  # The length of q is m
+    allocate_if(!DisI, workspace, :p , Sn, workspace.y)  # The length of p is n
+    allocate_if(!EisI, workspace, :wB, Sm, workspace.x)  # The length of wB is m
+    allocate_if(!FisI, workspace, :wA, Sn, workspace.y)  # The length of wA is n
     wA, wB, dA, dB, Δx, Δy = workspace.wA, workspace.wB, workspace.dA, workspace.dB, workspace.Δx, workspace.Δy
     x, y, V, U, gs, gc = workspace.x, workspace.y, workspace.V, workspace.U, workspace.gs, workspace.gc
     zt, R, stats = workspace.zt, workspace.R, workspace.stats

src/interface.jl

@@ -138,7 +138,7 @@ for (workspace, krylov, args, def_args, optargs, def_optargs, kwargs, def_kwargs
 
       krylov_solve(::Val{Symbol($krylov)}, $(def_args...); $(def_kwargs...)) where {T <: AbstractFloat, FC <: FloatOrComplex{T}} = $(krylov)($(args...); $(kwargs...))
     end
-  elseif krylov in (:diom, :dqgmres, :fom, :gmres, :fgmres, :gpmr)
+  elseif krylov in (:diom, :dqgmres, :fom, :gmres, :fgmres)
     @eval begin
       function $(krylov)($(def_args...); memory::Int = 20, $(def_kwargs...)) where {T <: AbstractFloat, FC <: FloatOrComplex{T}}
         start_time = time_ns()
@@ -164,6 +164,35 @@ for (workspace, krylov, args, def_args, optargs, def_optargs, kwargs, def_kwargs
           return results(workspace)
         end
 
+        krylov_solve(::Val{Symbol($krylov)}, $(def_args...), $(def_optargs...); memory::Int = 20, $(def_kwargs...)) where {T <: AbstractFloat, FC <: FloatOrComplex{T}} = $(krylov)($(args...), $(optargs...); memory, $(kwargs...))
+      end
+    end
+  elseif krylov == :gpmr
+    @eval begin
+      function $(krylov)($(def_args...); memory::Int = 20, $(def_kwargs...)) where {T <: AbstractFloat, FC <: FloatOrComplex{T}}
+        start_time = time_ns()
+        workspace = $workspace(KrylovConstructor(b, c, similar(b, 0), similar(c, 0)); memory)
+        elapsed_time = start_time |> ktimer
+        timemax -= elapsed_time
+        $(krylov!)(workspace, $(args...); $(kwargs...))
+        workspace.stats.timer += elapsed_time
+        return results(workspace)
+      end
+
+      krylov_solve(::Val{Symbol($krylov)}, $(def_args...); memory::Int = 20, $(def_kwargs...)) where {T <: AbstractFloat, FC <: FloatOrComplex{T}} = $(krylov)($(args...); memory, $(kwargs...))
+
+      if !isempty($optargs)
+        function $(krylov)($(def_args...), $(def_optargs...); memory::Int = 20, $(def_kwargs...)) where {T <: AbstractFloat, FC <: FloatOrComplex{T}}
+          start_time = time_ns()
+          workspace = $workspace(KrylovConstructor(b, c, similar(b, 0), similar(c, 0)); memory)
+          warm_start!(workspace, $(optargs...))
+          elapsed_time = start_time |> ktimer
+          timemax -= elapsed_time
+          $(krylov!)(workspace, $(args...); $(kwargs...))
+          workspace.stats.timer += elapsed_time
+          return results(workspace)
+        end
+
         krylov_solve(::Val{Symbol($krylov)}, $(def_args...), $(def_optargs...); memory::Int = 20, $(def_kwargs...)) where {T <: AbstractFloat, FC <: FloatOrComplex{T}} = $(krylov)($(args...), $(optargs...); memory, $(kwargs...))
       end
     end
@@ -196,6 +225,35 @@ for (workspace, krylov, args, def_args, optargs, def_optargs, kwargs, def_kwargs
         krylov_solve(::Val{Symbol($krylov)}, $(def_args...), $(def_optargs...); window::Int = 5, $(def_kwargs...)) where {T <: AbstractFloat, FC <: FloatOrComplex{T}} = $(krylov)($(args...), $(optargs...); window, $(kwargs...))
       end
     end
+  elseif krylov in (:tricg, :trimr)
+    @eval begin
+      function $(krylov)($(def_args...); $(def_kwargs...)) where {T <: AbstractFloat, FC <: FloatOrComplex{T}}
+        start_time = time_ns()
+        workspace = $workspace(KrylovConstructor(b, c, similar(b, 0), similar(c, 0)))
+        elapsed_time = start_time |> ktimer
+        timemax -= elapsed_time
+        $(krylov!)(workspace, $(args...); $(kwargs...))
+        workspace.stats.timer += elapsed_time
+        return results(workspace)
+      end
+
+      krylov_solve(::Val{Symbol($krylov)}, $(def_args...); $(def_kwargs...)) where {T <: AbstractFloat, FC <: FloatOrComplex{T}} = $(krylov)($(args...); $(kwargs...))
+
+      if !isempty($optargs)
+        function $(krylov)($(def_args...), $(def_optargs...); $(def_kwargs...)) where {T <: AbstractFloat, FC <: FloatOrComplex{T}}
+          start_time = time_ns()
+          workspace = $workspace(KrylovConstructor(b, c, similar(b, 0), similar(c, 0)))
+          warm_start!(workspace, $(optargs...))
+          elapsed_time = start_time |> ktimer
+          timemax -= elapsed_time
+          $(krylov!)(workspace, $(args...); $(kwargs...))
+          workspace.stats.timer += elapsed_time
+          return results(workspace)
+        end
+
+        krylov_solve(::Val{Symbol($krylov)}, $(def_args...), $(def_optargs...); $(def_kwargs...)) where {T <: AbstractFloat, FC <: FloatOrComplex{T}} = $(krylov)($(args...), $(optargs...); $(kwargs...))
+      end
+    end
   else
     @eval begin
       function $(krylov)($(def_args...); $(def_kwargs...)) where {T <: AbstractFloat, FC <: FloatOrComplex{T}}

src/krylov_show.jl

@@ -63,13 +63,13 @@ end
 
 Statistics of `workspace` are displayed if `show_stats` is set to true.
 """
-function show(io :: IO, workspace :: Union{KrylovWorkspace{T,FC,S}, BlockKrylovWorkspace{T,FC,S}}; show_stats :: Bool=true) where {T <: AbstractFloat, FC <: FloatOrComplex{T}, S <: AbstractVector{FC}}
+function show(io :: IO, workspace :: Union{KrylovWorkspaceNext{T,FC,Sm,S}, BlockKrylovWorkspace{T,FC,S}}; show_stats :: Bool=true) where {T <: AbstractFloat, FC <: FloatOrComplex{T}, S <: AbstractVector{FC}, Sm <: AbstractVector{FC}}
   type_workspace = typeof(workspace)
   name_workspace = string(type_workspace.name.name)
   name_stats = string(typeof(workspace.stats).name.name)
   nbytes = sizeof(workspace)
   storage = format_bytes(nbytes)
-  architecture = S <: Vector ? "CPU" : "GPU"
+  architecture = S <: Vector ? "CPU" : "GPU"   # FIXME cannot assume that all non-Vector types are GPU types
   l1 = max(length(name_workspace), length(string(FC)) + 11)  # length("Precision: ") = 11
   nchar = type_workspace <: Union{CgLanczosShiftWorkspace, FomWorkspace, DiomWorkspace, DqgmresWorkspace, GmresWorkspace, FgmresWorkspace, GpmrWorkspace, BlockGmresWorkspace} ? 8 : 0  # length("Vector{}") = 8
   l2 = max(ndigits(workspace.m) + 7, length(architecture) + 14, length(string(S)) + nchar)  # length("nrows: ") = 7 and length("Architecture: ") = 14