@fastmath doesn't apply to broadcasted infix operators and causes syntax errors

[Keluaa]

Currently using Julia 1.10.1.

The @fastmath macro ignores .+-like operators:

julia> @macroexpand @fastmath 1 + 2
:(Base.FastMath.add_fast(1, 2))

julia> @macroexpand @fastmath 1 .+ 2
:(1 .+ 2)

This is because .+ is a single Symbol in the AST:

julia> dump(:(1 .+ 2))
Expr
  head: Symbol call
  args: Array{Any}((3,))
    1: Symbol .+
    2: Int64 1
    3: Int64 2

But @fastmath only replaces non-broadcasted symbols.

The problem is even more apparent when using .+ alone, e.g. in a mapreduce, as it gives a syntax error:

julia> f(x) = (x-1, x+1)
f (generic function with 1 method)

julia> x = rand(Float64, 1000);

julia> @fastmath mapreduce(f, .+, x)
ERROR: syntax: invalid syntax Base.FastMath.add_fast
Stacktrace:
 [1] top-level scope
   @ REPL[19]:1

julia> mapreduce(f, Broadcast.BroadcastFunction(Base.FastMath.add_fast), x)
(-490.2706650247785, 1509.729334975223)

By hacking Base.FastMath.make_fastmath we can partially bypass this limitation:

const broadcast_fastmath = Dict(
    :var".+" => :add_fast,
    :var".-" => :sub_fast,
    :var".*" => :mul_fast,
    :var"./" => :div_fast,
    :var".==" => :eq_fast,
    :var".!=" => :ne_fast,
    :var".<" => :lt_fast,
    :var".<=" => :le_fast,
    :var".>" => :gt_fast,
    :var".>=" => :ge_fast,
    :var".^" => :pow_fast
)

function Base.FastMath.make_fastmath(s::Symbol)
    fast_symb = get(Base.FastMath.fast_op, s, :nothing)
    if fast_symb === :nothing
        fast_broadcast = get(broadcast_fastmath, s, :nothing)
        if fast_broadcast === :nothing
            return s
        end
        return :(Broadcast.BroadcastFunction(Base.FastMath.$fast_broadcast))
    end
    return :(Base.FastMath.$fast_symb)
end

Then:

julia> @macroexpand @fastmath 1 .- 2
:((Broadcast.BroadcastFunction(Base.FastMath.sub_fast))(1, 2))

But more work is needed in Base.FastMath.make_fastmath(::Expr) to support .+ outside of math expressions (e.g. for mapreduce), as well as broadcasted assignments .+=.