Fix performance regression and add some timing calls for future use

Project.toml

@@ -9,11 +9,13 @@ LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 
 [compat]
 CommonSolve = "0.2"
 LinearSolve = "2, 3"
 Reexport = "1"
+TimerOutputs = "0.5.29"
 julia = "1.6"
 
 [extras]

src/AlgebraicMultigrid.jl

@@ -1,5 +1,6 @@
 module AlgebraicMultigrid
 
+using TimerOutputs
 using Reexport
 using LinearAlgebra
 using LinearSolve

src/aggregation.jl

@@ -13,82 +13,119 @@ function smoothed_aggregation(A::TA,
                         diagonal_dominance = false,
                         keep = false,
                         coarse_solver = Pinv, kwargs...) where {T,V,bs,TA<:SparseMatrixCSC{T,V}}
+
+    @timeit_debug "prologue" begin
+
     n = size(A, 1)
     B = isnothing(B) ? ones(T,n) : copy(B)
     @assert size(A, 1) == size(B, 1)
 
-    #=max_levels, max_coarse, strength =
-        levelize_strength_or_aggregation(max_levels, max_coarse, strength)
-    max_levels, max_coarse, aggregate =
-        levelize_strength_or_aggregation(max_levels, max_coarse, aggregate)
-
-    improve_candidates =
-        levelize_smooth_or_improve_candidates(improve_candidates, max_levels)=#
-    # str = [stength for _ in 1:max_levels - 1]
-    # agg = [aggregate for _ in 1:max_levels - 1]
-    # sm = [smooth for _ in 1:max_levels]
-
     levels = Vector{Level{TA, TA, Adjoint{T, TA}}}()
     bsr_flag = false
     w = MultiLevelWorkspace(Val{bs}, eltype(A))
     residual!(w, size(A, 1))
 
+    end
+
     while length(levels) + 1 < max_levels && size(A, 1) > max_coarse
-        A, B, bsr_flag = extend_hierarchy!(levels, strength, aggregate, smooth,
+        @timeit_debug "extend_hierarchy!" A, B, bsr_flag = extend_hierarchy_sa!(levels, strength, aggregate, smooth,
                                 improve_candidates, diagonal_dominance,
                                 keep, A, B, symmetry, bsr_flag)
+        size(A, 1) == 0 && break
         coarse_x!(w, size(A, 1))
         coarse_b!(w, size(A, 1))
-        #=if size(A, 1) <= max_coarse
-            break
-        end=#
         residual!(w, size(A, 1))
     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, w)
+
+    @timeit_debug "coarse solver setup" cs = coarse_solver(A)
+    @timeit_debug "ml setup" ml = MultiLevel(levels, A, cs, presmoother, postsmoother, w)
+    return ml
 end
 
 struct HermitianSymmetry
 end
 
