libcsb.jl

@doc """
    $(TYPEDEF)

Matrix type for storing sparse matrices in the Compressed Sparse
Blocks format. The standard way of constructing `SparseMatrixCSB` is
to pass a `SparseMatrixCSC` object, see the constructors.

"""
mutable struct SparseMatrixCSB{Tv,Ti} <: SparseArrays.AbstractSparseMatrix{Tv,Ti}
    m::Ti
    n::Ti
    nz::Ti
    ptr::Ptr{Cvoid}
    function SparseMatrixCSB{Tv,Ti}(m::Ti, n::Ti, nz::Ti, ptr::Ptr{Cvoid}) where {Tv,Ti}
        A = new(m, n, nz, ptr)
        function f(X)
            @async deallocateCSB!(X)
        end
        return finalizer(f, A)
    end
end

CSBTYPES = Union{
    SparseMatrixCSB{Tv,Ti},
    LinearAlgebra.Adjoint{Tv,SparseMatrixCSB{Tv,Ti}},
    Transpose{Tv,SparseMatrixCSB{Tv,Ti}},
} where {Tv,Ti}



@doc """

$(TYPEDSIGNATURES)

Convert a `SparseMatrixCSC` matrix `A` into a `SparseMatrixCSB` matrix.

# Optional arguments

- `beta`: The size of each block in base-2 logarithm; if 0 the package
  decides the block size internally.

"""
SparseMatrixCSB(A::SparseMatrixCSC{Tv,Ti}, beta::Integer=0) where {Tv,Ti<:Integer} =
    prepareCSB(A, beta)


@doc """
$(TYPEDSIGNATURES)

Query the number of `Cilk` workers.
"""
getWorkers() = Int64(@ccall libcsb.getWorkers()::Cint)

@doc """
$(TYPEDSIGNATURES)

Set the number of `Cilk` workers to `np`.
"""
setWorkers(np::Int64) = @ccall libcsb.setWorkers(np::Int64)::Cvoid

function prepareCSB(A::SparseMatrixCSC{Tv,Ti}, beta=0) where {Tv,Ti<:Integer}
    size(A, 1) == 0 && error("Matrix too small.")
    size(A, 2) == 0 && error("Matrix too small.")
    nextpow(2, size(A, 1)) <= SLACKNESS * getWorkers() && error("Matrix too small.")
    nextpow(2, size(A, 2)) <= SLACKNESS * getWorkers() && error("Matrix too small.")
    mktemp() do path, io
        redirect_stdout(io) do
            SparseMatrixCSB{Tv,Ti}(
                Ti(size(A, 1)),
                Ti(size(A, 2)),
                Ti(nnz(A)),
                _prepareCSB(
                    A.nzval,
                    A.rowval .- one(Ti),
                    A.colptr .- one(Ti),
                    Ti(nnz(A)),
                    Ti(size(A, 1)),
                    Ti(size(A, 2)),
                    getWorkers(),
                    beta,
                ),
            )
        end
    end
end

for (Tv, Tvname) in ((Cdouble, "double"),)
    for (Ti, Tiname) in ((Cuint, "uint32"), (Cintmax_t, "int64"))
        @eval @inline function _prepareCSB(
            val::Vector{$Tv},
            row::Vector{$Ti},
            colptr::Vector{$Ti},
            nzmax::$Ti,
            m::$Ti,
            n::$Ti,
            workers::Integer,
            forcelogbeta::Integer,
        )
            return ccall(
                ($("prepareCSB_" * Tvname * "_" * Tiname), libcsb),
                Ptr{Cvoid},
                (Ptr{$Tv}, Ptr{$Ti}, Ptr{$Ti}, $Ti, $Ti, $Ti, Cint, Cint),
                val,
                row,
                colptr,
                nzmax,
                m,
                n,
                workers,
                forcelogbeta,
            )

        end

        @eval @inline function _gespmvCSB!(
            y::Vector{$Tv}, A::SparseMatrixCSB{$Tv,$Ti}, x::Vector{$Tv}
        )
            return ccall(
                ($("gespmv_" * Tvname * "_" * Tiname), libcsb),
                Ptr{Cvoid},
                (Ptr{Cvoid}, Ptr{$Tv}, Ptr{$Tv}),
                A.ptr,
                x,
                y,
            )

        end

        @eval @inline function _gespmvtCSB!(
            y::Vector{$Tv},
            A::LinearAlgebra.Adjoint{$Tv,SparseMatrixCSB{$Tv,$Ti}},
            x::Vector{$Tv},
        )
            return ccall(
                ($("gespmvt_" * Tvname * "_" * Tiname), libcsb),
                Ptr{Cvoid},
                (Ptr{Cvoid}, Ptr{$Tv}, Ptr{$Tv}),
                A.parent.ptr,
                x,
                y,
            )

        end
        @eval @inline function _gespmvtCSB!(
            y::Vector{$Tv},
            A::LinearAlgebra.Transpose{$Tv,SparseMatrixCSB{$Tv,$Ti}},
            x::Vector{$Tv},
        )
            return ccall(
                ($("gespmvt_" * Tvname * "_" * Tiname), libcsb),
                Ptr{Cvoid},
                (Ptr{Cvoid}, Ptr{$Tv}, Ptr{$Tv}),
                A.parent.ptr,
                x,
                y,
            )

        end

        for DIM in 2:32
            @eval @inline function _gespmmCSB!(
                y::Matrix{$Tv}, A::SparseMatrixCSB{$Tv,$Ti}, x::Matrix{$Tv}, ::Val{$DIM}
            )
                return ccall(
                    (
                        $("gespmm_" * Tvname * "_" * Tiname * "_" * string(DIM) * "_rhs"),
                        libcsb,
                    ),
                    Ptr{Cvoid},
                    (Ptr{Cvoid}, Ptr{$Tv}, Ptr{$Tv}, Cint, Cint),
                    A.ptr,
                    x,
                    y,
                    size(y, 1),
                    size(x, 1),
                )

            end
            @eval @inline function _gespmmtCSB!(
                y::Matrix{$Tv},
                A::LinearAlgebra.Adjoint{$Tv,SparseMatrixCSB{$Tv,$Ti}},
                x::Matrix{$Tv},
                ::Val{$DIM},
            )
                return ccall(
                    (
                        $("gespmmt_" * Tvname * "_" * Tiname * "_" * string(DIM) * "_rhs"),
                        libcsb,
                    ),
                    Ptr{Cvoid},
                    (Ptr{Cvoid}, Ptr{$Tv}, Ptr{$Tv}, Cint, Cint),
                    A.parent.ptr,
                    x,
                    y,
                    size(y, 1),
                    size(x, 1),
                )

            end
            @eval @inline function _gespmmtCSB!(
                y::Matrix{$Tv},
                A::LinearAlgebra.Transpose{$Tv,SparseMatrixCSB{$Tv,$Ti}},
                x::Matrix{$Tv},
                ::Val{$DIM},
            )
                return ccall(
                    (
                        $("gespmmt_" * Tvname * "_" * Tiname * "_" * string(DIM) * "_rhs"),
                        libcsb,
                    ),
                    Ptr{Cvoid},
                    (Ptr{Cvoid}, Ptr{$Tv}, Ptr{$Tv}, Cint, Cint),
                    A.parent.ptr,
                    x,
                    y,
                    size(y, 1),
                    size(x, 1),
                )

            end
        end

        @eval @inline function _deallocate!(A::SparseMatrixCSB{$Tv,$Ti})
            return ccall(
                ($("deallocate_" * Tvname * "_" * Tiname), libcsb),
                Ptr{Cvoid},
                (Ptr{Cvoid},),
                A.ptr,
            )

        end

    end
