From 10bba68e539846590be0d8ba2fc701e698f62692 Mon Sep 17 00:00:00 2001 From: sriharshakandala Date: Wed, 25 Oct 2023 15:38:31 -0700 Subject: [PATCH] This PR makes the following temporary changes: Use Float32 for DYAMOND configuration. This helps with kernel metadata size issues. Reduce h_elems to 30. This helps with the memory footprint on the A100 GPU. Compute pressure2ozone interpolation during initialization on the CPU for GPU runs. These temporary changes will help continue finding and debugging potential downstream issues on a single GPU process while fixing known memory footprint and efficiency issues in parallel. The changes can be reverted once appropriate fixes are merged. --- config/gpu_configs/gpu_aquaplanet_dyamond.yml | 4 ++-- .../radiation/radiation.jl | 14 ++++++++++++-- 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/config/gpu_configs/gpu_aquaplanet_dyamond.yml b/config/gpu_configs/gpu_aquaplanet_dyamond.yml index 6c485f00eb..58a96b2999 100644 --- a/config/gpu_configs/gpu_aquaplanet_dyamond.yml +++ b/config/gpu_configs/gpu_aquaplanet_dyamond.yml @@ -1,4 +1,4 @@ -h_elem: 60 +h_elem: 30 z_max: 55000.0 z_elem: 63 dz_bottom: 30.0 @@ -15,6 +15,6 @@ rayleigh_sponge: true dt_save_to_disk: "3hours" dt: "50secs" t_end: "1days" -FLOAT_TYPE: "Float64" +FLOAT_TYPE: "Float32" job_id: "gpu_aquaplanet_dyamond" toml: [toml/longrun_aquaplanet_dyamond.toml] diff --git a/src/parameterized_tendencies/radiation/radiation.jl b/src/parameterized_tendencies/radiation/radiation.jl index 0759b087e1..8e5eeff9fc 100644 --- a/src/parameterized_tendencies/radiation/radiation.jl +++ b/src/parameterized_tendencies/radiation/radiation.jl @@ -35,8 +35,10 @@ function radiation_model_cache( data_loader, ) context = ClimaComms.context(axes(Y.c)) + device = context.device (; idealized_h2o, idealized_insolation, idealized_clouds) = radiation_mode FT = Spaces.undertype(axes(Y.c)) + DA = ClimaComms.array_type(device){FT} rrtmgp_params = CAP.rrtmgp_params(params) if idealized_h2o && radiation_mode isa RRTMGPI.GrayRadiation error("idealized_h2o can't be used with $radiation_mode") @@ -82,8 +84,16 @@ function radiation_model_cache( input_center_pressure, input_center_volume_mixing_ratio_o3, ) - ᶜvolume_mixing_ratio_o3_field = - @. FT(pressure2ozone(default_cache.ᶜp)) + if device isa ClimaComms.CUDADevice + fv = Fields.field_values(default_cache.ᶜp) + fld_array = DA(pressure2ozone.(Array(parent(fv)))) + data = DataLayouts.rebuild(fv, fld_array) + ᶜvolume_mixing_ratio_o3_field = + Fields.Field(data, axes(default_cache.ᶜp)) + else + ᶜvolume_mixing_ratio_o3_field = + @. FT(pressure2ozone(default_cache.ᶜp)) + end center_volume_mixing_ratio_o3 = RRTMGPI.field2array(ᶜvolume_mixing_ratio_o3_field)