From 1134a39d8a21e1467a8217b34e5fa3fd49efdbbc Mon Sep 17 00:00:00 2001 From: sriharshakandala Date: Tue, 11 Mar 2025 14:43:30 -0700 Subject: [PATCH 1/6] Try RRTMGP branch. --- .buildkite/pipeline.yml | 1 + .../radiation/RRTMGPInterface.jl | 56 +++++++++++++------ 2 files changed, 41 insertions(+), 16 deletions(-) diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 7e21bfde97..edf7879854 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -30,6 +30,7 @@ steps: - echo "--- Instantiate .buildkite" - "julia --project=.buildkite -e 'using Pkg; Pkg.instantiate(;verbose=true)'" + - "julia --project=.buildkite -e 'using Pkg; Pkg.add(Pkg.PackageSpec(;name=\"RRTMGP\", rev=\"sk/refactor_fluxes_change_datalayout\"))'" - "julia --project=.buildkite -e 'using Pkg; Pkg.precompile()'" - "julia --project=.buildkite -e 'using CUDA; CUDA.precompile_runtime()'" - "julia --project=.buildkite -e 'using Pkg; Pkg.status()'" diff --git a/src/parameterized_tendencies/radiation/RRTMGPInterface.jl b/src/parameterized_tendencies/radiation/RRTMGPInterface.jl index e7925dc03f..3d1f0f9a06 100644 --- a/src/parameterized_tendencies/radiation/RRTMGPInterface.jl +++ b/src/parameterized_tendencies/radiation/RRTMGPInterface.jl @@ -529,14 +529,22 @@ function RRTMGPModel( fluxb_lw = radiation_mode isa GrayRadiation ? nothing : RRTMGP.Fluxes.FluxLW(ncol, nlay, FT, DA) - set_and_save!(flux_lw.flux_up, "face_lw_flux_up", t...) - set_and_save!(flux_lw.flux_dn, "face_lw_flux_dn", t...) - set_and_save!(flux_lw.flux_net, "face_lw_flux", t...) + set_and_save!(transpose(flux_lw.flux_up), "face_lw_flux_up", t...) + set_and_save!(transpose(flux_lw.flux_dn), "face_lw_flux_dn", t...) + set_and_save!(transpose(flux_lw.flux_net), "face_lw_flux", t...) if radiation_mode isa AllSkyRadiationWithClearSkyDiagnostics flux_lw2 = RRTMGP.Fluxes.FluxLW(ncol, nlay, FT, DA) - set_and_save!(flux_lw2.flux_up, "face_clear_lw_flux_up", t...) - set_and_save!(flux_lw2.flux_dn, "face_clear_lw_flux_dn", t...) - set_and_save!(flux_lw2.flux_net, "face_clear_lw_flux", t...) + set_and_save!( + transpose(flux_lw2.flux_up), + "face_clear_lw_flux_up", + t..., + ) + set_and_save!( + transpose(flux_lw2.flux_dn), + "face_clear_lw_flux_dn", + t..., + ) + set_and_save!(transpose(flux_lw2.flux_net), "face_clear_lw_flux", t...) end sfc_emis = DA{FT}(undef, nbnd_lw, ncol) @@ -592,20 +600,32 @@ function RRTMGPModel( fluxb_sw = radiation_mode isa GrayRadiation ? nothing : RRTMGP.Fluxes.FluxSW(ncol, nlay, FT, DA) - set_and_save!(flux_sw.flux_up, "face_sw_flux_up", t...) - set_and_save!(flux_sw.flux_dn, "face_sw_flux_dn", t...) - set_and_save!(flux_sw.flux_net, "face_sw_flux", t...) - set_and_save!(flux_sw.flux_dn_dir, "face_sw_direct_flux_dn", t...) + set_and_save!(transpose(flux_sw.flux_up), "face_sw_flux_up", t...) + set_and_save!(transpose(flux_sw.flux_dn), "face_sw_flux_dn", t...) + set_and_save!(transpose(flux_sw.flux_net), "face_sw_flux", t...) + set_and_save!( + transpose(flux_sw.flux_dn_dir), + "face_sw_direct_flux_dn", + t..., + ) if radiation_mode isa AllSkyRadiationWithClearSkyDiagnostics flux_sw2 = RRTMGP.Fluxes.FluxSW(ncol, nlay, FT, DA) - set_and_save!(flux_sw2.flux_up, "face_clear_sw_flux_up", t...) - set_and_save!(flux_sw2.flux_dn, "face_clear_sw_flux_dn", t...) set_and_save!( - flux_sw2.flux_dn_dir, + transpose(flux_sw2.flux_up), + "face_clear_sw_flux_up", + t..., + ) + set_and_save!( + transpose(flux_sw2.flux_dn), + "face_clear_sw_flux_dn", + t..., + ) + set_and_save!( + transpose(flux_sw2.flux_dn_dir), "face_clear_sw_direct_flux_dn", t..., ) - set_and_save!(flux_sw2.flux_net, "face_clear_sw_flux", t...) + set_and_save!(transpose(flux_sw2.flux_net), "face_clear_sw_flux", t...) end cos_zenith = DA{FT}(undef, ncol) @@ -633,9 +653,13 @@ function RRTMGPModel( sfc_alb_diffuse, ) - set_and_save!(similar(flux_lw.flux_net), "face_flux", t...) + set_and_save!(similar(transpose(flux_lw.flux_net)), "face_flux", t...) if radiation_mode isa AllSkyRadiationWithClearSkyDiagnostics - set_and_save!(similar(flux_lw2.flux_net), "face_clear_flux", t...) + set_and_save!( + similar(transpose(flux_lw2.flux_net)), + "face_clear_flux", + t..., + ) end if !(radiation_mode isa GrayRadiation) @assert RRTMGP.LookUpTables.get_n_gases(lookup_lw) == From 82b58515f9fedf3f875f5822d145d1e953ab1b46 Mon Sep 17 00:00:00 2001 From: sriharshakandala Date: Wed, 12 Mar 2025 11:11:16 -0700 Subject: [PATCH 2/6] Update ClimaCore branch in CI pipeline --- .buildkite/pipeline.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index edf7879854..67cacda1d4 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -30,6 +30,7 @@ steps: - echo "--- Instantiate .buildkite" - "julia --project=.buildkite -e 'using Pkg; Pkg.instantiate(;verbose=true)'" + - "julia --project=.buildkite -e 'using Pkg; Pkg.add(Pkg.PackageSpec(;name=\"ClimaCore\", rev=\"sk/update_for_RRTMGP_branch\"))'" - "julia --project=.buildkite -e 'using Pkg; Pkg.add(Pkg.PackageSpec(;name=\"RRTMGP\", rev=\"sk/refactor_fluxes_change_datalayout\"))'" - "julia --project=.buildkite -e 'using Pkg; Pkg.precompile()'" - "julia --project=.buildkite -e 'using CUDA; CUDA.precompile_runtime()'" From 5242739f0a5f4f29cdc8f08a95ef2952eb3db969 Mon Sep 17 00:00:00 2001 From: sriharshakandala Date: Wed, 12 Mar 2025 12:01:40 -0700 Subject: [PATCH 3/6] Add missing `parent_array_type` method. --- src/parameterized_tendencies/radiation/RRTMGPInterface.jl | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/parameterized_tendencies/radiation/RRTMGPInterface.jl b/src/parameterized_tendencies/radiation/RRTMGPInterface.jl index 3d1f0f9a06..e2716e936a 100644 --- a/src/parameterized_tendencies/radiation/RRTMGPInterface.jl +++ b/src/parameterized_tendencies/radiation/RRTMGPInterface.jl @@ -964,6 +964,11 @@ function set_array!(array, value::AbstractArray{<:Real}, symbol) end end +import LinearAlgebra +import ClimaCore.DataLayouts: parent_array_type +parent_array_type(::Type{<:LinearAlgebra.Transpose{T, P}}) where {T, P} = + parent_array_type(P) + function set_and_save!(array, name, views, domain_nlay, dict = nothing) domain_symbol = Symbol(name) From a44f3fdea423abc66f92be8d622ac523825fd0c1 Mon Sep 17 00:00:00 2001 From: sriharshakandala Date: Wed, 12 Mar 2025 13:25:45 -0700 Subject: [PATCH 4/6] Rm ClimaCore branch change --- .buildkite/pipeline.yml | 1 - 1 file changed, 1 deletion(-) diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 67cacda1d4..edf7879854 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -30,7 +30,6 @@ steps: - echo "--- Instantiate .buildkite" - "julia --project=.buildkite -e 'using Pkg; Pkg.instantiate(;verbose=true)'" - - "julia --project=.buildkite -e 'using Pkg; Pkg.add(Pkg.PackageSpec(;name=\"ClimaCore\", rev=\"sk/update_for_RRTMGP_branch\"))'" - "julia --project=.buildkite -e 'using Pkg; Pkg.add(Pkg.PackageSpec(;name=\"RRTMGP\", rev=\"sk/refactor_fluxes_change_datalayout\"))'" - "julia --project=.buildkite -e 'using Pkg; Pkg.precompile()'" - "julia --project=.buildkite -e 'using CUDA; CUDA.precompile_runtime()'" From 8b460f204bfdedd02caf667aa3a9ae102c1e5164 Mon Sep 17 00:00:00 2001 From: sriharshakandala Date: Wed, 12 Mar 2025 15:01:21 -0700 Subject: [PATCH 5/6] Add additional method for `safe_fill!` --- .../radiation/RRTMGPInterface.jl | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/src/parameterized_tendencies/radiation/RRTMGPInterface.jl b/src/parameterized_tendencies/radiation/RRTMGPInterface.jl index e2716e936a..428a29e55c 100644 --- a/src/parameterized_tendencies/radiation/RRTMGPInterface.jl +++ b/src/parameterized_tendencies/radiation/RRTMGPInterface.jl @@ -934,9 +934,17 @@ function RRTMGPModel( ) end +import LinearAlgebra +import ClimaCore.DataLayouts: parent_array_type +parent_array_type(::Type{<:LinearAlgebra.Transpose{T, P}}) where {T, P} = + parent_array_type(P) + +safe_fill!(array::LinearAlgebra.Transpose, value) = + fill!(transpose(array), value) +safe_fill!(array, value) = fill!(array, value) # This sets `array .= value`, but it allows `array` to be to be a `CuArray` # while `value` is an `Array` (in which case broadcasting throws an error). -set_array!(array, value::Real, symbol) = fill!(array, value) +set_array!(array, value::Real, symbol) = safe_fill!(array, value) function set_array!(array, value::AbstractArray{<:Real}, symbol) if ndims(array) == 2 if size(value) == size(array) @@ -947,7 +955,7 @@ function set_array!(array, value::AbstractArray{<:Real}, symbol) end elseif size(value) == (1, size(array, 2)) for (icol, col) in enumerate(eachcol(array)) - fill!(col, value[1, icol]) + safe_fill!(col, value[1, icol]) end else error("expected $symbol to be an array of size $(size(array)), \ @@ -964,11 +972,6 @@ function set_array!(array, value::AbstractArray{<:Real}, symbol) end end -import LinearAlgebra -import ClimaCore.DataLayouts: parent_array_type -parent_array_type(::Type{<:LinearAlgebra.Transpose{T, P}}) where {T, P} = - parent_array_type(P) - function set_and_save!(array, name, views, domain_nlay, dict = nothing) domain_symbol = Symbol(name) From a86efac80e27f4e4d62cab67fc0ece10d1de5593 Mon Sep 17 00:00:00 2001 From: sriharshakandala Date: Wed, 12 Mar 2025 17:17:25 -0700 Subject: [PATCH 6/6] Add `safe_fill` method for subarray. --- src/parameterized_tendencies/radiation/RRTMGPInterface.jl | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/src/parameterized_tendencies/radiation/RRTMGPInterface.jl b/src/parameterized_tendencies/radiation/RRTMGPInterface.jl index 428a29e55c..f86148e5b3 100644 --- a/src/parameterized_tendencies/radiation/RRTMGPInterface.jl +++ b/src/parameterized_tendencies/radiation/RRTMGPInterface.jl @@ -941,6 +941,14 @@ parent_array_type(::Type{<:LinearAlgebra.Transpose{T, P}}) where {T, P} = safe_fill!(array::LinearAlgebra.Transpose, value) = fill!(transpose(array), value) +function safe_fill!( + array::SubArray{T, 2, <:LinearAlgebra.Transpose}, + value, +) where {T} + subarray = view(transpose(parent(array)), reverse(array.indices)...) + fill!(subarray, value) + return nothing +end safe_fill!(array, value) = fill!(array, value) # This sets `array .= value`, but it allows `array` to be to be a `CuArray` # while `value` is an `Array` (in which case broadcasting throws an error).