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

virtual_size called too often #448

Closed radha-patel closed 6 months ago

radha-patel commented 6 months ago

size(A) is called three times in the below kernel when it only needs to be called once!

julia> A = rand(10, 10); x = rand(10); y = zeros(10)

julia> @finch_code begin 
                  for j=_, i=_
                      y[i] += A[i, j] * x[i]
                  end
              end
quote
    y = (ex.bodies[1]).body.body.lhs.tns.bind
    A = ((ex.bodies[1]).body.body.rhs.args[1]).tns.bind
    x = ((ex.bodies[1]).body.body.rhs.args[2]).tns.bind
    sugar_1 = size(y)
    y_mode1_stop = sugar_1[1]
    sugar_2 = size(A)
    A_mode1_stop = sugar_2[1]
    A_mode2_stop = sugar_2[2]
    A_mode1_stop == y_mode1_stop || throw(DimensionMismatch("mismatched dimension limits ($(A_mode1_stop) != $(y_mode1_stop))"))
    sugar_3 = size(x)
    x_mode1_stop = sugar_3[1]
    A_mode1_stop == x_mode1_stop || throw(DimensionMismatch("mismatched dimension limits ($(A_mode1_stop) != $(x_mode1_stop))"))
    result = nothing
    for j_3 = 1:A_mode2_stop
        sugar_4 = size(A)
        A_mode1_stop = sugar_4[1]
        A_mode2_stop = sugar_4[2]
        for i_5 = 1:A_mode1_stop
            sugar_7 = size(A)
            A_mode1_stop = sugar_7[1]
            A_mode2_stop = sugar_7[2]
            val = x[i_5]
            val_2 = A[i_5, j_3]
            y[i_5] += val_2 * val
        end
    end
    result = ()
    result
end
radha-patel commented 6 months ago

Seems that the issue comes from the fact that virtual_size in AbstractArrays pushes the size to ctx.code.preamble every time it is called. (And for this case in particular, virtual_size is called in popdim within an assert statement causing the repeated size measurements).

A potential solution is to cache the size of the array when we virtualize it.