Weird allocation issue for fill! with ArrayPartition in Julia < 1.6

General Usage Performance
,
1

Hello Julians,
I have an issue with allocations that I encountered while running my code on a cluster which uses Julia 1.5.4. Basically, once my partition array gets reasonably large, I encounter many allocations when calling the fill! function. (It also occurs if I iterate over this array explicitly).
Consider the following MWE:

using RecursiveArrayTools, BenchmarkTools

function test1()
    PartArray = ArrayPartition(
        zeros(2),
        zeros(200,200,30),
    )
    @btime fill!($PartArray,0.)
end
    
function test2()
    PartArray = ArrayPartition(
        zeros(2),
        zeros(200,200,30), 
        zeros(40,40,20,30),
    )
    @btime fill!($PartArray,0.)
end
test1();
test2();

On Julia 1.6.0 this returns:

julia> test1();
  2.449 ms (0 allocations: 0 bytes)

julia> test2();
  5.693 ms (0 allocations: 0 bytes)

Which is what I would expect. On Julia 1.5.4, however, the result is:

julia> test1();
  2.649 ms (0 allocations: 0 bytes)

julia> test2();
  74.203 ms (4318980 allocations: 65.90 MiB)

These allocations cause a lot of garbage collector time, so I’d like to get rid of them.
Is there a solution other than getting the admins to update the Julia version on the Cluster to 1.6?
Another interesting hint to the possible issue is that in both Julia versions the type of the partition Array seems to be unstable:

julia> isconcretetype( ArrayPartition(zeros(2)) )
false

Does anyone know how I can construct an ArrayPartition such that the result is a concrete type?
Best wishes,
Salmon

*Edit: Formatting

2

Ok, I figured out a workaround:

Even though one can iterate over a eachindex of an ArrayPartition, it is probably better to “unpack” its elements when doing this, since an ArrayPartition can contain arrays of different types.

This seems to work:

using RecursiveArrayTools,BenchmarkTools

function setZero!(PartArr::ArrayPartition)
    for arr in PartArr.x
        fill!(arr,0.)
    end
end
    
function test2()
    PartArray = ArrayPartition(
        zeros(2),
        zeros(200,200,30), 
        zeros(40,40,20,30),
    )
    @btime fill!($PartArray,0.)  #old version
    @btime setZero!($PartArray)
end

julia> test2()
  72.663 ms (4318980 allocations: 65.90 MiB) #using fill!
  1.574 ms (0 allocations: 0 bytes) #using setzero!

Its should also be possible to overload Base.fill! for an ArrayPartition to allow the efficient use of fill!.