In-place assignment, without allocation

I’m trying to do y = x[ind] in-place, where y, x, ind are all vectors. By in-place, I mean no memory is allocated. I’m able to achieve this using a hand-written for loop (assign2! below). I suspect there is a way to achieve this using broadcasting, which could potentially generate better code. However, I haven’t been able to find it. assign! below is my attempt at it. Is there a simple line of code using broadcasting that works (doesn’t allocate, & is fast)? Why doesn’t the @. find it?

using BenchmarkTools, Random

x = randn(100)
y = zeros(200)
ind = vcat(randperm(100), randperm(100))

assign!(y, x, ind) = @inbounds @. y = x[ind]

@btime assign!($y, $x, $ind);

function assign2!(y, x, ind)
    @simd for i in eachindex(ind)
        @inbounds y[i] = x[ind[i]]
    end
end

@btime assign2!($y, $x, $ind);
1 Like
assign!(y,x,ind) = y.=getindex.(Ref(x), ind)

is allocation free. The problem in your example is that x[ind] is not a broadcasted, hence it doesn’t fuse with the assignment.

By the way, the @simd is useless here.

3 Likes

EDIT: I think I misunderstood the question - sorry.
Granted, this has one allocation, but this seems pretty fast to me:

x = randn(10)
y = zeros(10)
@btime y .= x
  22.051 ns (1 allocation: 16 bytes)

Why is it that this does not allocate:

@btime assign!(y, x, 3)
  10.063 ns (0 allocations: 0 bytes)

… but this does?

@btime y .= getindex.(x, 3)
  19.719 ns (1 allocation: 32 bytes)

Regardless of the speed, it is strange to me that this allocates.

You need to do @btime $y .= $x, which shows no allocations.

4 Likes

I’m not yet good enough to unravel everything, but the best I can figure is if I wrote it in C, it would look something like:

int x[10];
int y[10]; // or whatever

for (int i = 0; i < 10; ++i)
    y[i] = x[i];

The allocation would be the loop variable?

It doesn’t, but you must remember to interpolate external variables in benchmark expressions

julia> @btime $y .= $x
  15.415 ns (0 allocations: 0 bytes)
2 Likes

No, “allocation” means heap allocation which need to be tracked and eventually garbage collected. That counter would live happily on the stack.

2 Likes