-function extend_hierarchy!(levels, strength, aggregate, smooth,
+function extend_hierarchy_sa!(levels, strength, aggregate, smooth,
                             improve_candidates, diagonal_dominance, keep,
                             A, B,
                             symmetry, bsr_flag)
 
     # Calculate strength of connection matrix
-    if symmetry isa HermitianSymmetry
+    @timeit_debug "strength" if symmetry isa HermitianSymmetry
         S, _T = strength(A, bsr_flag)
     else
         S, _T = strength(adjoint(A), bsr_flag)
     end
 
     # Aggregation operator
-    AggOp = aggregate(S)
+    @timeit_debug "aggregation" AggOp = aggregate(S)
     # b = zeros(eltype(A), size(A, 1))
 
     # Improve candidates
     b = zeros(size(A,1),size(B,2))
-    improve_candidates(A, B, b)
-    T, B = fit_candidates(AggOp, B)
+    @timeit_debug "improve candidates" improve_candidates(A, B, b)
+    @timeit_debug "fit candidates" T, B = fit_candidates(AggOp, B)
 
-    P = smooth(A, T, S, B)
-    R = construct_R(symmetry, P)
-    push!(levels, Level(A, P, R))
+    @timeit_debug "restriction setup" begin
+        P = smooth(A, T, S, B)
+        R = construct_R(symmetry, P)
+    end
 
-    A = R * A * P
+    @timeit_debug "RAP" RAP = R * A * P
 
-    dropzeros!(A)
+    push!(levels, Level(A, P, R))
 
     bsr_flag = true
 
-    A, B, bsr_flag
+    RAP, B, bsr_flag
 end
 construct_R(::HermitianSymmetry, P) = P'
 
-function fit_candidates(AggOp, B; tol=1e-10)
+function fit_candidates(AggOp, B::AbstractVector; tol=1e-10)
+
+    A = adjoint(AggOp)
+    n_fine, n_coarse = size(A)
+    n_col = n_coarse
+
+    R = zeros(eltype(B), n_coarse)
+    Qx = zeros(eltype(B), nnz(A))
+    # copy!(Qx, B)
+    for i = 1:size(Qx, 1)
+        Qx[i] = B[i]
+    end
+    # copy!(A.nzval, B)
+    for i = 1:n_col
+        for j in nzrange(A,i)
+            row = A.rowval[j]
+            A.nzval[j] = B[row]
+        end
+    end
+    k = 1
+    for i = 1:n_col
+        norm_i = norm_col(A, Qx, i)
+        threshold_i = tol * norm_i
+        if norm_i > threshold_i
+            scale = 1 / norm_i
+            R[i] = norm_i
+        else
+            scale = 0
+            R[i] = 0
+        end
+        for j in nzrange(A, i)
+            row = A.rowval[j]
+            # Qx[row] *= scale
+            #@show k
+            # Qx[k] *= scale
+            # k += 1
+            A.nzval[j] *= scale
+        end
+    end
+
+    # SparseMatrixCSC(size(A)..., A.colptr, A.rowval, Qx), R
+    A, R
+end
+
+function fit_candidates(AggOp, B::AbstractMatrix; tol=1e-10)
     A = adjoint(AggOp)
     n_fine, m = ndims(B) == 1 ? (length(B), 1) : size(B)
     n_fine2, n_agg = size(A)
@@ -102,7 +139,7 @@ function fit_candidates(AggOp, B; tol=1e-10)
         rows = A.rowval[A.colptr[agg]:A.colptr[agg+1]-1]
         M = @view B[rows, :]     # size(rows) × m
 
-
+        # TODO the code below can be optimized
         F = qr(M)
         r = min(length(rows), m)
         Qfull = Matrix(F.Q)
@@ -124,3 +161,17 @@ function fit_candidates(AggOp, B; tol=1e-10)
 
     return Qs, R
 end
+
+function norm_col(A, Qx, i)
+    s = zero(eltype(A))
+    for j in nzrange(A, i)
+        if A.rowval[j] > length(Qx)
+            val = 1
+        else
+            val = Qx[A.rowval[j]]
+        end
+        # val = A.nzval[A.rowval[j]]
+        s += val*val
+    end
+    sqrt(s)
+end

src/classical.jl

@@ -9,10 +9,9 @@ function ruge_stuben(_A::Union{TA, Symmetric{Ti, TA}, Hermitian{Ti, TA}},
                 max_coarse = 10,
                 coarse_solver = Pinv, kwargs...) where {Ti,Tv,bs,TA<:SparseMatrixCSC{Ti,Tv}}
 
-
     # fails if near null space `B` is provided
     haskey(kwargs, :B) && kwargs[:B] !== nothing && error("near null space `B` is only supported for smoothed aggregation AMG, not Ruge-Stüben AMG.")
-                
+
     if _A isa Symmetric && Ti <: Real || _A isa Hermitian
         A = _A.data
         symmetric = true
@@ -26,35 +25,37 @@ function ruge_stuben(_A::Union{TA, Symmetric{Ti, TA}, Hermitian{Ti, TA}},
     residual!(w, size(A, 1))
 
     while length(levels) + 1 < max_levels && size(A, 1) > max_coarse
-        A = extend_heirarchy!(levels, strength, CF, A, symmetric)
+        @timeit_debug "extend_hierarchy!" A = extend_hierarchy_rs!(levels, strength, CF, A, symmetric)
         coarse_x!(w, size(A, 1))
         coarse_b!(w, size(A, 1))
         residual!(w, size(A, 1))
     end
 
-    MultiLevel(levels, A, coarse_solver(A), presmoother, postsmoother, w)
+    @timeit_debug "coarse solver setup" cs = coarse_solver(A)
+    return MultiLevel(levels, A, cs, presmoother, postsmoother, w)
 end
 
-function extend_heirarchy!(levels, strength, CF, A::SparseMatrixCSC{Ti,Tv}, symmetric) where {Ti,Tv}
+function extend_hierarchy_rs!(levels, strength, CF, A::SparseMatrixCSC{Ti,Tv}, symmetric) where {Ti,Tv}
     if symmetric
         At = A
     else
         At = adjoint(A)
     end
-    S, T = strength(At)
-    splitting = CF(S)
-    P, R = direct_interpolation(At, T, splitting)
+    @timeit_debug "strength" S, T = strength(At)
+    @timeit_debug "splitting" splitting = CF(S)
+    @timeit_debug "interpolation" P, R = direct_interpolation(At, T, splitting)
+    @timeit_debug "RAP" RAP = R * A * P
     push!(levels, Level(A, P, R))
-    return R * A * P
+    return RAP
 end
 
 function direct_interpolation(At, T, splitting)
     T = typeof(At)(T)
     fill!(T.nzval, eltype(At)(1))
     T .= At .* T
 
-    Pp = rs_direct_interpolation_pass1(T, splitting)
-    Px, Pj, Pp = rs_direct_interpolation_pass2(At, T, splitting, Pp)
+    @timeit_debug "di pass1" Pp = rs_direct_interpolation_pass1(T, splitting)
+    @timeit_debug "di pass2" Px, Pj, Pp = rs_direct_interpolation_pass2(At, T, splitting, Pp)
     R = SparseMatrixCSC(isempty(Pj) ?  0 : maximum(Pj), size(At, 1), Pp, Pj, Px)
     P = R'
 

src/multilevel.jl

@@ -223,7 +223,7 @@ function _solve!(x, ml::MultiLevel, b::AbstractArray{T},
     itr = lvl = 1
     while itr <= maxiter && (!calculate_residual || normres > abstol)
         if length(ml) == 1
-            ml.coarse_solver(x, b)
+            @timeit_debug "Coarse solve" ml.coarse_solver(x, b)
         else
             __solve!(x, ml, cycle, b, lvl)
         end
@@ -259,27 +259,27 @@ end
 
 function __solve!(x, ml, cycle::Cycle, b, lvl)
     A = ml.levels[lvl].A
-    ml.presmoother(A, x, b)
+    @timeit_debug "Presmoother" ml.presmoother(A, x, b)
 
     res = ml.workspace.res_vecs[lvl]
-    mul!(res, A, x)
+    @timeit_debug "Residual eval" mul!(res, A, x)
     reshape(res, size(b)) .= b .- reshape(res, size(b))
 
     coarse_b = ml.workspace.coarse_bs[lvl]
-    mul!(coarse_b, ml.levels[lvl].R, res)
+    @timeit_debug "Restriction" mul!(coarse_b, ml.levels[lvl].R, res)
 
     coarse_x = ml.workspace.coarse_xs[lvl]
     coarse_x .= 0
     if lvl == length(ml.levels)
-        ml.coarse_solver(coarse_x, coarse_b)
+        @timeit_debug "Coarse solve" ml.coarse_solver(coarse_x, coarse_b)
     else
         coarse_x = __solve_next!(coarse_x, ml, cycle, coarse_b, lvl + 1)
     end
 
-    mul!(res, ml.levels[lvl].P, coarse_x)
+    @timeit_debug "Prolongation" mul!(res, ml.levels[lvl].P, coarse_x)
     x .+= res
 
-    ml.postsmoother(A, x, b)
+    @timeit_debug "Postsmoother" ml.postsmoother(A, x, b)
 
     x
 end

test/runtests.jl

@@ -333,9 +333,6 @@ ml = ruge_stuben(X)
 b = rand(27_000)
 @test AlgebraicMultigrid._solve(ml, b, reltol = 1e-10) ≈ X \ b rtol = 1e-10
 
-
-
-
 # LinearSolve precs interface
 @testset "LinearSolvePrecs" begin
 
@@ -350,7 +347,6 @@ for sz in [ (10,10), (20,20), (50,50)]
 
     strategy = KrylovJL_CG(precs = SmoothedAggregationPreconBuilder())
     @test solve(prob, strategy, atol=1.0e-14) ≈ u0 rtol = 1.0e-8
-
 end
 
 end