How to customize broadcast for a function that accepts keyword arguments?

General Usage
, ,
1

Suppose I have a function

julia> function f(n::Int)
           @assert n > 0
           n == 1 && return 1
           f(n-1) + 1
       end

and another function to cache the values of f(n) for a range of n:

julia> cachef(nr) = collect(nr)

Performance wise, broadcasting f over a range of n is much slower than calling cachef

julia> @btime cachef(1:1_000);
  736.652 ns (1 allocation: 7.94 KiB)

julia> @btime f.(1:1_000);
  4.881 ms (1 allocation: 7.94 KiB)

In this case, I may customize the broadcasting behavior of f as

julia> function Base.copy(B::Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1}, Tuple{Base.OneTo{Int64}}, typeof(f), Tuple{UnitRange{Int64}}})
           nr, = B.args
           cachef(nr)
       end

With this defined, I obtain

julia> @btime f.(1:1_000);
  721.657 ns (1 allocation: 7.94 KiB)

This shows that broadcasting now uses cachef. This works for functions that only accept positional arguments. How does one generalize this to functions that accept keyword arguments?

For example

julia> g(n; a=nothing) = f(n);

julia> @btime g.(1:1000, a=nothing);
  4.874 ms (1 allocation: 7.94 KiB)

So, this doesn’t inherit the broadcasting performance of f, which is unfortunate. Looking into the code generated:

julia> Meta.@lower g.(1:1000, a=nothing)
:($(Expr(:thunk, CodeInfo(
    @ none within `top-level scope`
1 ─ %1 = 1:1000
│   %2 = Base.broadcasted_kwsyntax
│   %3 = Core.tuple(:a)
│   %4 = Core.apply_type(Core.NamedTuple, %3)
│   %5 = Core.tuple(nothing)
│   %6 = (%4)(%5)
│   %7 = Core.kwfunc(%2)
│   %8 = (%7)(%6, %2, g, %1)
│   %9 = Base.materialize(%8)
└──      return %9
))))

Is Core.kwfunc returning an anonymous function here? Can one specialize broadcast for such a function? Given that none of this is documented, I am wary of using this, but then how does one specialize broadcast for a function that accepts keyword arguments?