`orchestrate_diagnostics` has lots of inference failures

[charleskawczynski]

orchestrate_diagnostics has lots of inference failures, and I've noticed that compilation time is significantly reduced when diagnostics are disabled. This very well may be the reason.

Found in https://buildkite.com/clima/climaatmos-ci/builds/20805#01922a42-bcbe-49b0-ae2d-2247cb3b2a49

││││││││││┌ orchestrate_diagnostics(integrator::ClimaTimeSteppers.DistributedODEIntegrator{…}, diagnostic_handler::ClimaDiagnostics.DiagnosticsHandler{…}) @ ClimaDiagnostics /central/scratch/esm/slurm-buildkite/climaatmos-ci/depot/default/packages/ClimaDiagnostics/Qy3We/src/clima_diagnostics.jl:195
  | │││││││││││┌ interpolate_field!(writer::ClimaDiagnostics.Writers.NetCDFWriter{…}, field::Any, diagnostic::ClimaDiagnostics.ScheduledDiagnostics.ScheduledDiagnostic{…}, u::ClimaCore.Fields.FieldVector{…}, p::ClimaAtmos.AtmosCache{…}, t::Float32) @ ClimaDiagnostics.Writers /central/scratch/esm/slurm-buildkite/climaatmos-ci/depot/default/packages/ClimaDiagnostics/Qy3We/src/netcdf_writer.jl:220
  | ││││││││││││┌ target_coordinates(space::ClimaCore.Spaces.ExtrudedFiniteDifferenceSpace, num_points::Tuple{…}, z_sampling_method::ClimaDiagnostics.Writers.LevelsMethod) @ ClimaDiagnostics.Writers /central/scratch/esm/slurm-buildkite/climaatmos-ci/depot/default/packages/ClimaDiagnostics/Qy3We/src/netcdf_writer_coordinates.jl:494
  | │││││││││││││┌  @ ClimaCore.Spaces /central/scratch/esm/slurm-buildkite/climaatmos-ci/depot/default/packages/ClimaCore/wHC4M/src/Spaces/finitedifference.jl:19
  | ││││││││││││││┌ ClimaCore.Grids.FiniteDifferenceGrid(topology::ClimaCore.Topologies.IntervalTopology) @ ClimaCore.Grids /central/scratch/esm/slurm-buildkite/climaatmos-ci/depot/default/packages/ClimaCore/wHC4M/src/Grids/finitedifference.jl:46
  | │││││││││││││││┌ get!(default::ClimaCore.Grids.var"#1#2"{…} where {…}, h::Dict{…}, key::Tuple{…}) @ Base ./dict.jl:479
  | ││││││││││││││││┌ (::ClimaCore.Grids.var"#1#2"{…} where {…})() @ ClimaCore.Grids /central/scratch/esm/slurm-buildkite/climaatmos-ci/depot/default/packages/ClimaCore/wHC4M/src/Grids/finitedifference.jl:47
  | │││││││││││││││││┌ _FiniteDifferenceGrid(topology::ClimaCore.Topologies.IntervalTopology) @ ClimaCore.Grids /central/scratch/esm/slurm-buildkite/climaatmos-ci/depot/default/packages/ClimaCore/wHC4M/src/Grids/finitedifference.jl:60
  | ││││││││││││││││││┌ (::Type{ClimaCore.DataLayouts.VF{_A, _B}} where {_A, _B})(::Type{Array{_A}} where _A, nelements::Any) @ ClimaCore.DataLayouts /central/scratch/esm/slurm-buildkite/climaatmos-ci/depot/default/packages/ClimaCore/wHC4M/src/DataLayouts/DataLayouts.jl:661
  | │││││││││││││││││││┌ (::Type{ClimaCore.DataLayouts.VF{_A, _B}} where {_A, _B})(array::Matrix) @ ClimaCore.DataLayouts /central/scratch/esm/slurm-buildkite/climaatmos-ci/depot/default/packages/ClimaCore/wHC4M/src/DataLayouts/DataLayouts.jl:649
  | ││││││││││││││││││││┌ (::Type{ClimaCore.DataLayouts.VF{_A, _B, Matrix{_A1}}} where {_A, _B, _A1})(array::Matrix) @ ClimaCore.DataLayouts /central/scratch/esm/slurm-buildkite/climaatmos-ci/depot/default/packages/ClimaCore/wHC4M/src/DataLayouts/DataLayouts.jl:642
  | │││││││││││││││││││││┌ convert(::Type{Array{NCDatasets.nc_vlen_t{T}, 2}}, data::Array{Vector{T}, 2}) where T @ NCDatasets /central/scratch/esm/slurm-buildkite/climaatmos-ci/depot/default/packages/NCDatasets/PkTUf/src/netcdf_c.jl:235
  | ││││││││││││││││││││││┌ iterate(e::Base.Iterators.Enumerate{I} where I<:(Array{Vector{T}, 2} where T)) @ Base.Iterators ./iterators.jl:205
  | │││││││││││││││││││││││┌ iterate(e::Base.Iterators.Enumerate{I} where I<:(Array{Vector{T}, 2} where T), state::Tuple{Int64}) @ Base.Iterators ./iterators.jl:206
  | ││││││││││││││││││││││││┌ iterate(A::Array{Vector{T}, 2} where T) @ Base ./array.jl:945
  | │││││││││││││││││││││││││┌ iterate(A::Array{Vector{T}, 2} where T, i::Int64) @ Base ./array.jl:945
  | ││││││││││││││││││││││││││ runtime dispatch detected: (A::Array{Vector{T}, 2} where T)[i::Int64]::Vector
  | │││││││││││││││││││││││││└────────────────────
  | │││││││││││││││││┌ _FiniteDifferenceGrid(topology::ClimaCore.Topologies.IntervalTopology) @ ClimaCore.Grids /central/scratch/esm/slurm-buildkite/climaatmos-ci/depot/default/packages/ClimaCore/wHC4M/src/Grids/finitedifference.jl:62
  | ││││││││││││││││││┌ vindex(v::Any) @ ClimaCore.DataLayouts /central/scratch/esm/slurm-buildkite/climaatmos-ci/depot/default/packages/ClimaCore/wHC4M/src/DataLayouts/DataLayouts.jl:1112
  | │││││││││││││││││││┌ CartesianIndex(::Int64, ::Int64, ::Int64, ::Integer, ::Int64) @ Base.IteratorsMD ./multidimensional.jl:79
  | ││││││││││││││││││││ runtime dispatch detected: CartesianIndex(index::Tuple{Int64, Int64, Int64, Integer, Int64})::CartesianIndex{5}
  | │││││││││││││││││││└────────────────────
  | │││││││││││││││││││┌ CartesianIndex(::Int64, ::Int64, ::Int64, ::Union{Integer, CartesianIndex}, ::Int64) @ Base.IteratorsMD ./multidimensional.jl:86
  | ││││││││││││││││││││┌ flatten(I::Tuple{Int64, Int64, Int64, Union{Integer, CartesianIndex}, Int64}) @ Base.IteratorsMD ./multidimensional.jl:89
  | │││││││││││││││││││││┌ flatten(I::Tuple{Int64, Int64, Any, Int64}) @ Base.IteratorsMD ./multidimensional.jl:89
  | ││││││││││││││││││││││┌ flatten(I::Tuple{Int64, Any, Int64}) @ Base.IteratorsMD ./multidimensional.jl:89
  | │││││││││││││││││││││││┌ flatten(I::Tuple{Any, Int64}) @ Base.IteratorsMD ./multidimensional.jl:89
  | ││││││││││││││││││││││││ runtime dispatch detected: Base.IteratorsMD.Tuple(%1::Any)::Tuple
  | │││││││││││││││││││││││└────────────────────
  | ││││││││││││││││││││┌ CartesianIndex(index::Tuple{Int64, Int64, Int64, Vararg{Union{Integer, CartesianIndex}}}) @ Base.IteratorsMD ./multidimensional.jl:90
  | │││││││││││││││││││││┌ CartesianIndex(::Int64, ::Int64, ::Int64, ::Vararg{Integer}) @ Base.IteratorsMD ./multidimensional.jl:79
  | ││││││││││││││││││││││ runtime dispatch detected: CartesianIndex(index::Tuple{Int64, Int64, Int64, Vararg{Integer}})::CartesianIndex
  | │││││││││││││││││││││└────────────────────
  | │││││││││││││││││││││┌ CartesianIndex(::Int64, ::Vararg{Union{Integer, CartesianIndex}}) @ Base.IteratorsMD ./multidimensional.jl:86
  | ││││││││││││││││││││││┌ flatten(I::Tuple{Int64, Vararg{Union{Integer, CartesianIndex}}}) @ Base.IteratorsMD ./multidimensional.jl:89
  | │││││││││││││││││││││││┌ flatten(I::Tuple{Vararg{Union{Integer, CartesianIndex}}}) @ Base.IteratorsMD ./multidimensional.jl:89
  | ││││││││││││││││││││││││ runtime dispatch detected: Base.IteratorsMD.Tuple(%1::Union{Integer, CartesianIndex})::Tuple
  | │││││││││││││││││││││││└────────────────────
  | │││││││││││││││││││││││┌ flatten(I::Tuple{Vararg{Union{Integer, CartesianIndex}}}) @ Base.IteratorsMD ./multidimensional.jl:89
  | ││││││││││││││││││││││││ runtime dispatch detected: Base.IteratorsMD.flatten(%14::Tuple{Vararg{Union{Integer, CartesianIndex}}})::Tuple
  | │││││││││││││││││││││││└────────────────────
  | ││││││││││││││││││││││┌ flatten(I::Tuple{Int64, Vararg{Union{Integer, CartesianIndex}}}) @ Base.IteratorsMD ./multidimensional.jl:89
  | │││││││││││││││││││││││ runtime dispatch detected: Base.IteratorsMD.flatten(%3::Tuple{Vararg{Union{Integer, CartesianIndex}}})::Tuple
  | ││││││││││││││││││││││└────────────────────
  | │││││││││││││││││││││┌ CartesianIndex(::Int64, ::Vararg{Union{Integer, CartesianIndex}}) @ Base.IteratorsMD ./multidimensional.jl:86
  | ││││││││││││││││││││││ failed to optimize due to recursion: CartesianIndex(::Int64, ::Vararg{Union{Integer, CartesianIndex}})
  | │││││││││││││││││││││└────────────────────
  | │││││││││││││││││││┌ CartesianIndex(::Int64, ::Int64, ::Int64, ::Union{Integer, CartesianIndex}, ::Int64) @ Base.IteratorsMD ./multidimensional.jl:86
  | ││││││││││││││││││││ failed to optimize due to recursion: CartesianIndex(::Int64, ::Int64, ::Int64, ::Union{Integer, CartesianIndex}, ::Int64)
  | │││││││││││││││││││└────────────────────
  | ││││││││││││││││││┌ vindex(v::Any) @ ClimaCore.DataLayouts /central/scratch/esm/slurm-buildkite/climaatmos-ci/depot/default/packages/ClimaCore/wHC4M/src/DataLayouts/DataLayouts.jl:1112
  | │││││││││││││││││││ runtime dispatch detected: ClimaCore.DataLayouts.CartesianIndex(1, 1, 1, v::Any, 1)::Any
  | ││││││││││││││││││└────────────────────
  | ││││││││││││││││││┌ setindex!(data::ClimaCore.DataLayouts.VF{_A, _B, Matrix{_A1}} where {_A, _B, _A1}, val::Any, I::CartesianIndex) @ ClimaCore.DataLayouts /central/scratch/esm/slurm-buildkite/climaatmos-ci/depot/default/packages/ClimaCore/wHC4M/src/DataLayouts/DataLayouts.jl:1164
  | │││││││││││││││││││┌ bounds_condition(data::ClimaCore.DataLayouts.VF{_A, _B, Matrix{_A1}} where {_A, _B, _A1}, I::CartesianIndex) @ ClimaCore.DataLayouts /central/scratch/esm/slurm-buildkite/climaatmos-ci/depot/default/packages/ClimaCore/wHC4M/src/DataLayouts/DataLayouts.jl:1009
  | ││││││││││││││││││││ runtime dispatch detected: (%2::Int64 ClimaCore.DataLayouts.:<= %7::Any)::Any
  | │││││││││││││││││││└────────────────────
  | ││││││││││││││││││┌ setindex!(data::ClimaCore.DataLayouts.VF{_A, _B, Matrix{_A1}} where {_A, _B, _A1}, val::Any, I::CartesianIndex) @ ClimaCore.DataLayouts /central/scratch/esm/slurm-buildkite/climaatmos-ci/depot/default/packages/ClimaCore/wHC4M/src/DataLayouts/DataLayouts.jl:1165
  | │││││││││││││││││││┌ set_struct!(array::Matrix{S}, val::Any, ::Val{2}, index::CartesianIndex{2}) where S @ ClimaCore.DataLayouts /central/scratch/esm/slurm-buildkite/climaatmos-ci/depot/default/packages/ClimaCore/wHC4M/src/DataLayouts/struct.jl:257
  | ││││││││││││││││││││┌ setindex!(::Matrix, ::Any, ::CartesianIndex{2}) @ Base ./multidimensional.jl:698
  | │││││││││││││││││││││┌ to_indices(A::Matrix, I::Tuple{CartesianIndex{2}}) @ Base ./multidimensional.jl:859
  | ││││││││││││││││││││││┌ to_indices(A::Matrix, inds::Tuple{}, I::Tuple{CartesianIndex{2}}) @ Base ./indices.jl:354
  | │││││││││││││││││││││││ runtime dispatch detected: Base._to_indices1(A::Matrix, inds::Tuple{}, %1::CartesianIndex{2})::Tuple{Int64, Int64}
  | ││││││││││││││││││││││└────────────────────
 ```

I had to truncate the message, since it was too long. I'm not sure how many of these are directly the result of the issue that is fixed in https://github.com/CliMA/ClimaCore.jl/pull/2004, but I can retry https://buildkite.com/clima/climaatmos-ci/builds/20805#01922a42-bcbe-49b0-ae2d-2247cb3b2a49 once that is merged/released.

[charleskawczynski]

I noticed that one of the failures in https://buildkite.com/clima/climaatmos-ci/builds/20805#01922a42-bcbe-49b0-ae2d-2247cb3b2a49 was an inference failure in

ClimaDiagnostics.jl/src/clima_diagnostics.jl

Lines 252 to 253 in c1cc7c3

	sciml_callback(integrator) =
	orchestrate_diagnostics(integrator, diagnostics_handler)

, which looks like it captures diagnostics_handler.

[charleskawczynski]

I re-ran this with some package updates (and with more type info) and it seems like the main issue is with the out parameter:

  ││││││││││││┌ compute_orog!(out::Any, state::ClimaCore.Fields.FieldVector{Float32, @NamedTuple{c::ClimaCore.Fields.Field{…}, f::ClimaCore.Fields.Field{…}}}, cache::ClimaAtmos.AtmosCache{Float32, Float32, ClimaAtmos.WallTimeEstimate, Dates.DateTime, ClimaAtmos.AtmosModel{ClimaAtmos.SingleColumnModel, ClimaAtmos.DryModel, ClimaAtmos.NoPrecipitation, ClimaAtmos.QuadratureCloud, Nothing, Nothing, Nothing, ClimaAtmos.RRTMGPInterface.GrayRadiation, Nothing, Nothing, Nothing, Bool, ClimaAtmos.UseAllTendency, ClimaAtmos.EDMFXModel{…}, Nothing, Nothing, Nothing, Nothing, Nothing, ClimaAtmos.Explicit, ClimaAtmos.Explicit, Nothing, Nothing, ClimaAtmos.ZonallySymmetricSST, ClimaAtmos.IdealizedInsolation, ClimaAtmos.PrescribedSurfaceTemperature, ClimaAtmos.ConstantAlbedo{…}, ClimaAtmos.AtmosNumerics{…}}, @NamedTuple{limiter::Nothing}, ClimaAtmos.Parameters.ClimaAtmosParameters{Float32, Thermodynamics.Parameters.ThermodynamicsParameters{…}, RRTMGP.Parameters.RRTMGPParameters{…}, Insolation.Parameters.InsolationParameters{…}, Nothing, Nothing, CloudMicrophysics.Parameters.WaterProperties{…}, SurfaceFluxes.Parameters.SurfaceFluxesParameters{…}, ClimaAtmos.Parameters.TurbulenceConvectionParameters{…}, ClimaAtmos.Parameters.SurfaceTemperatureParameters{…}}, @NamedTuple{ᶜΦ::ClimaCore.Fields.Field{…}, ᶠgradᵥ_ᶜΦ::ClimaCore.Fields.Field{…}, ᶜgradᵥ_ᶠΦ::ClimaCore.Fields.Field{…}, ᶜf³::ClimaCore.Fields.Field{…}, ᶠf¹²::Nothing, surface_ct3_unit::ClimaCore.Fields.Field{…}}, ClimaAtmos.SurfaceConditions.SurfaceState{Float32, ClimaAtmos.SurfaceConditions.ExchangeCoefficients{…}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, @NamedTuple{}, ClimaAtmos.SGSQuadrature{3, ClimaAtmos.GaussianQuad, StaticArraysCore.SVector{…}, StaticArraysCore.SVector{…}}, @NamedTuple{ᶜspecific::ClimaCore.Fields.Field{…}, ᶜu::ClimaCore.Fields.Field{…}, ᶠu³::ClimaCore.Fields.Field{…}, ᶜK::ClimaCore.Fields.Field{…}, ᶜts::ClimaCore.Fields.Field{…}, ᶜp::ClimaCore.Fields.Field{…}, ᶜh_tot::ClimaCore.Fields.Field{…}, ᶜmixing_length::ClimaCore.Fields.Field{…}, ᶜlinear_buoygrad::ClimaCore.Fields.Field{…}, ᶜstrain_rate_norm::ClimaCore.Fields.Field{…}, sfc_conditions::ClimaCore.Fields.Field{…}, ᶜgradᵥ_θ_virt::ClimaCore.Fields.Field{…}, ᶜgradᵥ_q_tot::ClimaCore.Fields.Field{…}, ᶜgradᵥ_θ_liq_ice::ClimaCore.Fields.Field{…}, cloud_diagnostics_tuple::ClimaCore.Fields.Field{…}}, @NamedTuple{ᶠtemp_scalar::ClimaCore.Fields.Field{…}, ᶜtemp_scalar::ClimaCore.Fields.Field{…}, ᶜtemp_scalar_2::ClimaCore.Fields.Field{…}, ᶜtemp_scalar_3::ClimaCore.Fields.Field{…}, ᶜtemp_scalar_4::ClimaCore.Fields.Field{…}, ᶜtemp_scalar_5::ClimaCore.Fields.Field{…}, ᶠtemp_field_level::ClimaCore.Fields.Field{…}, temp_field_level::ClimaCore.Fields.Field{…}, temp_field_level_2::ClimaCore.Fields.Field{…}, temp_field_level_3::ClimaCore.Fields.Field{…}, temp_data::ClimaCore.DataLayouts.VIJFH{…}, temp_data_face_level::ClimaCore.DataLayouts.IJFH{…}, temp_data_level::ClimaCore.DataLayouts.IJFH{…}, temp_data_level_2::ClimaCore.DataLayouts.IJFH{…}, temp_data_level_3::ClimaCore.DataLayouts.IJFH{…}, ᶜtemp_C12::ClimaCore.Fields.Field{…}, ᶜtemp_C3::ClimaCore.Fields.Field{…}, ᶜtemp_CT3::ClimaCore.Fields.Field{…}, ᶜtemp_CT123::ClimaCore.Fields.Field{…}, ᶠtemp_CT3::ClimaCore.Fields.Field{…}, ᶠtemp_CT12::ClimaCore.Fields.Field{…}, ᶠtemp_CT12ʲs::ClimaCore.Fields.Field{…}, ᶠtemp_C123::ClimaCore.Fields.Field{…}, ᶜtemp_UVWxUVW::ClimaCore.Fields.Field{…}, ᶠtemp_UVWxUVW::ClimaCore.Fields.Field{…}, sfc_temp_C3::ClimaCore.Fields.Field{…}, ∂ᶜK_∂ᶜuₕ::ClimaCore.Fields.Field{…}, ∂ᶜK_∂ᶠu₃::ClimaCore.Fields.Field{…}, ᶠp_grad_matrix::ClimaCore.Fields.Field{…}, ᶠbidiagonal_matrix_ct3::ClimaCore.Fields.Field{…}, ᶠbidiagonal_matrix_ct3_2::ClimaCore.Fields.Field{…}, ᶜadvection_matrix::ClimaCore.Fields.Field{…}, ᶜdiffusion_h_matrix::ClimaCore.Fields.Field{…}, ᶜdiffusion_u_matrix::ClimaCore.Fields.Field{…}, ᶠtridiagonal_matrix_c3::ClimaCore.Fields.Field{…}}, @NamedTuple{}, Bool, @NamedTuple{}, @NamedTuple{}, @NamedTuple{surface_rain_flux::ClimaCore.Fields.Field{…}, surface_snow_flux::ClimaCore.Fields.Field{…}}, @NamedTuple{}, @NamedTuple{}, @NamedTuple{}, @NamedTuple{}, @NamedTuple{}, @NamedTuple{}, @NamedTuple{}, @NamedTuple{orbital_data::Insolation.OrbitalData{…}, rrtmgp_model::ClimaAtmos.RRTMGPInterface.RRTMGPModel{…}, ᶠradiation_flux::ClimaCore.Fields.Field{…}}, @NamedTuple{}, Vector{ClimaCore.Geometry.WVector{…}}, Vector{ClimaCore.Geometry.WVector{…}}, @NamedTuple{col_integrated_precip_energy_tendency::@NamedTuple{}}, String}, time::Float32, hypsography::ClimaCore.Grids.Flat) @ ClimaAtmos.Diagnostics /dev/CliMA/ClimaAtmos.jl/src/diagnostics/core_diagnostics.jl:792

So, we may be able to fix this when we move to using LazyBroadcast. I can't imagine that we don't have the same issue with the function, too.

[charleskawczynski]

This is similar to the manual dispatch example:

using BenchmarkTools

func(x::Int) = 1;
func(x::Float64) = 1;

# The type of the arrays is `Vector{Any}`, so that the compiler cannot choose the
# required method for `func` at compile time.
const x_int = Any[1 for i in 1:1000];
const x_float = Any[1.0 for i in 1:1000];
const x_mixed = Any[iseven(i) ? 1 : 1.0 for i in 1:1000];

function ex_union_split(x)
    for y in x
        if y isa Int
            func(y)
        elseif y isa Float64
            func(y)
        end
    end
end

function ex_runtime_dispatch(x)
    for y in x
        func(y)
    end
end

function main(x_int,x_float,x_mixed)
    print("Manual  union split with array of Int:     "); @btime ex_union_split(x_int)
    print("Runtime dispatch    with array of Int:     "); @btime ex_runtime_dispatch(x_int)
    print("Manual  union split with array of Float64: "); @btime ex_union_split(x_float)
    print("Runtime dispatch    with array of Float64: "); @btime ex_runtime_dispatch(x_float)
    print("Manual  union split with mixed array:      "); @btime ex_union_split(x_mixed)
    print("Runtime dispatch    with mixed array:      "); @btime ex_runtime_dispatch(x_mixed)
end
main(x_int,x_float,x_mixed)

[charleskawczynski]

A function-analogy might be:

using BenchmarkTools

foo(x::Int) = 1;
bar(x::Float64) = 1;

# The type of the arrays is `Vector{Any}`, so that the compiler cannot choose the
# required method for `func` at compile time.
const x_int = Any[foo for i in 1:1000];
const x_float = Any[bar for i in 1:1000];
const x_mixed = Any[iseven(i) ? foo : bar for i in 1:1000];

function ex_union_split(x::X) where {X}
    for f in x
        ft = typeof(f)
        if ft <: typeof(foo)
            z = f(1)
        elseif ft <: typeof(bar)
            z = f(1.0)
        end
    end
end

function ex_runtime_dispatch(x)
    for f in x
        a = typeof(f) <: typeof(foo) ? 1 : 1.0
        f(a)
    end
end

# function main(x_int,x_float,x_mixed)
#     ex_union_split(x_int)
#     ex_runtime_dispatch(x_int)
#     ex_union_split(x_float)
#     ex_runtime_dispatch(x_float)
#     ex_union_split(x_mixed)
#     ex_runtime_dispatch(x_mixed)

#     print("Manual  union split with array of Int:     "); @btime ex_union_split(x_int)
#     print("Runtime dispatch    with array of Int:     "); @btime ex_runtime_dispatch(x_int)
#     print("Manual  union split with array of Float64: "); @btime ex_union_split(x_float)
#     print("Runtime dispatch    with array of Float64: "); @btime ex_runtime_dispatch(x_float)
#     print("Manual  union split with mixed array:      "); @btime ex_union_split(x_mixed)
#     print("Runtime dispatch    with mixed array:      "); @btime ex_runtime_dispatch(x_mixed)
# end
# main(x_int,x_float,x_mixed)

using JET
@test_opt ex_union_split(x_int)
@test_opt ex_union_split(x_float)
@test_opt ex_union_split(x_mixed)

But this didn't seem to work, I had to put the function into a wrapper object for it to work:

using BenchmarkTools

foo(x::Int) = 1;
bar(x::Float64) = 1;
struct FunctionWrapper{F}
    f::F
end
@inline (fw::FunctionWrapper)(args...; kwargs...) = fw.f(args...; kwargs...)
# The type of the arrays is `Vector{Any}`, so that the compiler cannot choose the
# required method for `func` at compile time.
const x_int = Any[FunctionWrapper(foo) for i in 1:1000];
const x_float = Any[FunctionWrapper(bar) for i in 1:1000];
const x_mixed = Any[iseven(i) ? FunctionWrapper(foo) : FunctionWrapper(bar) for i in 1:1000];

function ex_union_split(x::X) where {X}
    for f in x
        if f isa FunctionWrapper{typeof(foo)}
            z = f(1)
        elseif f isa FunctionWrapper{typeof(bar)}
            z = f(1.0)
        end
    end
end

function ex_runtime_dispatch(x)
    for f in x
        a = f isa FunctionWrapper{typeof(foo)} ? 1 : 1.0
        f(a)
    end
end

function main(x_int,x_float,x_mixed)
    ex_union_split(x_int)
    ex_runtime_dispatch(x_int)
    ex_union_split(x_float)
    ex_runtime_dispatch(x_float)
    ex_union_split(x_mixed)
    ex_runtime_dispatch(x_mixed)

    print("Manual  union split with array of Int:     "); @btime ex_union_split(x_int)
    print("Runtime dispatch    with array of Int:     "); @btime ex_runtime_dispatch(x_int)
    print("Manual  union split with array of Float64: "); @btime ex_union_split(x_float)
    print("Runtime dispatch    with array of Float64: "); @btime ex_runtime_dispatch(x_float)
    print("Manual  union split with mixed array:      "); @btime ex_union_split(x_mixed)
    print("Runtime dispatch    with mixed array:      "); @btime ex_runtime_dispatch(x_mixed)
end
main(x_int,x_float,x_mixed)