Use mul! in SYMMLQ

Author: amontoison

src/krylov_solvers.jl

@@ -133,18 +133,22 @@ The outer constructors
 may be used in order to create these vectors.
 """
 mutable struct SymmlqSolver{T,S} <: KrylovSolver{T,S}
-  x     :: S
-  Mvold :: S
-  Mv    :: S
-  w̅     :: S
+  x       :: S
+  Mvold   :: S
+  Mv      :: S
+  Mv_next :: S
+  w̅       :: S
+  v       :: Union{S, Nothing}
 
   function SymmlqSolver(n, m, S)
-    T     = eltype(S)
-    x     = S(undef, n)
-    Mvold = S(undef, n)
-    Mv    = S(undef, n)
-    w̅     = S(undef, n)
-    solver = new{T,S}(x, Mvold, Mv, w̅)
+    T       = eltype(S)
+    x       = S(undef, n)
+    Mvold   = S(undef, n)
+    Mv      = S(undef, n)
+    Mv_next = S(undef, n)
+    w̅       = S(undef, n)
+    v       = nothing
+    solver = new{T,S}(x, Mvold, Mv, Mv_next, w̅, v)
     return solver
   end
 

src/symmlq.jl

@@ -48,19 +48,22 @@ function symmlq!(solver :: SymmlqSolver{T,S}, A, b :: AbstractVector{T};
 
   m, n = size(A)
   m == n || error("System must be square")
-  size(b, 1) == m || error("Inconsistent problem size")
+  length(b) == m || error("Inconsistent problem size")
   (verbose > 0) && @printf("SYMMLQ: system of size %d\n", n)
 
-  # Test M == Iₘ
-  MisI = isa(M, opEye)
+  # Tests M == Iₙ
+  MisI = isa(M, opEye) || (M == I)
 
   # Check type consistency
   eltype(A) == T || error("eltype(A) ≠ $T")
   ktypeof(b) == S || error("ktypeof(b) ≠ $S")
   MisI || (eltype(M) == T) || error("eltype(M) ≠ $T")
 
   # Set up workspace.
-  x, Mvold, Mv, w̅ = solver.x, solver.Mvold, solver.Mv, solver.w̅
+  !MisI && isnothing(solver.v) && (solver.v = S(undef, n))
+  x, Mvold, Mv, Mv_next, w̅ = solver.x, solver.Mvold, solver.Mv, solver.Mv_next, solver.w̅
+  v = MisI ? Mv : solver.v
+  vold = MisI ? Mvold : solver.v
 
   ϵM = eps(T)
   x .= zero(T)
@@ -69,7 +72,7 @@ function symmlq!(solver :: SymmlqSolver{T,S}, A, b :: AbstractVector{T};
   # Initialize Lanczos process.
   # β₁ M v₁ = b.
   Mvold .= b
-  vold = M * Mvold
+  MisI || mul!(vold, M, Mvold)
   β₁ = @kdot(m, vold, Mvold)
   β₁ == 0 && return (x, SimpleStats(true, true, [zero(T)], [zero(T)], "x = 0 is a zero-residual solution"))
   β₁ = sqrt(β₁)
@@ -79,10 +82,10 @@ function symmlq!(solver :: SymmlqSolver{T,S}, A, b :: AbstractVector{T};
 
   w̅ .= vold
 
-  Mv .= A * vold
+  mul!(Mv, A, vold)
   α = @kdot(m, vold, Mv) + λ
   @kaxpy!(m, -α, Mvold, Mv)  # Mv = Mv - α * Mvold
-  v = M * Mv
+  MisI || mul!(v, M, Mv)
   β = @kdot(m, v, Mv)
   β < 0 && error("Preconditioner is not positive definite")
   β = sqrt(β)
@@ -177,13 +180,13 @@ function symmlq!(solver :: SymmlqSolver{T,S}, A, b :: AbstractVector{T};
 
     # Generate next Lanczos vector
     oldβ = β
-    Mv_next = A * v
+    mul!(Mv_next, A, v)
     α = @kdot(m, v, Mv_next) + λ
     @kaxpy!(m, -oldβ, Mvold, Mv_next)
     @. Mvold = Mv
     @kaxpy!(m, -α, Mv, Mv_next)
     @. Mv = Mv_next
-    v = M * Mv
+    MisI || mul!(v, M, Mv)
     β = @kdot(m, v, Mv)
     β < 0 && error("Preconditioner is not positive definite")
     β = sqrt(β)

src/variants.jl

@@ -34,7 +34,7 @@ for fn in (:cgls, :cgne, :lnlq, :craig, :craigmr, :crls, :crmr, :lslq, :lsqr, :l
 end
 
 # Variants where matrix-vector products with A are only required
-for fn in (:symmlq, :cgs, :bicgstab)
+for fn in (:cgs, :bicgstab)
   @eval begin
     $fn(A :: AbstractMatrix{T}, b :: AbstractVector{T}; kwargs...) where T <: AbstractFloat =
       $fn(PreallocatedLinearOperator(A, storagetype=ktypeof(b), symmetric=true), b; wrap_preconditioners(kwargs, ktypeof(b))...)

test/test_alloc.jl

@@ -23,18 +23,19 @@ function test_alloc()
   UniformScaling_bytes = @allocated cg(L, b, M=M2)
   @test 0.99 * UniformScaling_bytes ≤ opEye_bytes ≤ 1.01 * UniformScaling_bytes
 
-  # without preconditioner and with Ap preallocated, SYMMLQ needs 4 n-vectors: x_lq, vold, v, w̅ (= x_cg)
-  storage_symmlq(n) = 4 * n
+  # SYMMLQ needs:
+  # 5 n-vectors: x, Mvold, Mv, Mv_next, w̅
+  storage_symmlq(n) = 5 * n
   storage_symmlq_bytes(n) = 8 * storage_symmlq(n)
 
   expected_symmlq_bytes = storage_symmlq_bytes(n)
-  symmlq(A, b)  # warmup
-  actual_symmlq_bytes = @allocated symmlq(A, b)
+  symmlq(L, b)  # warmup
+  actual_symmlq_bytes = @allocated symmlq(L, b)
   @test actual_symmlq_bytes ≤ 1.1 * expected_symmlq_bytes
 
-  solver = SymmlqSolver(A, b)
-  symmlq!(solver, A, b)  # warmup
-  inplace_symmlq_bytes = @allocated symmlq!(solver, A, b)
+  solver = SymmlqSolver(L, b)
+  symmlq!(solver, L, b)  # warmup
+  inplace_symmlq_bytes = @allocated symmlq!(solver, L, b)
   @test (VERSION < v"1.5") || (inplace_symmlq_bytes == 672)
 
   # CG needs: