CliMA / ClimaCore.jl

CliMA model dycore
https://clima.github.io/ClimaCore.jl/dev
Apache License 2.0
87 stars 8 forks source link

Specialize on `copyto!(::Field, ::typeof(broadcasted(identity, ::Field)))` #1887

Open charleskawczynski opened 3 months ago

charleskawczynski commented 3 months ago

We can easily specialize on @. x = y where typeof(x) == typeof(y) to dispatch into the copyto!(::Field, ::Field) method, which is ~2x faster. There aren't many cases where we do this, but Base and CUDA do this-- it's still probably worth it.

charleskawczynski commented 3 months ago

From Base.Broadcast:

    # Performance optimization: broadcast!(identity, dest, A) is equivalent to copyto!(dest, A) if indices match
    if bc.f === identity && bc.args isa Tuple{AbstractArray} # only a single input argument to broadcast!
        A = bc.args[1]
        if axes(dest) == axes(A)
            return copyto!(dest, A)
        end
    end