Broadcasted indexing

[jlchan]

Would it be difficult to add support for broadcasted indexing? For example, accessing arrays/tuples of arrays

function foo(x)
    @avx for i = 1:length(first(x))
        x_i = getindex.(x,i)
    end
end

using @avx gives an Expression not recognized error.

[chriselrod]

If x is some form of tuple, it'd be possible using the same approach it uses for CartesianIndexing, i.e. I could make it so that when you run

function foo(x)
    @avx for i = 1:length(first(x))
        x_i = getindex.(x,i)
        y_i = exp.(x_i)
        setindex!.(x, y_i, i)
    end
end

x = (rand(100), randn(100), randexp(100))
foo(x)

It can basically expand

x_i = getindex.(x,i)

into

x_i_1 = getindex(x_1, i)
x_i_2 = getindex(x_2, i)
x_i_3 = getindex(x_3, i)

y_i_1 = exp(x_i_1)
y_i_2 = exp(x_i_2)
y_i_3 = exp(x_i_3)

setindex!(x_1, y_i_1, i)
setindex!(x_2, y_i_2, i)
setindex!(x_3, y_i_2, i)

But for now, this works and has the same effect:

using LoopVectorization, Base.Cartesian
@generated function foo!(x::Tuple{Vararg{Array,N}}) where {N}
    quote
        @nextract $N x x
        @avx for i = 1:length(first(x))
            @nexprs $N n -> x_n[i] = exp(x_n[i])
        end
    end
end

using Random
x = (rand(100), randn(100), randexp(100))
foo!(x)
x

If x is an array...yes, I do plan on supporting that.
What I'd like to do eventually is support a wider range of loopy operations into loops (including, but not limited to, broadcasting):

@avx begin
    y = exp.(w)
    z = exp.(A * x)
    b = log.(y .+ z)
    s = sum(abs2, b)
end

With also the option to fuse loops, but there'd preferably be some sort of syntax for saying w and A*x are the same long, although in this particular example it'd also be possible to infer from the y .+ z.
It would also need the ability to specify whether the outputs (y,z, b, and s) are returned. If not, it will not have to allocate temporaries.

This would also support nesting within loops.

This is all a ways away, but something I do plan on supporting eventually.

[jlchan]

Thanks for the detailed response! This is helpful (also for getting a sense for how LoopVectorization works).