Creating types inside broadcast expressions can lead to allocations

[charleskawczynski]

Reproducer:

#=
julia --project=perf
include("../perf/thermo_inference_repro.jl")
include("../../perf/thermo_inference_repro.jl")
=#
using Revise
import ClimaCore;
import ClimaCore.Fields as Fields
import ClimaComms
@isdefined(TU) || include(joinpath(pkgdir(ClimaCore), "test", "TestUtilities", "TestUtilities.jl"));
import .TestUtilities as TU;
FT = Float64;
# device = ClimaComms.device()
device = ClimaComms.CPUSingleThreaded()
space = TU.CenterExtrudedFiniteDifferenceSpace(FT; context=ClimaComms.context(device), zelem=25, helem=10);
@show device

struct PhasePartition{FT <: Real}
    tot::FT
    liq::FT
    ice::FT
end

struct DummyPhaseEquil{FT}
    q_tot::FT
end

@inline PhasePartition(ts::DummyPhaseEquil{FT}) where {FT} =
    PhasePartition(FT(0.01), FT(0.001), FT(0))

function test_allocs_inline1!(x)
    (; ts, q_tot) = x
    @. q_tot = PhasePartition(ts).tot # allocates!
    return nothing
end

@inline total_specific_humidity(ts) = PhasePartition(ts).tot
function test_allocs_inline2!(x)
    (; ts, q_tot) = x
    @. q_tot = total_specific_humidity(ts) # allocation-free
    return nothing
end

x = (;
    ts   = Fields.Field(DummyPhaseEquil{FT}, space),
    q_tot = Fields.Field(FT, space),
)
test_allocs_inline1!(x) # compile
test_allocs_inline2!(x) # compile

p_allocated1 = @allocated test_allocs_inline1!(x)
p_allocated2 = @allocated test_allocs_inline2!(x)
@show p_allocated1 # 5760128
@show p_allocated2 # 0
@info "Done"

We should add a broken test for this.

[charleskawczynski]

This is an issue in Julia Base broadcasting, and I confirmed that, on the CPU, this is fixed in Julia 1.11.