add preallocation and block support machinery

Author: mohamed82008

src/aggregation.jl

@@ -1,4 +1,4 @@
-function smoothed_aggregation(A::SparseMatrixCSC{T,V},
+function smoothed_aggregation(A::SparseMatrixCSC{T,V}, ::Type{Val{bs}}=Val{1},
                         symmetry = HermitianSymmetry(),
                         strength = SymmetricStrength(),
                         aggregate = StandardAggregation(),
@@ -10,7 +10,7 @@ function smoothed_aggregation(A::SparseMatrixCSC{T,V},
                         max_coarse = 10,
                         diagonal_dominance = false,
                         keep = false,
-                        coarse_solver = Pinv) where {T,V}
+                        coarse_solver = Pinv) where {T,V,bs}
 
     n = size(A, 1)
     # B = kron(ones(n, 1), eye(1))
@@ -29,19 +29,23 @@ function smoothed_aggregation(A::SparseMatrixCSC{T,V},
 
     levels = Vector{Level{T,V}}()
     bsr_flag = false
+    w = MultiLevelWorkspace(Val{bs}, eltype(A))
 
     while length(levels) + 1 < max_levels && size(A, 1) > max_coarse
+        residual!(w, size(A, 1))
         A, B, bsr_flag = extend_hierarchy!(levels, strength, aggregate, smooth,
                                 improve_candidates, diagonal_dominance,
                                 keep, A, B, symmetry, bsr_flag)
+        coarse_x!(w, size(A, 1))
+        coarse_b!(w, size(A, 1))
         #=if size(A, 1) <= max_coarse
             break
         end=#
     end
     #=A, B = extend_hierarchy!(levels, strength, aggregate, smooth,
                             improve_candidates, diagonal_dominance,
                             keep, A, B, symmetry)=#
-    MultiLevel(levels, A, coarse_solver(A), presmoother, postsmoother)
+    MultiLevel(levels, A, coarse_solver(A), presmoother, postsmoother, w)
 end
 
 struct HermitianSymmetry

src/classical.jl

@@ -7,23 +7,29 @@ struct Solver{S,T,P,PS}
     max_coarse::Int64
 end
 
-function ruge_stuben(A::SparseMatrixCSC{Ti,Tv};
+function ruge_stuben(A::SparseMatrixCSC{Ti,Tv}, ::Type{Val{bs}}=Val{1};
                 strength = Classical(0.25),
                 CF = RS(),
                 presmoother = GaussSeidel(),
                 postsmoother = GaussSeidel(),
                 max_levels = 10,
-                max_coarse = 10) where {Ti,Tv}
+                max_coarse = 10,
+                coarse_solver = Pinv) where {Ti,Tv,bs}
 
     s = Solver(strength, CF, presmoother,
                 postsmoother, max_levels, max_levels)
 
     levels = Vector{Level{Ti,Tv}}()
+    w = MultiLevelWorkspace(Val{bs}, eltype(A))
 
     while length(levels) + 1 < max_levels && size(A, 1) > max_coarse
+        residual!(w, size(A, 1))
         A = extend_heirarchy!(levels, strength, CF, A)
+        coarse_x!(w, size(A, 1))
+        coarse_b!(w, size(A, 1))
     end
-    MultiLevel(levels, A, presmoother, postsmoother)
+
+    MultiLevel(levels, A, coarse_solver(A), presmoother, postsmoother, w)
 end
 
 function extend_heirarchy!(levels::Vector{Level{Ti,Tv}}, strength, CF, A::SparseMatrixCSC{Ti,Tv}) where {Ti,Tv}

src/multilevel.jl

@@ -4,12 +4,48 @@ struct Level{T,V}
     R::SparseMatrixCSC{T,V}
 end
 
-struct MultiLevel{S, Pre, Post, Ti, Tv}
+struct MultiLevel{S, Pre, Post, Ti, Tv, TW}
     levels::Vector{Level{Ti,Tv}}
     final_A::SparseMatrixCSC{Ti,Tv}
     coarse_solver::S
     presmoother::Pre
     postsmoother::Post
+    workspace::TW
+end
+
+struct MultiLevelWorkspace{TX, bs}
+    coarse_xs::Vector{TX}
+    coarse_bs::Vector{TX}
+    res_vecs::Vector{TX}
+end
+function MultiLevelWorkspace(::Type{Val{bs}}, ::Type{T}) where {bs, T<:Number}
+    if bs === 1
+        TX = Vector{T}
+    else
+        TX = Matrix{T}
+    end
+    MultiLevelWorkspace{TX, bs}(TX[], TX[], TX[])
+end
+function residual!(m::MultiLevelWorkspace{TX, bs}, n) where {TX, bs}
+    if bs === 1
+        push!(m.res_vecs, TX(undef, n))
+    else
+        push!(m.res_vecs, TX(undef, n, bs))
+    end
+end
+function coarse_x!(m::MultiLevelWorkspace{TX, bs}, n) where {TX, bs}
+    if bs === 1
+        push!(m.coarse_xs, TX(undef, n))
+    else
+        push!(m.coarse_xs, TX(undef, n, bs))
+    end
+end
+function coarse_b!(m::MultiLevelWorkspace{TX, bs}, n) where {TX, bs}
+    if bs === 1
+        push!(m.coarse_bs, TX(undef, n))
+    else
+        push!(m.coarse_bs, TX(undef, n, bs))
+    end
 end
 
 abstract type CoarseSolver end
@@ -19,8 +55,6 @@ struct Pinv{T} <: CoarseSolver
 end
 (p::Pinv)(x, b) = mul!(x, p.pinvA, b)
 
-MultiLevel(l::Vector{Level{Ti,Tv}}, A::SparseMatrixCSC{Ti,Tv}, presmoother, postsmoother) where {Ti,Tv} =
-    MultiLevel(l, A, Pinv(A), presmoother, postsmoother)
 Base.length(ml) = length(ml.levels) + 1
 
 function Base.show(io::IO, ml::MultiLevel)

test/gmg.jl

@@ -4,16 +4,20 @@ function multigrid(A::SparseMatrixCSC{T,V}; max_levels = 10, max_coarse = 10,
                     presmoother = GaussSeidel(), postsmoother = GaussSeidel()) where {T,V}
 
     levels = Vector{Level{T,V}}()
+    w = AlgebraicMultigrid.MultiLevelWorkspace(Val{1}, eltype(A))
 
     while length(levels) + 1 < max_levels && size(A, 1) > max_coarse
+        AlgebraicMultigrid.residual!(w, size(A, 1))
         A = extend!(levels, A)
+        AlgebraicMultigrid.coarse_x!(w, size(A, 1))
+        AlgebraicMultigrid.coarse_b!(w, size(A, 1))
         #=if size(A, 1) <= max_coarse
             # push!(levels, Level(A,spzeros(T,0,0),spzeros(T,0,0)))
             break
         end=#
     end
 
-    MultiLevel(levels, A, presmoother, postsmoother)
+    MultiLevel(levels, A, Pinv(A), presmoother, postsmoother, w)
 end
 
 function extend!(levels, A::SparseMatrixCSC{Ti,Tv}) where {Ti,Tv}