Comparison row / column

jacobi.jl

import Base: start, next, done

using BenchmarkTools

mutable struct RowJacobi{matT,vecT}
    A::matT
    x::vecT
    next::vecT
    b::vecT
    maxiter::Int
end

mutable struct ColumnJacobi{matT,vecT}
    A::matT
    x::vecT
    next::vecT
    b::vecT
    maxiter::Int
end

start(::Union{RowJacobi,ColumnJacobi}) = 1
done(it::Union{RowJacobi,ColumnJacobi}, iteration::Int) = iteration > it.maxiter

function next(j::RowJacobi, iteration::Int)
    n = size(j.A, 1)
    
    for row = 1 : n
        σ = 0.0
        for col = 1 : row - 1
            @inbounds σ += j.A[row, col] * j.x[col]
        end

        for col = row + 1 : n
            @inbounds σ += j.A[row, col] * j.x[col]
        end

        @inbounds j.next[row] = (j.b[row] - σ) / j.A[row, row]
    end

    j.x, j.next = j.next, j.x

    nothing, iteration + 1
end

function next(j::ColumnJacobi, iteration::Int)
    n = size(j.A, 1)

    copy!(j.next, j.b)
    
    for col = 1 : n
        @simd for row = 1 : col - 1
            @inbounds j.next[row] -= j.A[row, col] * j.x[col]
        end

        @simd for row = col + 1 : n
            @inbounds j.next[row] -= j.A[row, col] * j.x[col]
        end
    end

    for col = 1 : n
        @inbounds j.x[col] = j.next[col] / j.A[col, col]
    end

    nothing, iteration + 1
end

function row_jacobi!(x::AbstractVector, A::DenseMatrix, b::AbstractVector; maxiter = 10)
    iterable = RowJacobi(A, x, similar(x), b, maxiter)

    for _ = iterable end

    x
end

function column_jacobi!(x::AbstractVector, A::DenseMatrix, b::AbstractVector; maxiter = 10)
    iterable = ColumnJacobi(A, x, similar(x), b, maxiter)

    for _ = iterable end

    x
end

function bench()
    ns = round.(Int, logspace(1, log10(500), 10))
    by_col = []
    by_row = []

    for n = ns
        println(n)
        A = rand(n, n) + n * I
        b = A * ones(n)

        # Make sure they compute the same thing.
        x_c = column_jacobi!(zeros(n), A, b; maxiter = 20)
        x_r = row_jacobi!(zeros(n), A, b; maxiter = 20)

        @assert x_c ≈ x_r

        push!(by_col, @benchmark column_jacobi!(x0, $A, $b, maxiter = 20) setup = (x0 = zeros($n)))
        push!(by_row, @benchmark row_jacobi!(x0, $A, $b, maxiter = 20) setup = (x0 = zeros($n)))
    end

    plot_stuff(ns, by_col, by_row), ns, by_col, by_row
end

function plot_stuff(ns, by_col, by_row)
    take_time(x) = x.time

    col_time = take_time.(minimum.(by_col))
    row_time = take_time.(minimum.(by_row))

    plot(ns, row_time ./ col_time, xscale = :log10, label = "By column", marker = :+, 
         xlabel = "Matrix width", ylabel = "Speedup by row / by col", legend = :topleft)
end