the_bench.jl

the_bench.jl

import Base: start, next, done

using Base: BitInteger
using BenchmarkTools

struct OpenRange{T}
    start::T
    stop::T
end

→(start::T, stop::T) where {T<:BitInteger} = OpenRange(start, stop)

# Iterators
start(r::OpenRange{<:BitInteger}) = r.start
next(::OpenRange{T}, state) where {T<:BitInteger} = state, state + T(1)
done(r::OpenRange{<:BitInteger}, state) = state == r.stop

"""
Simplified A_mul_B!
"""
function mul_impl_1!(y::StridedVector, A::SparseMatrixCSC{Tv,Ti}, x::StridedVector) where {Tv,Ti}
    fill!(y, zero(eltype(y)))

    @inbounds for i = 1 : A.n
        xval = x[i]
        for j = (A.colptr[i] → A.colptr[i+1])
            y[A.rowval[j]] += A.nzval[j] * xval
        end
    end

    y
end

"""
Simplified A_mul_B!
"""
function mul_impl_2!(y::StridedVector, A::SparseMatrixCSC, x::StridedVector)
    fill!(y, zero(eltype(y)))
    
    @inbounds for i = 1:A.n
        xval = x[i]
        for j = A.colptr[i]:(A.colptr[i+1] - 1)
            y[A.rowval[j]] += A.nzval[j] * xval
        end
    end
    y
end

"""
Benchmark a mv-product with a matrix of order n and k super- and sub-diagonals.
"""
function bench(n = 100_000, k = 2)
    B = spdiagm([fill(rand(Float32), n - abs(i)) for i = -k:k], -k:k)

    for Ti in (Int32, UInt32, Int64, UInt64)
        A = convert(SparseMatrixCSC{Float32,Ti}, B)
        x = rand(Float32, n)
        y1 = similar(x)
        y2 = similar(x)

        @assert mul_impl_1!(y1, A, x) == mul_impl_2!(y2, A, x)

        a = @benchmark mul_impl_1!($y1, $A, $x)
        b = @benchmark mul_impl_2!($y2, $A, $x)

        println(Ti, ":\t", a, '\t', b)
    end
end

bench()