diff --git a/perf/jet.jl b/perf/jet.jl index 85fd9fe3f9..021395f528 100644 --- a/perf/jet.jl +++ b/perf/jet.jl @@ -1,3 +1,12 @@ +#= +# to get a target configuration: +julia --project=perf +import ClimaAtmos as CA +using Revise; include(joinpath(pkgdir(CA), "perf", "common.jl")) +config = TargetJobConfig("gpu_prognostic_edmfx_aquaplanet") + +include(joinpath(pkgdir(CA), "perf", "jet.jl")) +=# redirect_stderr(IOContext(stderr, :stacktrace_types_limited => Ref(false))) import Random import HDF5, NCDatasets, CUDA @@ -6,16 +15,31 @@ import ClimaAtmos as CA include("common.jl") -config_dict = Dict("z_elem" => 63, "dt" => "10secs", "t_end" => "3600secs") -config = AtmosCoveragePerfConfig(config_dict) +if !(@isdefined config) + config_dict = Dict("z_elem" => 63, "dt" => "10secs", "t_end" => "3600secs") + config = AtmosCoveragePerfConfig(config_dict) +end + +import JET + +function jet_test(f, args) + f(args...) # compile first + JET.@test_opt ignored_modules = (HDF5, CUDA, NCDatasets) f(args...) +end simulation = CA.get_simulation(config) (; integrator) = simulation -import JET - import SciMLBase -SciMLBase.step!(integrator) # Make sure no errors -JET.@test_opt ignored_modules = (HDF5, CUDA, NCDatasets) SciMLBase.step!( - integrator, -) +W = get_W(integrator); +(; u, p, dt, t) = integrator; + +# jet_test(wfact_fun(integrator), (W, u, p, dt, t)) +# jet_test(LA.ldiv!, (X, W, u)) +# jet_test(implicit_fun(integrator), implicit_args(integrator)) +# jet_test(remaining_fun(integrator), remaining_args(integrator)) +# jet_test(CA.additional_tendency!, (X, u, p, t)) +# jet_test(CA.hyperdiffusion_tendency!, remaining_args(integrator)) +# jet_test(CA.dss!, (u, p, t)) +# jet_test(CA.set_precomputed_quantities!, (u, p, t)) +jet_test(SciMLBase.step!, (integrator,))