make indexing expressions participate in dot syntax fusion

[StefanKarpinski]

Consider an expression like v[p.(v)] .*= 2 where p is a predicate (i.e. returns a boolean). Currently Julia's semantics are that p.(v) is constructed and then used to index into v. However, the optimal implementation of this would to loop through v, testing each element x in v to see if p(x) is true and then multiplying it by 2 if that's the case. While it might be possible for a compiler to optimize the current semantics to this behavior in some circumstances, arguably, the dotted application of p in the indexing expression should participate in dot fusion of the whole expression, and this expression should simply mean the optimal implementation. Marking as 1.0 since if we're going to change the semantics of this expression, we should do so before then.

[KristofferC]

Ref #19169

[stevengj]

I would tend to prefer v.[p.(v)] for something like this. If we lowered a.[i] to getindex.(a, i), then v.[p.(v)] would already work on the right-hand-side of assignments since getindex.(v, p.(v)) is fusing.

[JeffBezanson]

OK, I guess this is just a parser issue then. I can allow parsing x.[...].

[JeffBezanson]

Update: already parses, so just needs a case in the broadcast lowering code.

[mbauman]

It'll take quite a bit more than that, particularly in the logical indexing case like Stefan proposed in his example. The naive broadcast will just repeatedly pass a true or false index to the array. See my comment at #19169 (comment).

[JeffBezanson]

Yeah, I just realized that and was about to post again. a[i] is just not generally equivalent to getindex.(a, i).

[StefanKarpinski]

Writing this as v.[p.(v)] .*= 2 would be fine as well as long as there is a way to write it.

[JeffBezanson]

a.[i] would mean [ getindex(a[1], i[1]), getindex(a[2], i[2]), ... ] or [ getindex(a[1], i), getindex(a[2], i), ... ] if i is a scalar. So we'd have to change something else to make that work, but I'm not sure what.

[StefanKarpinski]

Since the v.[p.(v)] syntax doesn't mean anything and gives an invalid syntax error, this is not 1.0-blocking:

julia> v.[p.(v)]
ERROR: syntax: invalid syntax v.[p.((v,))]

[mbauman]

Now that I've recently thought through this, I figured I'd leave a note for posterity (and future me) on the difficulties here:

An expression like A.[B .& C, D] .* E would naively become something like broadcast((b,c,d,e)->A[b & c, d] * e, B, C, D, E). This obviously cannot work, and it would be extremely difficult to come up with a transformation that could work:

• As I note above, it would repeatedly pass true or false as indexes. We need some way of transforming them into the indices where they are true. But this is hardly the biggest challenge.
• Broadcast needs to know that B and C are involved in a subexpression that becomes a logical index, and then skip values for which that subexpression produces falses. That is we cannot walk through all four arrays synchronously; D and E must only yield values when the logical subexpression is true.
• Indexing with a two-dimensional logical mask yields a vector of results. Broadcast would need to treat all variables participating in the logical subexpression as a single vector.
• Broadcast doesn't know the result size without running through and counting the trues in the logical subexpression ahead of time.
• Broadcast would have a hard time doing the same kinds of "outer" bounds checking that non-scalar indexing currently performs.

Logical indexing is really the worst case here. But there are other problems. : would have to be handled specially; other fancy index types like .. intervals that some packages have defined would have similar troubles.

Were we to do this, we'd only be able to support arrays of integer indices or we'd have to disable fusion from propagating inside the .[] syntax entirely (we'd broadcast over the indices after they've been transformed into collections of integers by to_indices). If we chose the latter semantics and also transformed the indices to be orthogonal to each other, then this is effectively the status quo with @views @.… except maybe a bit worse since right now broadcast can trust that accesses into views are inbounds! The benefit, on the other hand, would be that this is semantically the same as APL indexing… and we could deprecate nonscalar getindex and setindex! entirely.

[StefanKarpinski]

Nice writeup. Seems like the best course of action is to leave this as an error until we can handle it.

[mbauman]

Latest thoughts:

In order to get around the difficulties I enumerated in my last post, I had been thinking I'd like to define A.[B .& C, D] .* E as:

broadcast((x,y,z)->A[x,y] * z, orthogonalize(to_indices(A, B .& C, D)...)..., E)

where orthogonalize adds leading singleton dimensions to each successive argument such that the number of leading singletons is equal to the sum of the dimensionalities of all the preceding arguments. We effectively sacrifice our ability to fuse individual index computations with the rest of the expression in order to gain APL indexing semantics and fancy index types. So B .& C allocates an array, but the accesses into A need neither a view nor an allocation — they're done on-demand.

Here's the unfortunate part: orthogonalize requires reshape to do its work. And reshape is a wrapper much like view (or for Array, it's a new header linked to the same data). As long as a view requires allocations, so would this syntax transform. Then the million dollar question is: why would we do all this extra work to define a new syntax that's effectively identical to @views @. A[B & C, D] * E in both functionality and performance!? Let's not do that.

[pabloferz]

Sorry, closed it by github-mobile-interface accident.