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

@@ -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]

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)