non-allocating cumsum for SVector

[jkrumbiegel]

Hi, I noticed while optimizing some code that the only allocating bit that was left was a cumsum over an SVector. Because I thought it was an interesting challenge, I tried to make my first @generated function to fix this. Because I feel like it's really bad style (didn't know how to do it with quoting and interpolating), I'm posting it here first, before making a PR..

So this is the default behavior:

>>> svec = SVector(1, 2, 3, 4, 5)

>>> @btime cumsum($svec)

13.754 ns (1 allocation: 48 bytes)
5-element MArray{Tuple{5},Int64,1,5} with indices SOneTo(5):
  1
  3
  6
 10
 15

Now my function:

@generated function mycumsum(svec::SVector{N, T}) where {N, T}

    args = Any[]

    i = 1
    push!(args, :(v1 = svec[1]))

    for i in 2:N
        push!(args,
            Expr(
                Symbol("="),
                Symbol("v$i"),
                Expr(
                    Symbol("call"),
                    Symbol("+"),
                    Expr(
                        Symbol("ref"),
                        Symbol("svec"),
                        i
                    ),
                    Symbol("v$(i-1)")
                )
            )
        )
    end

    variables = [Symbol("v$i") for i in 1:N]
    push!(args, Expr(Symbol("call"), :SVector, variables...))

    exp = Expr(Symbol("block"), args...)

end

This is the expression that this function generates for svec:

exp = quote
    v1 = svec[1]
    v2 = svec[2] + v1
    v3 = svec[3] + v2
    v4 = svec[4] + v3
    v5 = svec[5] + v4
    SVector(v1, v2, v3, v4, v5)
end

And this is the result when using it:

@btime mycumsum($svec)
  0.031 ns (0 allocations: 0 bytes)
5-element SArray{Tuple{5},Int64,1,5} with indices SOneTo(5):
  1
  3
  6
 10
 15

So if I'm not mistaken, this is 444 times faster than the default version?

[KristofferC]

So if I'm not mistaken, this is 444 times faster than the default version?

A result like 0.031 ns means that no computation took place and the whole thing got calculated at compile time and optimized away. WIth some napkin math, say a CPU runs at 1 GHz and does one operation per Hz, that would be 1 ns per operations so 0.031ns for the sum is clearly unreasonable.

[jkrumbiegel]

Haha that makes sense, so how do I measure the actual effect?

[KristofferC]

One workaround that currently seems to work is writing it as

julia> @btime mycumsum($(Ref(svec))[])
  1.853 ns (0 allocations: 0 bytes)

[jkrumbiegel]

Hey that's still 8.36 times faster :) And no allocation

[jkrumbiegel]

Do you maybe have a comment about what I need to do to turn this into PR material?

[KristofferC]

Could write it a bit shorter like

@generated function mycumsum(svec::SVector{N, T}) where {N, T}
    return quote
        accum = zero(T)
        SVector{N, T}( $([:(accum += svec[$(i-1)]) for i in 2:N+1]...))
    end
end

perhaps.

[jkrumbiegel]

that is concise!

[c42f]

cumsum is implemented in #702 so that will also fix this.

I feel like #702 is asking a lot of the compiler and is a very different style from the rest of the package so I haven't felt fully confident to merge it yet. But those worries are probably unjustified, I just need to take a closer look.

[jkrumbiegel]

oh I see, I did search issues for cumsum but it wasn't in the title :)

[c42f]

No problem, it can be hard to find these things.

Anyway, now we get to close an extra issue when #702 merges 😀