Matvecs with matrices in PLINK v1 BEM format

m012.jl

import Base: convert, rand, size, getindex, A_mul_B!

immutable PLINK1Matrix <: AbstractMatrix{UInt8}
    m :: Int
    n :: Int
    data :: Array{UInt8}
end

size(M::PLINK1Matrix, i::Int) = i==1 ? M.m : i==2 ? M.n : error()
size(M::PLINK1Matrix) = (M.m, M.n)

@inline function getindex(M::PLINK1Matrix, i::Int, j::Int)
    offset = (i-1)*M.n+(j-1) #Assume row major
    byteoffset = (offset >> 2) + 4 #Skip the magic bytes
    crumboffset = offset & 0b00000011
    @inbounds rawbyte = M.data[byteoffset]
    rawcrumb = (rawbyte >> 6-2crumboffset) & 0b11
    ###############################
    # Raw encoding # Return value #
    ###############################
    #     00       #       0      #
    #     01       #       0 (NA) # XXX should be missing
    #     10       #       1      #
    #     11       #       2      #
    ###############################
    ifelse(rawcrumb==0, rawcrumb, rawcrumb-0x01)
end

function rand(::Type{PLINK1Matrix}, m::Int, n::Int)
    numbytes = m*n
    nrand = ceil(Int, numbytes+3)
    data = rand(UInt8, nrand)
    #Add magic bytes
    data[1] = 0x6c
    data[2] = 0x1b
    data[3] = 0x01
    PLINK1Matrix(m, n, data)
end

function convert(::Type{Matrix{Float64}}, M::PLINK1Matrix)
    A = zeros(M.m, M.n)
    @inbounds for i=1:M.m,j=1:M.n
        A[i,j] = M[i,j]
    end
    A
end

function A_mul_B!{T}(y::AbstractVector{T},
	M::PLINK1Matrix, b::AbstractVector{T})

    y[:] = zero(T)
    @fastmath @inbounds for i=1:M.m
  	δy = zero(T)
        @simd for j=1:4:M.n
            x = M[i,j]
  	    z = b[j]
            x2 = M[i,j+1]
  	    z2 = b[j+1]
            x3 = M[i,j+2]
  	    z3 = b[j+2]
            x4 = M[i,j+3]
  	    z4 = b[j+3]
  	    δy += x*z + x2*z2 + x3*z3 + x4*z4
        end
	y[i] += δy
    end
    y
end

m=8000
n=4000
M=rand(PLINK1Matrix,m,n)
b=rand(n)

#Trigger compile
rand(PLINK1Matrix,1,1)*rand(1)
@time M*b

#Compare against dense BLAS
Mf=convert(Matrix{Float64}, M)
@time Mf*b

y=zeros(m)
@time Base.LinAlg.generic_matvecmul!(y, 'N', Mf, b)

y=zeros(m)
@time Base.LinAlg.generic_matvecmul!(y, 'N', M, b)