Lazy broadcasts for reductions / more broadcasting customization docs

[chethega]

Another simple example for broadcast customization, that demonstrates how to reduce broadcast results without allocating intermediaries.

I am not sure; maybe the lazy broadcast even wants to become part of base, as a new feature (after bikeshedding with respect to names)?

[mbauman]

Yeah, I think this should just be a feature instead of a doc example.

This is a really clever way of going about it - I had been thinking we’d use a macro and need to re-implement some of the lowering transform to get this done. Even though this doesn’t need a macro I think the macro interface makes a lot of sense still.

[chethega]

Ok, then let the great bikeshedding begin. We could have

struct lazy end
struct _Lazy{T}
  arg::T
end
Base.Broadcast.materialize(ell::_Lazy) = ell.arg
Base.Broadcast.broadcasted(::Type{lazy}, arg) = _Lazy(arg)

and export lazy and don't export _Lazy. Then, the manual would point at lazy twice: Once in the user manual for its purpose of reducing broadcast results without materializing them, and once in the section about broadcast customization, as an educational tool for exploring how the broadcast machinery works.

Regarding a macro: That might be nice, but I don't have any opinions on how this would look like. On syntax/naming, a loosely binding unary ascii operator would be perfect: The parentheses of lazy.($expr) are annoying to write.

We also need a warning in the docstring that lazy performs no DimensionMismatch checks ; these are on access. This means that a lazy broadcast might violate conventions on validity of @inbounds and on what operations can throw.

There would be an argument for a variant of lazy broadcasts that eagerly checks DimensionMismatch and possibly non-emptyness, and records this into the type. Then it would be easier to work with the result. But I like the fact that this feature is a 4-liner ;)

In order to make reduction of lazy broadcasts perform well, we need to work a little on reductions. However, I would say that we can push this new feature, and later improve performance of typical reductions, as user reports come in.

At least some @propagate_inbounds are probably in order, type stability is awkward, and successful simd is often more important than avoidance of the intermediate array. We can work on all, any, fold, etc afterwards, worst case by adding methods for Broadcasted, and potentially by future compiler improvements (when using the lazy pattern, it is not entirely trivial to coax the compiler into avoiding the tiny allocation of the Broadcasted object).

[vchuravy]

I would actually like to see Lazy reductions developed outside of Base first. I think the design space is rather large and I don't want us to settle on a particular implementation eagerly.

(Full disclosure, Peter and I have been working and thinking about lazy reductions and broadcast for a while now)

[chethega]

Is there an issue/PR/archive/package with some results of your thoughts? Or should we use this PR for a discussion?

Avoiding to lock in premature design decisions certainly makes sense. Then we would just put this as an example into the documentation, and potentially update it once the new feature is polished?

[tkf]

This API is discussed in #19198. FYI, I also have a proof-of-concept implementation of a DSL for fused map-filter-reduce based on broadcast and Transducers.jl: https://github.com/tkf/Mathy.jl

Note that using reduce etc. with broadcasting is slow at the moment because they use slow iterator interface. So, I think it makes sense to use fast eachindex-based loop in Base.mapreduce etc. as otherwise it's difficult to implement it without type piracy.

[chethega]

Because the question occasionally pops up on discourse, also gave a more explicit example for "wrapping in a Ref to behave like a scalar", and clarified that the convention collect(x) == collect(Base.broadcastable(x)) only holds if the LHS is well-defined, and not for e.g. missing. The arguments against exporting something like lazy/air at this moment are quite convincing, so this is done from my side.

[mbauman]

I still don't like documenting this before we implement it. If we're going to document it, we might as well just implement it in Base and export it.

Right now, only broadcast implementors are touching Broadcasted objects. With this change we expose the API to a much larger surface and it becomes near-impossible to make meaningful changes. So either we decide that this is what we want and we export it, or we wait for more design iterations to be done.

[chethega]

The example is in the manual section for people who want to customize broadcasts. This necessitates a minimum understanding of the broadcasting machinery, and I find that copy-paste-able examples that can be used interactively are a great way to get started understanding something.

Something like lazy/air is what I would have liked to see as an early example in the docs, in order to better understand how all this currently works. Upon reconsideration, I would also like an initial explanation of the lowering: Dot-syntax is syntactic sugar for Base.broadcasted, Base.materialize, Base.materialize! and Base.dotview, plus a code excerpt for function broadcasted(f, arg1, arg2, args...), which demonstrates together with the broadcastable(x::String)=Ref(x) excerpt how broadcastable works. I think users who want to customize broadcasts and use macros on broadcast expressions need this anyway?

And while some minor changes of broadcast lowering would be great going forward (there were some issues where unneeded views where allocated), let's not kid ourselves that this could be done without breakage. Most of the other syntactic sugars have their lowering explained in the docs (getfield, getindex, setindex, getproperty, maybe even for-loops in the iteration protocol?).

The explicit excerpt about wrapping in Ref was added for the same reason: Displayed with triple back-ticks instead of single back-ticks in order to allow users to quickly scan for the most important example ("I want to be a scalar" / "I am a wrapper and want to broadcast as parent" are probably the most common broadcast customizations).

[tkf]

I would also like an initial explanation of the lowering

I think this is a good idea and it's better to focus on it rather than introducing utility functions/macros in the doc. I think it would be more appropriate to mention Meta.@lower usage with something like this:

julia> Meta.@lower x .+ y .* z
:($(Expr(:thunk, CodeInfo(
    @ none within `top-level scope'
1 ─ %1 = Base.broadcasted(*, y, z)
│   %2 = Base.broadcasted(+, x, %1)
│   %3 = Base.materialize(%2)
└──      return %3
))))

and demonstrate how broadcasting expressions are evaluated.

[tkf]

See my PR #31088 for adding the macro to Base.

[chethega]

Superseded by #31088.