Julep: Syntax for reduction (overloadable splatting?)

[tkf]

I suggest a mechanism to customize splatting behavior such that code like

+((x .- y).^2...) / length(x)  # compute MSE

can be executed efficiently without any allocations.

The idea is inspired by discussion in #29114 by @c42f et al.

Idea

Like dot-call syntax, I suggest to lower splatting to a series of function calls that are overloadable. One possibility is

# f(a...) is expanded to:
apply(f, Arguments(VA(splattable(a))))

# f(a, b..., c, d...) is expanded to:
apply(f, Arguments(a, VA(splattable(b)), c, VA(splattable(d))))

(To avoid recursion, this lowering should happen only when the call includes splatting.)

When the dot-call syntax appears in the splatting operand, I suggest not materialize the dot-call. That is to say, for example, op(f.(xs)...) is lowered to

bc = broadcasted(f, xs)  # `bc` not materialized
apply(op, Arguments(VA(splattable(bc))))

This let us evaluate op(f.(xs)...) without any allocation once reduce/foldl supports Broadcasted object (#31020 started tackle this).

Interface

The lowering above requires the following interface functions and types:

function apply end

splattable(x) = x

struct VA{T}
    args::T
end

struct Arguments{T <: Tuple}
    args::T
    Arguments(args...) = new{typeof(args)}(args)
end

• apply must be dispatched on the first argument type and may be dispatched on the second argument type.
• Per-vararg processing should be done via splattable. This is analogous to broadcastable (@yurivish suggested this in Performance of splatting a number #29114 (comment)).
• The type VA (whose name can/should be improved) must be used only for defining apply; its constructor must not be overloaded. This is for making it hard to break splatting semantics.
• The constructor for type Arguments must not be overloaded for the same reason.

Using current Core._apply, the default apply can be implemented as

apply(f, x) = Core._apply(Core._apply, (f,), _default_splattables(x.args))
_default_splattables(args) = map(x -> x isa VA ? materialize(x.args) : (x,), args)

Example overloads

Associative binary operators

Many useful operations can be expressed using splatting into associative operators

+(xs...)      == sum(xs)
*(xs...)      == prod(xs)
min(xs...)    == minimum(xs)
max(xs...)    == maximum(xs)

or "mapped-splatting"

+(xs.^2...)      == norm(xs)
+(xs .* ys...)   == dot(xs, ys)

(This is reminiscent of the "big operator" in Fortress.)

There are also various other associative binary operators in Base. Invoking reduce with splatting could be useful:

const AssociativeOperator = Union{
    typoef(*),
    typoef(+),
    typoef(&),
    typoef(|),
    typoef(min),
    typoef(max),
    typoef(intersect),
    typoef(union),
    typoef(vcat),
    typoef(hcat),
    typoef(merge),
    # what else?
}

apply(op::AssociativeOperator, args::Arguments{Tuple{<:VA}}) =
    reduce(op, args.args[1].args)

For example, concatenating vectors would be efficiently done via vcat(vectors...) thanks to #27188.

It also is possible to support

op :: AssociativeOperator
op(a, bs..., cs...)

such that it is computed as

op(op(a, reduce(op, bs)), reduce(op, cs))

This may be implemented as

apply(op::AssociativeOperator, args::Arguments) = mapreduce(apply1, op, args.args)
_apply1(op::AssociativeOperator, args::VA) = reduce(op, args.args)
_apply1(op, x) = x

Note that, since reduce would degrade to foldl when the input is not an array (or not Broadcasted after #31020), we can also use it to fuse filtering with reduction

+((x for x in xs if x > 0)...)

Non-associative binary functions

Splatting is useful for non-associative binary functions:

const BinaryFunction = Union{
    typoef(/),
    typoef(-),
    typoef(intersect!),
    typoef(union!),
    typoef(merge!),
    typoef(append!),
    typoef(push!),
    # what else?
}

function apply(op::BinaryFunction, args::Arguments{Tuple{Any, <:VA}})
    @assert !(args.args[1] isa VA)
    return foldl(op, args.args[2].args; init=args.args[1])
end

or more generally

function apply(op::BinaryFunction, args::Arguments)
    @assert !(args.args[1] isa VA)
    return foldl(op, flatten(x isa VA ? x.args : (x,) for x in args.args[2:end]); init=args.args[1])
end

Matrix-vector multiplications

Not sure how many people need this, but *(matrices..., vector) can be (somewhat) efficiently evaluated by defining

apply(::typeof(*), args::Arguments{Tuple{<:VA, <:AbstractVector}}) =
    foldr(*, args.args[1].args; init=args.args[2])

(Of course, allocation could be much more minimized if we really want this.)

Similar optimization can be done for ∘(fs...); but I'm not sure about the exact usecase.

Higher-order functions (map(f, iters..) etc.)

As this mechanism let any function optimize splatting, higher-order functions that may call splatting of given function can be optimized by defining their own apply specialization. For example, map(f, iters..) can be specialized as

function apply(::typeof(map), args::Arguments{Tuple{Any, <:VA}})
    f = args.args[1]
    iters = args.args[2].args
    return map(splat(f), _zipsplat(iters))
end

where _zipsplat(iters) behaves like zip(iters...) but its element type does not have to be a Tuple.

Defining a similar overload for broadcasted may be possible provided that the object returned by _zipsplat(iters) is indexable. This let us nest reduction inside mapping and avoid allocation in some cases:

vector .= +.(eachcol(matrix)...)

Other map-like functions including map! and foreach can also implement this overload.

print-like functions

print, println and write can be invoked with varargs. We can make, e.g., println(xs...) more compiler friendly and efficient when xs is a generic iterator. Note that apply(string, Arguments(xs)) can also be implemented in terms of apply(print, Arguments(io, xs)).

splattable

splattable may be used to solve performance problem discussed in #29114:

splattable(x::Number) = (x,)
splattable(x::StaticArray) = Tuple(x)

In case of Broadcasted, it can be used for calling instantiate:

splattable(x::Broadcasted) = instantiate(x)

[JeffBezanson]

I'm pretty strongly against this. Actually, I think the idea is not best described as "overloadable splatting", but as "change ... to syntax for reduce instead of splatting". And that's not a bad idea --- maybe we should just add some syntax for reduce.

Splatting is supposed to be completely orthogonal to function calling. Calling functions is the core operation in the language, and splatting is just an alternate way to provide arguments. If this proposal were adopted, we would need some other way to call a function on arguments taken from a container. For example, this change makes it impossible to implement utilities like

∘(f, g) = (x...)->f(g(x...))

because there would be no way to generically take some arguments, and then pass all of them on to another function as if it were called directly with them.

So in short I'd favor coming up with some syntax for reduce rather than re-purposing ... and removing any ability to have a first-class representation of an argument list.

[tkf]

Would it be a problem if there is a way to seal methods (#31222)? I imagine that #31222 allows us to construct lowering such that it is impossible to change the behavior when splatting tuples.

I understand that function call is too fundamental so that you don't want to play any magic with it. At the same time, since it is very fundamental, many people understand its semantics; including the users who are not very familiar with functional programming construct like reduce. Although..., if there is some syntax that resembles splatting, maybe that's enough? (I am a bit ambivalent about a new syntax for reduce. It's certainly is a huge win if it has a concise syntax that interacts nicely with broadcasting.)

[c42f]

concise syntax that interacts nicely with broadcasting

Yes if we can get syntax for mapreduce which reuses broadcast syntax for the mapping part that could be compelling.

Maybe we just need more dots in more places 😬

+...(xs.^2)   # mapreduce(x->x^2, +, xs)
              # [Edit: really reduce(+, broadcasted(literal_pow, ^, xs, Val(2)))]

(motivation: an ellipsis indicates more items in the list. In this case "many +'s")

[tkf]

@c42f Moving ... to outside is interesting! I thought comprehensions like +((x for x in xs if x > 0)...) could become hard to read when the expression is large. Putting ... right next to the binary operator is a good idea because you know that it is invoking a reduction right away: +...(x for x in xs if x > 0).

By the way, we only need the syntax for reduce and foldl because mapreduce and mapfoldl can be "derived" from reduce and foldl. That is to say, mapreduce(f, op, xs) is equivalent to reduce(op, broadcasted(f, xs)) once we have #31020. Generally speaking, it is actually redundant to define mapreduce/mapfoldl if you have transducers or iterator transforms (which include broadcasted).

[c42f]

we only need the syntax for reduce and foldl because mapreduce and mapfoldl can be "derived" from reduce and foldl

Agreed; all I really mean is that you want to be able compile to something as efficient as mapreduce using existing broadcast notation combined with some reduction syntax so that lowering can avoid emitting the intervening Base.materialize.

[c42f]

I checked whether +...(xs.^2) is currently a parse error.

To my surprise:

julia> :(+...(xs.^2))
:((Core.:(@int128_str))(".") .. xs .^ 2)

Haha! What.

[c42f]

Ok, harder problem: can we get syntax for reduce(f, A; init=x, dims=d)?

For init, the following might, maybe, be possible:

0 +... xs  # Now we're talking!
f...(0, xs)

For dims, the composable way seems to be broadcasting the reduction across slices from eachslice ().

(+...).(eachslice(A, dims=2))    # reduce(+, A, dims=1) if A is a matrix
+....(eachslice(A, dims=2))      # ?? Is it even possible to parse this? :-(

(+...).(eachslice(A, dims=2).^2)  # Well, it should work with broadcast at least

That's... a lot of dots...

Though a bit tangential, this discussion also reminds me of @mbauman's comment at #31217 (comment)

[c42f]

Dealing with init makes me think that the rule should be something like

f...(a, xs; h=1, g=2)   # reduce((x1,x2)->f(x1, x2; h=1, g=2), xs, init=a)

but that naturally leads to wondering about other arguments to the "binary" operator

f...(a, xs, b, c; h=1, g=2)  # Unsure whether this can have a sensible meaning.

[tkf]

That's... a lot of dots...

Maybe throw more dots? How about

+...(x.^2 .+ 1)[., ., :]

# Equivalent to:
dropdims(reduce(+, (for x.^2 .+ 1); dims=(1, 2)); dims=(1, 2))
# `.`s inside `[., ., :]` become `dims`

and

y .= +...(x.^2 .+ 1)[., ., :]

# Equivalent to:
reducedim!(+, reshape(fill!(y, 0), (1, 1, size(y)...)), (for x.^2 .+ 1))
# `.`s insert singleton dimensions in the destination array

(I'm using for broadcasting_expression notation from #31553)

The dropdims in the first example is motivated by the second example. But I'm not super sure about this.

Regarding init, I think it is the reason why splatting is so appealing:

+(a, b...)           # reduce(+, (a, reduce(+, b))) or reduce(+, b; init=a)
+(a, b..., c..., d)  # reduce(+, (a, reduce(+, b), reduce(+, c), d))
push!([], b...)      # foldl(push!, b; init=[])

[tkf]

Actually, [.] can be used instead of splatting?:

+(a, b[.])           # => reduce(+, (a, reduce(+, b))) or reduce(+, b; init=a)
+(a, b[.], c[.], d)  # => reduce(+, (a, reduce(+, b), reduce(+, c), d))
push!([], b[.])      # => foldl(push!, b; init=[])
write(io, a[.])      # => foreach(x -> write(io, x), a)

+(x.^2 .+ 1)[., ., :]
# => dropdims(reduce(+, (for x.^2 .+ 1); dims=(1, 2)); dims=(1, 2))

y .= +(x.^2 .+ 1)[., ., :]
# => reducedim!(+, reshape(fill!(y, 0), (1, 1, size(y)...)), (for x.^2 .+ 1))

It also motivates dropdims.