Problem with broadcasting, flatten, and type-stability

General Usage
1

I’m writing a custom AbstractArray implementation which supports broadcasting. This implementation makes use of Broadcast.flatten, but I’m finding that this introduces type-instabilities once expressions go beyond a few operations.

Here’s a MWE. Rather than give you my whole custom AbstractArray type, you can see it by just overloading normal Array broadcasting to insert a call to flatten:

Broadcast.preprocess(dest::Array, bc::Broadcast.Broadcasted{Nothing}) = Broadcast.flatten(bc)

Then here’s an expression which is just complicated enough to trigger it:

x = rand(2,2)
foo(x) = @. (x + 2*x + 2*x)*x
@code_warntype foo(x)

Scrolling through the long output, you’ll see the red. Although the return type is inferred, the instabilities within result in e.g. a 1000x performance hit on a 128x128 array (checked on Julia 1.1 and release-1.2/8a84ba5018)

Am I doing anything stupid here, and if not does anyone have strategies to remedy this or make things easier on inference? Should I file an Issue? Thanks.

2

Oh yea, I thought this sounded familiar, I realize now I already filed it a year ago to much interesting discussion :slight_smile: inference failure for broadcast implementations that use `flatten` (eg SparseArrays) · Issue #27988 · JuliaLang/julia · GitHub

Unfortunately still no solution. The underlying thing being the fairly weird fact that the following does not infer, but it does if you start a fresh session and uncomment the commented line:

using Test
using Base.Broadcast: cat_nested, Broadcasted
#@inferred cat_nested(1, Broadcasted(*, (2, 3)), 4, 5, 6, 7)
@inferred cat_nested(Broadcasted(*, (1, Broadcasted(*, (2, 3)), 4, 5, 6, 7)))
2 Likes