Memory Allocation on Views for Arrays

Hi there,

I would like to make a function that returns a view on an an 2-dimensional array. The view would change based on if the array has more rows and columns. MWE:

using BenchmarkTools
data1 = randn(100, 2)
data2 = randn(2, 100)

function grab(data, index::Integer, byrow::Bool)
    return byrow ? view(data, index, :) : view(data, :, index)
end

grab(data1, 10, true) #2-element view(::Matrix{Float64}, 10, :) with eltype Float64
grab(data2, 10, false) #2-element view(::Matrix{Float64}, :, 10) with eltype Float64

@btime grab($data1, $10, $true) #7.900 ns (1 allocation: 48 bytes)
@btime grab($data2, $10, $false) #8.000 ns (1 allocation: 48 bytes)
@code_warntype grab(data1, 10, true) #Body::Union{SubArray{Float64, 1, Matrix{Float64}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, SubArray{Float64, 1, Matrix{Float64}, Tuple{Int64, Base.Slice{Base.OneTo{Int64}}}, true}}

Now for some reason, I cannot make this function without allocations. Is there any clever trick to make the allocation go away? If I just define the row and column function separately, there is no allocation:

function grabcols(data, index::Integer)
    return view(data, index, :)
end
function grabrows(data, index::Integer)
    return view(data, :, index)
end

grabcols(data1, 10)
grabrows(data2, 10)
@btime grabcols($data1, $10) #2.300 ns (0 allocations: 0 bytes)
@btime grabrows($data2, $10) #2.300 ns (0 allocations: 0 bytes)

However, I am not able to make this allocation free by including the Boolean (and neither with defining a separate ByRow/ByCol struct and dispatching it). Is there any better way?

function grab2(data, index::Integer, byrow::Bool)
    return byrow ? grabcols(data, index) : grabrows(data, index)
end

grab2(data1, 10, true)
grab2(data2, 10, false)
@btime grab2($data1, $10, $true) #7.800 ns (1 allocation: 48 bytes)
@btime grab2($data2, $10, $false) #8.000 ns (1 allocation: 48 bytes)

You can dispatch on a type only for that (or on a value), for example:

julia> struct ByRow end

julia> struct ByCol end

julia> grab3(data, index::Integer, ::Type{T}) where T <: ByCol = grabcols(data, index)
grab3 (generic function with 3 methods)

julia> grab3(data, index::Integer, ::Type{T}) where T <: ByRow = grabrows(data, index)
grab3 (generic function with 4 methods)

julia> @btime grab3($data1, 10, ByCol)
  2.212 ns (0 allocations: 0 bytes)
2-element view(::Matrix{Float64}, 10, :) with eltype Float64:
 -0.8648160633699784
 -0.9484554544215145

julia> @btime grab3($data2, 10, ByRow)
  2.212 ns (0 allocations: 0 bytes)
2-element view(::Matrix{Float64}, :, 10) with eltype Float64:
 -0.18251801961171704
 -1.6380287971079939

But this will be useful if you can specialize the calling function for the operation/type you want. Otherwise it will create a dynamic dispatch and allocate something as well.

3 Likes

Something like this

using BenchmarkTools
data1 = randn(100, 2)
data2 = randn(2, 100)

function grab(data, index::Integer, byrow::Bool)
    return byrow ? view(data, index:index, 1:size(data, 2)) : view(data, 1:size(data, 1), index:index)
end

grab(data1, 10, true) 
grab(data2, 10, false) 

@btime grab($data1, $10, $true) 
@btime grab($data2, $10, $false) 
@code_warntype grab(data1, 10, true)

unifies the different view types for the prize of a different representation of the result.

  3.400 ns (0 allocations: 0 bytes)
  3.400 ns (0 allocations: 0 bytes)

Not sure if it helps, though.

1 Like

Thanks a lot, that worked perfectly!

1 Like

Thank you!