willow-ahrens / Finch.jl

Sparse tensors in Julia and more! Datastructure-driven array programing language.
http://willowahrens.io/Finch.jl/
MIT License
151 stars 12 forks source link

Improve the Sparse level's sort procedure #600

Closed kylebd99 closed 3 weeks ago

kylebd99 commented 3 weeks ago

This uses a two-step sort where we first use tbl.ptr to bucket the idxs & values. Then, we sort each bucket individually. Some initial testing showed a 2x matrix transpose improvement:

using Finch
using BenchmarkTools

n = 100000
nnz = 1000000
A = Tensor(Dense(SparseList(Element(0.0))), fsprand(n, n, nnz))
C = Tensor(Dense(Sparse(Element(0.0))))

eval(@finch_kernel mode=:fast function tp(A, C)
    C .= 0
    for i=_
        for j=_
            C[i,j] = A[j,i]
        end
    end
end)

hash_times = []

@benchmark tp($A, $C)
willow-ahrens commented 3 weeks ago

this looks good! We can merge this now, and I'll see if I can tweak it later.