Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Sparse matrix - vector multiplication with threads.
# Gives roughly a 2x speedup using JULIA_NUM_THREADS=2
# requires 0.7, due to closure type inference bug in 0.6
using Random, SparseArrays
srand(1); X = sprand(20_000,100_000,0.1); v = ones(100_000);
function mul(X,v)
m,n = size(X)
U = zeros(size(X,1), Threads.nthreads())
@assert n == length(v)
Threads.@threads for col = 1:n
thread = Threads.threadid()
@inbounds a = v[col]
@inbounds for i = X.colptr[col]:(X.colptr[col + 1] - 1)
row = X.rowval[i]
@inbounds U[row, thread] = muladd(X.nzval[i], a, U[row, thread])
end
end
sum(U,2)
end
@time X*v;
@time X*v;
@time X*v;
@time mul(X,v);
@time mul(X,v);
@time mul(X,v);
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment