willow-ahrens / Finch.jl

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

Usage with index set construction #318

Closed ericphanson closed 6 months ago

ericphanson commented 9 months ago

Is it possible to write a Finch kernel for

function subXTv!(XTv::AbstractVector, x::AbstractMatrix, v::AbstractVector, I::Vector{Int}, sizeI::Int)
    @views begin
        xx = x[:, I[1:sizeI]]
        vv = v[1:sizeI]
    end
    return mul!(XTv, xx, vv)
end

where x is sparse, and v is dense?

I can write a non-Finch SparseMatrixCSC kernel with

cols = 1:sizeI
vals = nonzeros(x)
rows = rowvals(x)
XTv .= 0
for col in cols
    b_value = @inbounds v[col]
    nzr = nzrange(x, I[col])
    @fastmath @inbounds @simd for i in nzr
        row = rows[i]
        val = vals[i]
        XTv[row] += val * b_value
    end
end

though it needs a bunch of extra workspace to parallelize (I need to have n_tasks copies of XTv, accumulate into each separately over a subset of the columns, then add them together).

wraith1995 commented 9 months ago

@ericphanson This is very possible. You can certainly read out indirect views of x and a portion of v as you suggest because Finch supports indirect access so you can write x[k, I[i]] (or if statements (e.g. if j == I[i], if you prefer)) and because Finch supports loops over a range j = 1:sizeI. With these ingredients, you should be able to write a version of your kernel. One complication might be the format of XTv. If this is dense, it is easy. If it is sparse. then you might need to write more complex code.

We are in the process of adding parallelism, but the optimization of using workspaces to parallelize the kernel might be possible already.

ericphanson commented 9 months ago

Ah great! This does indeed work:

function subXTv!(XTv::AbstractVector, x::AbstractFiber, v::AbstractVector, I::Vector{Int}, sizeI::Int)
    vv = view(v, 1:sizeI)
    II = view(I, 1:sizeI)
    @finch begin
        XTv .= 0
        for i=_, j=_
         XTv[i] += x[i,II[j]]*vv[j]
        end
    end
    return nothing
end

I had been trying without the views and couldn't get the indices to work:

function subXTv!(XTv::AbstractVector, x::AbstractFiber, v::AbstractVector, I::Vector{Int}, sizeI::Int)
    @finch begin
        XTv .= 0
        for i=_, j=1:sizeI
         XTv[i] += x[i,I[j]]*v[j]
        end
    end
    return nothing
end

which yields

  Got exception outside of a @test
  DimensionMismatch: mismatched dimension limits (200 != 100)
  Stacktrace:
    [1] macro expansion
      @ ~/.julia/packages/Finch/R62UH/src/environment.jl:109 [inlined]
    [2] macro expansion
      @ ~/.julia/packages/Finch/R62UH/src/execute.jl:67 [inlined]
    [3] macro expansion
      @ ~/.julia/packages/Finch/R62UH/src/environment.jl:73 [inlined]
    [4] macro expansion
      @ ~/.julia/packages/Finch/R62UH/src/util.jl:76 [inlined]
    [5] execute(ex::Finch.FinchNotation.BlockInstance{Tuple{Finch.FinchNotation.DeclareInstance{Finch.FinchNotation.TagInstance{Finch.FinchNotation.VariableInstance{:XTv}, Vector{Float64}}, Finch.FinchNotation.LiteralInstance{0}}, Finch.FinchNotation.LoopInstance{Finch.FinchNotation.IndexInstance{:i}, Finch.FinchNotation.Dimensionless, Finch.FinchNotation.LoopInstance{Finch.FinchNotation.IndexInstance{:j}, Finch.FinchNotation.CallInstance{Finch.FinchNotation.TagInstance{Finch.FinchNotation.VariableInstance{:extent}, Finch.FinchNotation.LiteralInstance{Finch.FinchNotation.extent}}, Tuple{Finch.FinchNotation.LiteralInstance{1}, Finch.FinchNotation.TagInstance{Finch.FinchNotation.VariableInstance{:sizeI}, Int64}}}, Finch.FinchNotation.AssignInstance{Finch.FinchNotation.AccessInstance{Finch.FinchNotation.TagInstance{Finch.FinchNotation.VariableInstance{:XTv}, Vector{Float64}}, Finch.FinchNotation.LiteralInstance{Finch.FinchNotation.Updater()}, Tuple{Finch.FinchNotation.TagInstance{Finch.FinchNotation.VariableInstance{:i}, Finch.FinchNotation.IndexInstance{:i}}}}, Finch.FinchNotation.TagInstance{Finch.FinchNotation.VariableInstance{:+}, Finch.FinchNotation.LiteralInstance{+}}, Finch.FinchNotation.CallInstance{Finch.FinchNotation.TagInstance{Finch.FinchNotation.VariableInstance{:*}, Finch.FinchNotation.LiteralInstance{*}}, Tuple{Finch.FinchNotation.AccessInstance{Finch.FinchNotation.TagInstance{Finch.FinchNotation.VariableInstance{:x}, Fiber{DenseLevel{Int64, SparseListLevel{Int64, Vector{Int64}, Vector{Int64}, ElementLevel{0.0, Float64, Int64, Vector{Float64}}}}}}, Finch.FinchNotation.LiteralInstance{Finch.FinchNotation.Reader()}, Tuple{Finch.FinchNotation.TagInstance{Finch.FinchNotation.VariableInstance{:i}, Finch.FinchNotation.IndexInstance{:i}}, Finch.FinchNotation.AccessInstance{Finch.FinchNotation.TagInstance{Finch.FinchNotation.VariableInstance{:I}, Vector{Int64}}, Finch.FinchNotation.LiteralInstance{Finch.FinchNotation.Reader()}, Tuple{Finch.FinchNotation.TagInstance{Finch.FinchNotation.VariableInstance{:j}, Finch.FinchNotation.IndexInstance{:j}}}}}}, Finch.FinchNotation.AccessInstance{Finch.FinchNotation.TagInstance{Finch.FinchNotation.VariableInstance{:v}, Vector{Float64}}, Finch.FinchNotation.LiteralInstance{Finch.FinchNotation.Reader()}, Tuple{Finch.FinchNotation.TagInstance{Finch.FinchNotation.VariableInstance{:j}, Finch.FinchNotation.IndexInstance{:j}}}}}}}}}}}, opts::NamedTuple{(), Tuple{}})
      @ Finch ~/.julia/packages/Finch/R62UH/src/util.jl:68
    [6] macro expansion
      @ ~/.julia/packages/Finch/R62UH/src/execute.jl:173 [inlined]
    [7] subXTv!(XTv::Vector{Float64}, x::Fiber{DenseLevel{Int64, SparseListLevel{Int64, Vector{Int64}, Vector{Int64}, ElementLevel{0.0, Float64, Int64, Vector{Float64}}}}}, v::Vector{Float64}, I::Vector{Int64}, sizeI::Int64)

in my test code


@testset "subXTv!" begin
    n, m = 100, 200
    I = collect(1:m)
    sizeI = m ÷ 2
    Xv = Vector{Float64}(undef, n)
    A = sprand(n, m, 0.01)
    v = ones(m)
    for x in (Matrix(A), sparse(A), fiber(A))
        subXTv!(Xv, x, v, I, sizeI)
        @test Xv ≈ A[:, I[1:sizeI]] * v[1:sizeI]
    end
end

However, since v and I are dense, the views seem harmless enough, although I haven't benchmarked yet.

wraith1995 commented 9 months ago

@ericphanson You shouldn't need the views - there could be a bug in the dimensioinalization computation. In fact, I think there is based on this. I'll mark this as bug and discuss with others.

willow-ahrens commented 6 months ago

No, there's no bug, the dimensions of all the uses of j need to match unless you exempt them from dimensionalization with ~. I'll link the appropriate issue about docs surrounding dimensionalization and close this once we have dimension docs. https://github.com/willow-ahrens/Finch.jl/issues/335