Possible to do non-allocating Setfield operations on OffsetArray(::StaticVector)?

General Usage
1

I have some code which initializes an array with a loop in a way where each element depends on the last element. I’m trying to make this work with StaticVectors so it can be non-allocating. I’ve found that I can use Setfield.@set!(since I need to depend on the previous elements) and that part seems to work fine, e.g.

function foo(x₁)
    x = @SVector(zeros(Int,5))
    @set! x[1] = x₁
    for i=2:5
        @set! x[i] = 2*x[i-1]+1
    end
    x
end
@code_llvm foo(0)

generates what to my naive eye looks like optimal code (which, in a lot of ways, is already pretty amazing to me given what Setfield does; Julia 1.6, btw).

For code clarity in my real case, I also want to have this be an OffsetArray though. Something like:

function offset_foo(x₀)
    x = OffsetVector(@SVector(zeros(Int,5)), 0:4)
    @set! x[0] = x₀
    for i=1:4
        @set! x[i] = 2*x[i-1]+1
    end
    x
end

This allocates unfortunately. I think it may related to bounds checking in OffsetArray.get/setindex, but I can’t figure out how to turn it off. Pehaps unsurprisingly, neither of

@set! @inbounds x[0] = x₀
@inbounds @set! x[0] = x₀

seem to work (one is an error, the other doesn’t do anything).

Any ideas how I can turn off bounds checking here? Or if there’s other better approaches to solving this problem?

EDIT: Actually sorry I don’t think its the boundschecking, its just that @set! is turning the OffsetArray(::StaticVector) into a OffsetArray(::Vector). Still welcome any suggestions for a way around this though.

2

This behavior might get closer to what you want after OffsetArrays.jl/pull/213 is merged. After this,

julia> offset_foo(1)
5-element OffsetArray(::MArray{Tuple{5},Int64,1,5}, 0:4) with eltype Int64 with indices 0:4:
  1
  3
  7
 15
 31

The return type is an offset MArray, so it is static. It’s not allocation free though, but allocations do come down.

Before:

julia> @btime offset_foo(1);
  337.063 ns (5 allocations: 640 bytes)

After:

julia> @btime offset_foo(1);
  134.921 ns (5 allocations: 240 bytes)
2 Likes