end

function deallocateCSB!(A::SparseMatrixCSB)
    A.ptr == C_NULL && error("Invalid CSB object")
    _deallocate!(A)
    A.ptr = C_NULL
    A.m = A.n = A.nz = 0
    return nothing
end



# overrides

function mul!(y::AbstractVecOrMat, A::SparseMatrixCSB, x::AbstractVector)

    @assert size(y, 1) == size(A, 1)
    @assert size(x, 1) == size(A, 2)

    fill!(y, 0)
    _gespmvCSB!(y, A, x)
    return y

end

function mul!(y::AbstractVecOrMat, A::SparseMatrixCSB, x::AbstractMatrix)

    @assert size(y, 1) == size(A, 1)
    @assert size(x, 1) == size(A, 2)
    @assert size(y, 2) == size(x, 2)

    fill!(y, 0)
    if size(x, 2) == 1
        _gespmvCSB!(vec(y), A, vec(x))
    elseif size(x, 2) > 32
        throw(
            DimensionMismatch(
                "This CSB wrapper has been compiled to support up to 32 columns at a time."
            ),
        )
    else
        _gespmmCSB!(y, A, x, Val(size(x, 2)))
    end
    return y

end

# transpose

function mul!(
    y::AbstractVecOrMat,
    A::Union{Adjoint{Tv,SparseMatrixCSB{Tv,Ti}},Transpose{Tv,SparseMatrixCSB{Tv,Ti}}},
    x::AbstractVector,
) where {Tv,Ti}

    @assert size(y, 1) == size(A, 1)
    @assert size(x, 1) == size(A, 2)

    fill!(y, 0)
    _gespmvtCSB!(y, A, x)
    return y

end

function mul!(
    y::AbstractVecOrMat,
    A::Union{Adjoint{Tv,SparseMatrixCSB{Tv,Ti}},Transpose{Tv,SparseMatrixCSB{Tv,Ti}}},
    x::AbstractMatrix,
) where {Tv,Ti}

    @assert size(y, 1) == size(A, 1)
    @assert size(x, 1) == size(A, 2)
    @assert size(y, 2) == size(x, 2)

    fill!(y, 0)
    if size(x, 2) == 1
        _gespmvtCSB!(vec(y), A, vec(x))
    elseif size(x, 2) > 32
        throw(
            DimensionMismatch(
                "This CSB wrapper has been compiled to support up to 32 columns at a time."
            ),
        )
    else
        _gespmmtCSB!(y, A, x, Val(size(x, 2)))
    end
    return y

end

size(A::SparseMatrixCSB) = (A.m, A.n)
nnz(A::SparseMatrixCSB) = A.nz
nnz(A::LinearAlgebra.Adjoint{Tv,SparseMatrixCSB{Tv,Ti}}) where {Tv,Ti} = A.parent.nz
nnz(A::LinearAlgebra.Transpose{Tv,SparseMatrixCSB{Tv,Ti}}) where {Tv,Ti} = A.parent.nz

function Base.print_matrix(io::IO, S::CSBTYPES)

end

function array_summary(io::IO, S::CSBTYPES, dims::Tuple{Vararg{Base.OneTo}})
    xnnz = nnz(S)
    m, n = size(S)
    print(
        io,
        m,
        "×",
        n,
        " ",
        typeof(S),
        " with ",
        xnnz,
        " stored ",
        xnnz == 1 ? "entry" : "entries",
    )
    return nothing
end

function print_array(io::IO, A::CSBTYPES)

end