diff --git a/src/soca/Fields/soca_fields_mod.F90 b/src/soca/Fields/soca_fields_mod.F90 index 79360a94..9e41c360 100644 --- a/src/soca/Fields/soca_fields_mod.F90 +++ b/src/soca/Fields/soca_fields_mod.F90 @@ -18,6 +18,7 @@ module soca_fields_mod datetime_create, datetime_diff use duration_mod, only: duration, duration_to_string use fckit_configuration_module, only: fckit_configuration +use iso_c_binding, only: c_ptr use logger_mod use kinds, only: kind_real use oops_variables_mod, only: oops_variables @@ -29,7 +30,11 @@ module soca_fields_mod ! SOCA I/O use soca_io_mod, only : soca_io_reader, soca_io_writer, & - soca_io_file_exists, soca_io_var_exists + soca_io_file_exists, soca_io_var_exists, & + soca_io_writers_commit_ensemble, & + soca_io_readers_commit_ensemble, & + soca_io_ensemble_root_pe, & + soca_io_ensemble_batch_size ! SOCA modules use soca_fields_metadata_mod, only : soca_field_metadata @@ -40,6 +45,9 @@ module soca_fields_mod implicit none private +public :: soca_fields_write_ensemble +public :: soca_fields_read_ensemble + ! ------------------------------------------------------------------------------ ! ------------------------------------------------------------------------------ @@ -164,6 +172,49 @@ module soca_fields_mod real(kind=kind_real), allocatable :: data(:,:,:) end type varwrapper +! ------------------------------------------------------------------------------ +! Polymorphic pointer wrapper so soca_fields_*_ensemble can take heterogeneous +! arrays of soca_state and soca_increment (both extend soca_fields). Used only +! by the bulk I/O entrypoints; built on the Fortran side from arrays of C++ +! registry keys. +type, public :: soca_fields_ptr_t + class(soca_fields), pointer :: p => null() +end type soca_fields_ptr_t + +! ------------------------------------------------------------------------------ +! One writer's worth of varwrapper state in the ensemble write path. Per-writer +! arrays are nested under this type so a single 1-D allocatable indexed by +! "writer index" carries all the per-(member, domain) vars buffers. +type :: ensemble_vars_t + type(varwrapper), allocatable :: vars(:) +end type ensemble_vars_t + +! ------------------------------------------------------------------------------ +! One member's worth of read state, carried across the soca_fields_read split +! into prepare (config + queue) and finalize (copy-to-atlas + post-processing). +! For the ensemble path, an array of plans is built up across members so that +! all per-domain readers can be flushed in one bulk soca_io_readers_commit_ensemble +! call, then each member's plan is finalized to drive its post-processing. +type :: read_plan_t + integer :: iread = 0 + ! per-active-domain reader + scratch buffers. readers(ri) reads into + ! domain_vars(ri)%vars(n)%data; same ri indexes both. Size 0 when + ! iread is not 1 or 3 (the no-read paths). + type(soca_io_reader), allocatable :: readers(:) + type(ensemble_vars_t), allocatable :: domain_vars(:) + ! vertical remap target (allocatable so its absence flags "no remap"). + ! target-ness comes from declaring plan instances with the target attribute. + real(kind=kind_real), allocatable :: h_common(:,:,:) + ! CICE category-vars deferred-read state (handled in finalize via read_seaice). + ! seaice_categories_vars is only initialized (oops_variables() empty_ctor) + ! when iread==1 or 3; finalize matches that guard before destructing. + type(oops_variables) :: seaice_categories_vars + character(len=:), allocatable :: ice_filename + ! Probe results that drive post-processing + logical :: compute_icethickness = .false. + logical :: compute_snowthickness = .false. +end type read_plan_t + ! ------------------------------------------------------------------------------ ! ------------------------------------------------------------------------------ @@ -468,307 +519,395 @@ subroutine soca_fields_read(self, f_conf, vdate) type(fckit_configuration), intent(in) :: f_conf type(datetime), intent(inout) :: vdate - character(len=:), allocatable :: str, basename, filename - integer :: iread = 0 - real(kind=kind_real), allocatable, target :: h_common(:,:,:) !< layer thickness to remap to (target for soca_io_reader pointer association) - type(soca_io_reader) :: reader - integer :: d, f, i, j, k, n, idx - type(remapping_CS) :: remapCS - type(oops_variables) :: seaice_categories_vars - type(varwrapper), allocatable, target :: vars(:) ! target so vars(n)%data sections are valid pointer targets for soca_io enqueue - type(atlas_field) :: afield1, afield2, afield3, afield4 - real(kind=kind_real), pointer :: adata1(:,:), adata2(:,:), adata3(:,:), adata4(:,:) - real(kind=kind_real), allocatable :: h_common_ij(:), hocn_ij(:), varocn_ij(:), varocn2_ij(:) - logical :: compute_icethickness, compute_snowthickness, remap_this_point + type(read_plan_t), target :: plan + integer :: r - character(len=3), dimension(5) :: domains + call soca_fields_read_prepare(self, f_conf, vdate, plan) + + ! Drain the per-domain readers serially in the single-state path. + ! soca_fields_read_ensemble flushes all members' readers together via + ! soca_io_readers_commit_ensemble instead. + if (allocated(plan%readers)) then + do r = 1, size(plan%readers) + call plan%readers(r)%commit() + end do + end if + + call soca_fields_read_finalize(self, plan) +end subroutine soca_fields_read + + +! ------------------------------------------------------------------------------ +!> Build a read_plan_t for one soca_fields member: handle pre-read config +!! (Identity, vdate from "date", h_common for vertical remap), open per-domain +!! readers, and enqueue reads onto them. The readers are NOT committed here -- +!! callers do that, either per-reader (single-state path) or in bulk via +!! soca_io_readers_commit_ensemble (ensemble path). +!! +!! \param[inout] self fields to populate +!! \param[in] f_conf config block (basename, *_filename, read_from_file, ...) +!! \param[inout] vdate set from f_conf%"date" if Identity or iread==1 +!! \param[out] plan per-domain readers + scratch + post-processing flags +!! \param[in] reader_pe optional reader_pe for all readers built here +!! (h_common + per-domain). Ensemble path passes +!! soca_io_ensemble_root_pe(m, nmembers) so different members' +!! readers run on different PEs concurrently. +!! \relates soca_fields_mod::soca_fields +subroutine soca_fields_read_prepare(self, f_conf, vdate, plan, reader_pe) + class(soca_fields), target, intent(inout) :: self + type(fckit_configuration), intent(in) :: f_conf + type(datetime), intent(inout) :: vdate + type(read_plan_t), target, intent(inout) :: plan + integer, optional, intent(in) :: reader_pe + + character(len=:), allocatable :: str, basename, filename + type(soca_io_reader) :: hreader + integer :: d, f, i, n, idx, nreaders, ri type(soca_field_metadata) :: field_meta - domains = [character(len=3) :: "ocn", "sfc", "ice", "wav", "bio"] + character(len=3), dimension(5) :: domains - if ( f_conf%has("read_from_file") ) call f_conf%get_or_die("read_from_file", iread) + domains = [character(len=3) :: "ocn", "sfc", "ice", "wav", "bio"] + ! Default iread to 1 when 'basename' is set: the original soca_fields_read had + ! `integer :: iread = 0` which (Fortran gotcha) gave it implicit SAVE, so + ! callers like saber's "model file" path (basename + ocn_filename but no + ! read_from_file) silently relied on iread being inherited from a prior call. + ! Inferring iread=1 from the presence of basename keeps that behavior + ! deterministic without forcing a YAML schema change. + plan%iread = 0 + if (f_conf%has("basename")) plan%iread = 1 + if (f_conf%has("read_from_file")) call f_conf%get_or_die("read_from_file", plan%iread) + + ! Vertical remap target: read h_common immediately on its own reader. + ! Kept serial-per-member in v1; lives outside plan%readers so the bulk + ! ensemble commit isn't blocked on this small single-field read. + if (f_conf%has("remap_filename")) then + call f_conf%get_or_die("remap_filename", str) + allocate(plan%h_common(self%geom%isd:self%geom%ied, & + self%geom%jsd:self%geom%jed, self%geom%nzo)) + plan%h_common = 0.0_kind_real + if (present(reader_pe)) then + call hreader%init(self%geom%Domain%mpp_domain, str, reader_pe=reader_pe) + else + call hreader%init(self%geom%Domain%mpp_domain, str) + end if + call hreader%enqueue('h', plan%h_common) + call hreader%commit() + end if - ! Check if vertical remapping needs to be applied - if ( f_conf%has("remap_filename") ) then - call f_conf%get_or_die("remap_filename", str) - allocate(h_common(self%geom%isd:self%geom%ied, self%geom%jsd:self%geom%jed, self%geom%nzo)) - h_common = 0.0_kind_real + ! Identity (unit increment); independent of iread. Sets vdate from "date". + if (f_conf%has("Identity")) then + call f_conf%get_or_die("Identity", i) + if (i == 1) call self%ones() + call f_conf%get_or_die("date", str) + call datetime_set(str, vdate) + end if - ! Read common vertical coordinate from file - call reader%init(self%geom%Domain%mpp_domain, str) - call reader%enqueue('h', h_common) - call reader%commit() + if (plan%iread /= 1 .and. plan%iread /= 3) then + allocate(plan%readers(0), plan%domain_vars(0)) + return end if - ! Create unit increment - if ( f_conf%has("Identity") ) then - call f_conf%get_or_die("Identity", i) - if ( i==1 ) call self%ones() - call f_conf%get_or_die("date", str) - call datetime_set(str, vdate) + plan%seaice_categories_vars = oops_variables() + + if (plan%iread == 1) then + call f_conf%get_or_die("date", str) + call datetime_set(str, vdate) end if + call f_conf%get_or_die("basename", basename) + + ! Probe for compute_ice/snow thickness from the ice file (if any). + if (f_conf%get("ice_filename", str)) then + filename = trim(basename) // trim(str) + plan%ice_filename = filename + field_meta = self%geom%fields_metadata%get("sea_ice_thickness") + if (.not. soca_io_var_exists(filename, field_meta%io_name)) then + plan%compute_icethickness = .true. + end if + field_meta = self%geom%fields_metadata%get("sea_ice_snow_thickness") + if (.not. soca_io_var_exists(filename, field_meta%io_name)) then + plan%compute_snowthickness = .true. + end if + end if + + ! Pass 1: count domains that have both a filename in config and matching fields. + nreaders = 0 + do d = 1, size(domains) + if (.not. f_conf%get(domains(d)//"_filename", str)) cycle + n = 0 + do i = 1, size(self%fields) + if (self%fields(i)%metadata%io_file == domains(d)) n = n + 1 + end do + if (n > 0) nreaders = nreaders + 1 + end do - ! iread = 1 (state) or 3 (increment): Read restart file - if (iread==1 .or. iread==3) then - seaice_categories_vars = oops_variables() + allocate(plan%readers(nreaders)) + allocate(plan%domain_vars(nreaders)) - ! Set vdate if reading state - if (iread==1) then - call f_conf%get_or_die("date", str) - call datetime_set(str, vdate) + ! Pass 2: init each reader, enqueue all matching vars. Vars whose data are + ! not allocated (the deferred-CICE-category case) keep allocated(data)=.false. + ! and finalize's copy-to-atlas loop skips them. + ri = 0 + do d = 1, size(domains) + if (.not. f_conf%get(domains(d)//"_filename", str)) cycle + filename = trim(basename) // trim(str) + + n = 0 + do i = 1, size(self%fields) + if (self%fields(i)%metadata%io_file == domains(d)) n = n + 1 + end do + if (n == 0) cycle + + ri = ri + 1 + allocate(plan%domain_vars(ri)%vars(n)) + if (present(reader_pe)) then + call plan%readers(ri)%init(self%geom%Domain%mpp_domain, filename, reader_pe=reader_pe) + else + call plan%readers(ri)%init(self%geom%Domain%mpp_domain, filename) end if - call f_conf%get_or_die("basename", basename) - - ! handle constant fields first - do f=1,size(self%fields) - if (self%fields(f)%metadata%io_file == "CONSTANT") then - afield1 = self%afieldset%field(self%fields(f)%name) - call afield1%data(adata1) - adata1(:,:) = self%fields(f)%metadata%constant_value + + n = 0 + do f = 1, size(self%fields) + if (self%fields(f)%metadata%io_file /= domains(d)) cycle + if (domains(d) == "ice" .and. self%fields(f)%metadata%categories > 0) then + ! CICE-history-style aggregated cat+lev files take a separate read path + ! in finalize via self%read_seaice. SOCA-written files have each + ! category as a separate var and read through here normally. + if (soca_io_file_exists(filename) .and. & + soca_io_var_exists(filename, self%fields(f)%metadata%io_name)) then + else + call plan%seaice_categories_vars%push_back(self%fields(f)%name) + cycle + end if end if + n = n + 1 + + associate (v => plan%domain_vars(ri)%vars(n)) + v%field => self%fields(f) + v%afield = self%afieldset%field(v%field%name) + call v%afield%data(v%adata) + allocate(v%data(self%geom%isd:self%geom%ied, & + self%geom%jsd:self%geom%jed, v%field%nz)) + if (v%field%nz == 1) then + if ((self%fields(f)%metadata%name == "sea_ice_thickness") & + .and. plan%compute_icethickness) then + field_meta = self%geom%fields_metadata%get("sea_ice_volume") + call plan%readers(ri)%enqueue(field_meta%io_name, v%data(:,:,1)) + else if ((self%fields(f)%metadata%name == "sea_ice_snow_thickness") & + .and. plan%compute_snowthickness) then + field_meta = self%geom%fields_metadata%get("sea_ice_snow_volume") + call plan%readers(ri)%enqueue(field_meta%io_name, v%data(:,:,1)) + else + call plan%readers(ri)%enqueue(v%field%metadata%io_name, v%data(:,:,1)) + end if + else + call plan%readers(ri)%enqueue(v%field%metadata%io_name, v%data(:,:,:)) + end if + end associate end do + ! unused slots at tail of plan%domain_vars(ri)%vars (skipped CICE cases) + ! keep allocated(v%data)=.false.; finalize handles them. + end do +end subroutine soca_fields_read_prepare - ! determine whether we'll need to compute ice thickness or snow thickness - compute_icethickness = .false. - compute_snowthickness = .false. - if(f_conf%get("ice_filename", str)) then - filename = trim(basename) // trim(str) - field_meta = self%geom%fields_metadata%get("sea_ice_thickness") - if ((.not. soca_io_var_exists(filename, field_meta%io_name))) then - compute_icethickness = .true. - endif - field_meta = self%geom%fields_metadata%get("sea_ice_snow_thickness") - if ((.not. soca_io_var_exists(filename, field_meta%io_name))) then - compute_snowthickness = .true. - endif - endif - ! for each separate domain, check if a filename is provided - do d=1, size(domains) - if(f_conf%get(domains(d)//"_filename", str)) then - filename = trim(basename) // trim(str) +! ------------------------------------------------------------------------------ +!> Drain a populated read_plan_t into self: copy per-domain scratch into atlas +!! fields with land-mask handling, run the existing post-processing chain +!! (CICE-category deferred read, ice/snow thickness compute, mid-layer depth, +!! mixed-layer depth, optional vertical remap), and release scratch. +!! +!! Assumes readers in plan%readers have already been committed. +!! \relates soca_fields_mod::soca_fields +subroutine soca_fields_read_finalize(self, plan) + class(soca_fields), target, intent(inout) :: self + type(read_plan_t), target, intent(inout) :: plan - ! determine how many variables will be read in with this file - n = 0 - do i=1,size(self%fields) - if (self%fields(i)%metadata%io_file == domains(d)) n = n + 1 - end do - if (n == 0) cycle - allocate(vars(n)) + integer :: d, f, i, j, k, n, idx, ri + type(remapping_CS) :: remapCS + type(atlas_field) :: afield1, afield2, afield3, afield4 + real(kind=kind_real), pointer :: adata1(:,:), adata2(:,:), adata3(:,:), adata4(:,:) + real(kind=kind_real), allocatable :: h_common_ij(:), hocn_ij(:), varocn_ij(:), varocn2_ij(:) + logical :: remap_this_point - ! for each variable, setup to read - call reader%init(self%geom%Domain%mpp_domain, filename) - n = 0 - do f=1,size(self%fields) - if (self%fields(f)%metadata%io_file == domains(d)) then - if (domains(d) == "ice" .and. self%fields(f)%metadata%categories > 0) then - ! check if the file was constructed by soca or comes from the CICE history - ! The CICE history aggregates the category and level in 1 array - ! The SOCA io considers categories to be separate variables and will index the naming - if(soca_io_file_exists(filename) .and. soca_io_var_exists(filename, self%fields(f)%metadata%io_name)) then - else - call seaice_categories_vars%push_back(self%fields(f)%name) - cycle - end if - end if - n = n + 1 - - vars(n)%field => self%fields(f) - vars(n)%afield = self%afieldset%field(vars(n)%field%name) - call vars(n)%afield%data(vars(n)%adata) - allocate(vars(n)%data(& - self%geom%isd:self%geom%ied, self%geom%jsd:self%geom%jed, vars(n)%field%nz)) - if (vars(n)%field%nz == 1) then - ! special handling when ice thickness is requested but only ice volume and - ! ice concentration are available - if ((self%fields(f)%metadata%name == "sea_ice_thickness") .and. compute_icethickness) then - field_meta = self%geom%fields_metadata%get("sea_ice_volume") - call reader%enqueue(field_meta%io_name, vars(n)%data(:,:,1)) - elseif ((self%fields(f)%metadata%name == "sea_ice_snow_thickness") .and. compute_snowthickness) then - field_meta = self%geom%fields_metadata%get("sea_ice_snow_volume") - call reader%enqueue(field_meta%io_name, vars(n)%data(:,:,1)) - else - call reader%enqueue(vars(n)%field%metadata%io_name, vars(n)%data(:,:,1)) - endif - else - call reader%enqueue(vars(n)%field%metadata%io_name, vars(n)%data(:,:,:)) - end if - end if - end do + ! Identity-only / no-read path still needs h_common dealloc; nothing else to do. + if (plan%iread /= 1 .and. plan%iread /= 3) then + if (allocated(plan%h_common)) deallocate(plan%h_common) + return + end if - ! read - call reader%commit() - - ! copy back into atlas fields, filling land with fillvalue - do n=1,size(vars) - if (.not. allocated(vars(n)%data)) cycle ! skip special ice fields - vars(n)%adata(:,:) = 0.0 - do j=self%geom%jsc, self%geom%jec - do i=self%geom%isc, self%geom%iec - idx = self%geom%atlas_ij2idx(i,j) - if (associated(vars(n)%field%mask)) then - if (vars(n)%field%mask(i,j) == 0) then - vars(n)%adata(:, idx) = vars(n)%field%metadata%fillvalue - else - vars(n)%adata(:, idx) = vars(n)%data(i,j,:) - end if + ! Set constant-valued atlas fields (independent of any read). + do f = 1, size(self%fields) + if (self%fields(f)%metadata%io_file == "CONSTANT") then + afield1 = self%afieldset%field(self%fields(f)%name) + call afield1%data(adata1) + adata1(:,:) = self%fields(f)%metadata%constant_value + end if + end do + + ! Copy per-domain scratch buffers into atlas fields with land-mask handling. + do ri = 1, size(plan%readers) + associate (vars => plan%domain_vars(ri)%vars) + do n = 1, size(vars) + if (.not. allocated(vars(n)%data)) cycle ! deferred CICE-category slot + vars(n)%adata(:,:) = 0.0 + do j = self%geom%jsc, self%geom%jec + do i = self%geom%isc, self%geom%iec + idx = self%geom%atlas_ij2idx(i,j) + if (associated(vars(n)%field%mask)) then + if (vars(n)%field%mask(i,j) == 0) then + vars(n)%adata(:, idx) = vars(n)%field%metadata%fillvalue else vars(n)%adata(:, idx) = vars(n)%data(i,j,:) end if - end do + else + vars(n)%adata(:, idx) = vars(n)%data(i,j,:) + end if end do - call vars(n)%afield%set_dirty() end do + call vars(n)%afield%set_dirty() + end do + do n = 1, size(vars) + if (.not. allocated(vars(n)%data)) cycle + deallocate(vars(n)%data) + call vars(n)%afield%final() + end do + end associate + deallocate(plan%domain_vars(ri)%vars) + end do - ! done, cleanup - do i=1,size(vars) - if (.not. allocated(vars(i)%data)) cycle - deallocate(vars(i)%data) - call vars(i)%afield%final() - end do - deallocate(vars) - end if - end do + ! Deferred-CICE-category read (its own internal I/O cycle; kept serial in v1). + if (plan%seaice_categories_vars%nvars() > 0) then + call self%read_seaice(plan%ice_filename, plan%seaice_categories_vars) + end if - ! read sea ice variables with category and/or levels dimensions - if (seaice_categories_vars%nvars() > 0) then - call f_conf%get_or_die("ice_filename", str) - filename = trim(basename) // trim(str) - call self%read_seaice(filename, seaice_categories_vars) - end if + ! Compute ice thickness from ice_volume / ice_area_fraction. + if (plan%compute_icethickness .and. self%afieldset%has("sea_ice_thickness")) then + afield1 = self%afieldset%field("sea_ice_thickness") + afield2 = self%afieldset%field("sea_ice_area_fraction") + call afield1%data(adata1) + call afield2%data(adata2) + do j = self%geom%jsc, self%geom%jec + do i = self%geom%isc, self%geom%iec + idx = self%geom%atlas_ij2idx(i,j) + if (adata2(1,idx) > 0.0) then + adata1(1,idx) = adata1(1,idx) / adata2(1,idx) + else + adata1(1,idx) = 0.0_kind_real + end if + end do + end do + call afield1%set_dirty() + end if - ! compute ice thickness if needed - if (compute_icethickness .and. self%afieldset%has("sea_ice_thickness")) then - afield1 = self%afieldset%field("sea_ice_thickness") - afield2 = self%afieldset%field("sea_ice_area_fraction") - call afield1%data(adata1) - call afield2%data(adata2) - do j=self%geom%jsc, self%geom%jec - do i=self%geom%isc, self%geom%iec - idx = self%geom%atlas_ij2idx(i,j) - if (adata2(1,idx) > 0.0) then - adata1(1,idx) = adata1(1,idx) / adata2(1,idx) - else - adata1(1,idx) = 0.0_kind_real - end if - end do + ! Compute snow thickness from snow_volume / ice_area_fraction. + if (plan%compute_snowthickness .and. self%afieldset%has("sea_ice_snow_thickness")) then + afield1 = self%afieldset%field("sea_ice_snow_thickness") + afield2 = self%afieldset%field("sea_ice_area_fraction") + call afield1%data(adata1) + call afield2%data(adata2) + do j = self%geom%jsc, self%geom%jec + do i = self%geom%isc, self%geom%iec + idx = self%geom%atlas_ij2idx(i,j) + if (adata2(1,idx) > 0.0) then + adata1(1,idx) = adata1(1,idx) / adata2(1,idx) + else + adata1(1,idx) = 0.0_kind_real + end if end do - call afield1%set_dirty() - end if - ! compute snow thickness if needed - if (compute_snowthickness .and. self%afieldset%has("sea_ice_snow_thickness")) then - afield1 = self%afieldset%field("sea_ice_snow_thickness") - afield2 = self%afieldset%field("sea_ice_area_fraction") - call afield1%data(adata1) - call afield2%data(adata2) - do j=self%geom%jsc, self%geom%jec - do i=self%geom%isc, self%geom%iec - idx = self%geom%atlas_ij2idx(i,j) - if (adata2(1,idx) > 0.0) then - adata1(1,idx) = adata1(1,idx) / adata2(1,idx) - else - adata1(1,idx) = 0.0_kind_real - end if + end do + call afield1%set_dirty() + end if + + ! TODO: this should live in a variable-change class, not here. + ! Mid-layer depth from layer thickness (running half-thickness sum). + if (self%afieldset%has("sea_water_depth")) then + afield1 = self%afieldset%field("sea_water_depth") + afield2 = self%afieldset%field("sea_water_cell_thickness") + call afield1%data(adata1) + call afield2%data(adata2) + do j = self%geom%jsc, self%geom%jec + do i = self%geom%isc, self%geom%iec + idx = self%geom%atlas_ij2idx(i,j) + adata1(:,idx) = 0.5 * adata2(:,idx) + do k = 2, afield1%shape(1) + adata1(k,idx) = adata1(k,idx) + sum(adata2(1:k-1,idx)) end do end do - call afield1%set_dirty() - end if + end do + call afield1%set_dirty() + end if - ! initialize mid-layer depth from layer thickness - ! TODO, this shouldn't live here, it should be part of the variable change class only - if (self%afieldset%has("sea_water_depth")) then - afield1 = self%afieldset%field("sea_water_depth") - afield2 = self%afieldset%field("sea_water_cell_thickness") - call afield1%data(adata1) - call afield2%data(adata2) - do j=self%geom%jsc, self%geom%jec - do i=self%geom%isc, self%geom%iec - idx = self%geom%atlas_ij2idx(i,j) - adata1(:,idx) = 0.5 * adata2(:,idx) - do k=2,afield1%shape(1) - adata1(k,idx) = adata1(k,idx) + sum(adata2(1:k-1,idx)) - end do - end do + ! TODO: move out. Mixed-layer depth from T/S/depth. + if (self%afieldset%has("ocean_mixed_layer_thickness")) then + afield1 = self%afieldset%field("ocean_mixed_layer_thickness") + afield2 = self%afieldset%field("sea_water_salinity") + afield3 = self%afieldset%field("sea_water_potential_temperature") + afield4 = self%afieldset%field("sea_water_depth") + call afield1%data(adata1) + call afield2%data(adata2) + call afield3%data(adata3) + call afield4%data(adata4) + adata1(:,:) = 0.0_kind_real + do j = self%geom%jsc, self%geom%jec + do i = self%geom%isc, self%geom%iec + if (self%geom%mask2d(i,j) == 0) cycle + idx = self%geom%atlas_ij2idx(i,j) + adata1(1,idx) = soca_mld(adata2(:,idx), adata3(:,idx), adata4(:,idx), & + self%geom%lon(i,j), self%geom%lat(i,j)) end do - call afield1%set_dirty() - end if + end do + call afield1%set_dirty() + end if - ! Compute mixed layer depth TODO: Move somewhere else ... - if (self%afieldset%has("ocean_mixed_layer_thickness")) then - afield1 = self%afieldset%field("ocean_mixed_layer_thickness") - afield2 = self%afieldset%field("sea_water_salinity") - afield3 = self%afieldset%field("sea_water_potential_temperature") - afield4 = self%afieldset%field("sea_water_depth") - call afield1%data(adata1) + ! Optional vertical remap onto h_common. + if (allocated(plan%h_common)) then + call initialize_remapping(remapCS, 'PCM') + allocate(h_common_ij(self%geom%nzo), hocn_ij(self%geom%nzo), & + varocn_ij(self%geom%nzo), varocn2_ij(self%geom%nzo)) + afield1 = self%afieldset%field("sea_water_cell_thickness") + call afield1%data(adata1) + do n = 1, size(self%fields) + if (.not. self%fields(n)%metadata%vert_interp) cycle + if (self%geom%f_comm%rank() == 0) then + call oops_log%info("vertically remapping "//trim(self%fields(n)%name)) + end if + afield2 = self%afieldset%field(self%fields(n)%name) call afield2%data(adata2) - call afield3%data(adata3) - call afield4%data(adata4) - adata1(:,:) = 0.0_kind_real - do j=self%geom%jsc, self%geom%jec - do i=self%geom%isc, self%geom%iec - if (self%geom%mask2d(i,j)==0) cycle + do j = self%geom%jsc, self%geom%jec + do i = self%geom%isc, self%geom%iec idx = self%geom%atlas_ij2idx(i,j) - adata1(1,idx) = soca_mld(adata2(:,idx), adata3(:,idx), adata4(:,idx), & - self%geom%lon(i,j), self%geom%lat(i,j)) - end do - end do - call afield1%set_dirty() - end if - - ! Remap layers if needed - if (allocated(h_common)) then - call initialize_remapping(remapCS,'PCM') - - ! allocate things - allocate(h_common_ij(self%geom%nzo), hocn_ij(self%geom%nzo), & - varocn_ij(self%geom%nzo), varocn2_ij(self%geom%nzo)) - afield1 = self%afieldset%field("sea_water_cell_thickness") - call afield1%data(adata1) - - ! for each field that should be remapped - do n=1,size(self%fields) - if (.not. self%fields(n)%metadata%vert_interp) cycle - if ( self%geom%f_comm%rank() == 0 ) then - call oops_log%info("vertically remapping "//trim(self%fields(n)%name)) - end if - afield2 = self%afieldset%field(self%fields(n)%name) - call afield2%data(adata2) - - ! for each grid point - do j=self%geom%jsc, self%geom%jec - do i=self%geom%isc, self%geom%iec - idx = self%geom%atlas_ij2idx(i,j) - remap_this_point = .not. associated(self%fields(n)%mask) - if (associated(self%fields(n)%mask)) then - remap_this_point = self%fields(n)%mask(i,j) .gt. 0.0 - end if - if (remap_this_point) then - h_common_ij(:) = h_common(i,j,:) - hocn_ij(:) = adata1(:, idx) - varocn_ij(:) = adata2(:, idx) - call remapping_core_h(remapCS, self%geom%nzo, h_common_ij, varocn_ij, & - self%geom%nzo, hocn_ij, varocn2_ij) - adata2(:, idx) = varocn2_ij - else - adata2(:, idx) = 0.0_kind_real - end if - end do + remap_this_point = .not. associated(self%fields(n)%mask) + if (associated(self%fields(n)%mask)) then + remap_this_point = self%fields(n)%mask(i,j) .gt. 0.0 + end if + if (remap_this_point) then + h_common_ij(:) = plan%h_common(i,j,:) + hocn_ij(:) = adata1(:, idx) + varocn_ij(:) = adata2(:, idx) + call remapping_core_h(remapCS, self%geom%nzo, h_common_ij, varocn_ij, & + self%geom%nzo, hocn_ij, varocn2_ij) + adata2(:, idx) = varocn2_ij + else + adata2(:, idx) = 0.0_kind_real + end if end do - call afield2%set_dirty() end do - - ! cleanup - call end_remapping(remapCS) - deallocate(h_common_ij, hocn_ij, varocn_ij, varocn2_ij) - end if + call afield2%set_dirty() + end do + call end_remapping(remapCS) + deallocate(h_common_ij, hocn_ij, varocn_ij, varocn2_ij) end if - ! cleanup - if (allocated(h_common)) deallocate(h_common) + ! Release the empty-ctor'd seaice_categories_vars handle from prepare so + ! the C++-side variables_ allocation doesn't leak per-member. + call plan%seaice_categories_vars%destruct() + + if (allocated(plan%h_common)) deallocate(plan%h_common) call afield1%final() call afield2%final() call afield3%final() call afield4%final() -end subroutine soca_fields_read +end subroutine soca_fields_read_finalize ! Populate an empty oop_variable instance with the unique CICE variables @@ -907,10 +1046,77 @@ subroutine soca_fields_read_seaice(self, filename, seaice_categories_vars) end subroutine soca_fields_read_seaice +! ------------------------------------------------------------------------------ +!> Build one per-domain writer for `fld`: enqueue every field whose io_file +!! matches `dom`, copying each atlas field's compute-domain slice into a fresh +!! varwrapper buffer (masked cells set to the field's fillvalue). `vars` carries +!! the buffers and has the TARGET attribute because the writer holds pointers +!! into them; the caller must keep `vars` alive (and unmutated) until the +!! writer is committed. Leaves `vars` unallocated and `writer` untouched when no +!! field maps to `dom`. Shared by soca_fields_write_rst (single state) and +!! soca_fields_write_ensemble (bulk), so the atlas->compute-domain copy and the +!! enqueue logic live in exactly one place. +!! \relates soca_fields_mod::soca_fields +subroutine soca_fields_enqueue_domain(fld, dom, filename, writer, vars, root_pe) + class(soca_fields), intent(inout) :: fld + character(len=*), intent(in) :: dom + character(len=*), intent(in) :: filename + type(soca_io_writer), intent(inout) :: writer + type(varwrapper), allocatable, target, intent(out) :: vars(:) + integer, optional, intent(in) :: root_pe + + integer :: i, j, f, n, idx + + ! How many fields land in this domain's file? + n = 0 + do f = 1, size(fld%fields) + if (fld%fields(f)%metadata%io_file == dom) n = n + 1 + end do + if (n == 0) return + + if (present(root_pe)) then + call writer%init(fld%geom%Domain%mpp_domain, filename, root_pe=root_pe) + else + call writer%init(fld%geom%Domain%mpp_domain, filename) + end if + + allocate(vars(n)) + n = 0 + do f = 1, size(fld%fields) + if (fld%fields(f)%metadata%io_file /= dom) cycle + n = n + 1 + + ! copy atlas field into a per-PE compute-domain temporary + vars(n)%afield = fld%aFieldset%field(fld%fields(f)%name) + call vars(n)%afield%data(vars(n)%adata) + allocate(vars(n)%data(fld%geom%isd:fld%geom%ied, & + fld%geom%jsd:fld%geom%jed, vars(n)%afield%shape(1))) + + ! copy, setting masked values to fillvalue + if (associated(fld%fields(f)%mask)) vars(n)%data = fld%fields(f)%metadata%fillvalue + do j = fld%geom%jsc, fld%geom%jec + do i = fld%geom%isc, fld%geom%iec + idx = fld%geom%atlas_ij2idx(i, j) + if (associated(fld%fields(f)%mask)) then + if (fld%fields(f)%mask(i, j) == 0) cycle + end if + vars(n)%data(i, j, :) = vars(n)%adata(:, idx) + end do + end do + + ! enqueue with the writer + if (vars(n)%afield%shape(1) == 1) then + call writer%enqueue(fld%fields(f)%metadata%io_name, vars(n)%data(:,:,1)) + else + call writer%enqueue(fld%fields(f)%metadata%io_name, vars(n)%data(:,:,:)) + end if + end do +end subroutine soca_fields_enqueue_domain + + ! ------------------------------------------------------------------------------ !> Save soca fields in a restart format !! -!! TODO this can be generalized even more !! \relates soca_fields_mod::soca_fields subroutine soca_fields_write_rst(self, f_conf, vdate) class(soca_fields), target, intent(inout) :: self !< Fields @@ -919,14 +1125,14 @@ subroutine soca_fields_write_rst(self, f_conf, vdate) integer, parameter :: max_string_length=800 type(soca_io_writer) :: writer - integer :: i, j, k, idx, d, f, n - type(soca_field), pointer :: field + integer :: d, n logical :: date_cols character(len=3), allocatable :: domains(:) character(len=:), allocatable :: domain_filename - type(varwrapper), allocatable, target :: vars(:) ! target so vars(n)%data sections are valid pointer targets for soca_io enqueue + ! target so vars(n)%data sections are valid pointer targets for soca_io enqueue + type(varwrapper), allocatable, target :: vars(:) ! Get date IO format (colons or not?) date_cols = .true. @@ -941,46 +1147,8 @@ subroutine soca_fields_write_rst(self, f_conf, vdate) do d=1,size(domains) domain_filename = soca_genfilename(f_conf,max_string_length,vdate,date_cols,domains(d)) - ! count the number of vars that we will write in this file - n = 0 - do f=1,size(self%fields) - if (self%fields(f)%metadata%io_file == domains(d)) n = n +1 - end do - if (n == 0) cycle - allocate(vars(n)) - - ! copy atlas fields into per-PE compute-domain temporaries and enqueue with the writer - call writer%init(self%geom%Domain%mpp_domain, domain_filename) - n=0 - do f=1,size(self%fields) - if (self%fields(f)%metadata%io_file /= domains(d)) cycle - n = n + 1 - - ! create memory - vars(n)%afield = self%aFieldset%field(self%fields(f)%name) - call vars(n)%afield%data(vars(n)%adata) - allocate(vars(n)%data(self%geom%isd:self%geom%ied, & - self%geom%jsd:self%geom%jed, vars(n)%afield%shape(1))) - - ! copy, setting masked values to fillvalue - if (associated(self%fields(f)%mask)) vars(n)%data = self%fields(f)%metadata%fillvalue - do j=self%geom%jsc, self%geom%jec - do i=self%geom%isc, self%geom%iec - idx = self%geom%atlas_ij2idx(i,j) - if (associated(self%fields(f)%mask)) then - if (self%fields(f)%mask(i,j) == 0) cycle - end if - vars(n)%data(i,j,:) = vars(n)%adata(:, idx) - end do - end do - - ! enqueue with the writer - if (vars(n)%afield%shape(1) == 1) then - call writer%enqueue(self%fields(f)%metadata%io_name, vars(n)%data(:,:,1)) - else - call writer%enqueue(self%fields(f)%metadata%io_name, vars(n)%data(:,:,:)) - end if - end do + call soca_fields_enqueue_domain(self, domains(d), domain_filename, writer, vars) + if (.not. allocated(vars)) cycle ! no fields in this domain ! write the file call writer%commit() @@ -994,6 +1162,220 @@ subroutine soca_fields_write_rst(self, f_conf, vdate) end do end subroutine soca_fields_write_rst +! ------------------------------------------------------------------------------ +!> Bulk write multiple states/increments in one ensemble pass. Each member's +!! per-domain files are written from a rotated writer PE (assigned by +!! soca_io_ensemble_root_pe) so disk writes happen concurrently across the +!! world communicator. The geometry.io.'ensemble batch size' knob bounds how +!! many members are committed per pass (batch 1 serializes, like a serial write). +!! +!! Behavior per-member is identical to soca_fields_write_rst: both build each +!! per-domain writer via soca_fields_enqueue_domain (same domain set, filename, +!! mask/fillvalue handling and byte-level netCDF output). Only the dispatch +!! (rotated root + ensemble orchestrator) changes. +!! +!! \param fields_ptrs polymorphic pointers to soca_state or soca_increment +!! \param f_confs per-member configurations (must align 1:1 with fields_ptrs) +!! \param vdates per-member valid dates +!! \relates soca_fields_mod::soca_fields +subroutine soca_fields_write_ensemble(fields_ptrs, c_confs, vdates) + type(soca_fields_ptr_t), intent(inout) :: fields_ptrs(:) + type(c_ptr), intent(in) :: c_confs(:) + type(datetime), intent(inout) :: vdates(:) + + integer, parameter :: max_string_length = 800 + integer :: nmembers, m, d, f, n, wi, nwriters, root_pe + integer :: batch_size, mb_start, mb_end, nbatch + character(len=3), allocatable :: domains(:) + character(len=:), allocatable :: domain_filename + logical :: date_cols + class(soca_fields), pointer :: fld + type(soca_io_writer), allocatable, target :: writers(:) + type(ensemble_vars_t), allocatable, target :: vars_arr(:) + type(fckit_configuration) :: f_conf + + nmembers = size(fields_ptrs) + if (nmembers == 0) return + if (size(c_confs) /= nmembers .or. size(vdates) /= nmembers) then + call abor1_ftn("soca_fields_write_ensemble: input array size mismatch") + end if + + domains = [character(len=3) :: "ocn", "sfc", "ice", "wav", "bio"] + + ! Resolve the batch size: process the ensemble in waves of at most batch_size + ! members so peak memory is bounded by one wave's writers + scratch rather + ! than the whole ensemble. 0 / >= nmembers means a single batch (the default, + ! maximum overlap, byte-identical to the unbatched path). + batch_size = soca_io_ensemble_batch_size() + if (batch_size <= 0 .or. batch_size > nmembers) batch_size = nmembers + + do mb_start = 1, nmembers, batch_size + mb_end = min(mb_start + batch_size - 1, nmembers) + nbatch = mb_end - mb_start + 1 + + ! Pass 1: count writers (one per (member, domain) with vars) in this batch. + nwriters = 0 + do m = mb_start, mb_end + fld => fields_ptrs(m)%p + do d = 1, size(domains) + n = 0 + do f = 1, size(fld%fields) + if (fld%fields(f)%metadata%io_file == domains(d)) n = n + 1 + end do + if (n > 0) nwriters = nwriters + 1 + end do + end do + if (nwriters == 0) cycle + + allocate(writers(nwriters)) + allocate(vars_arr(nwriters)) + + ! Pass 2: build writers and enqueue. Construct fckit_configuration on the fly + ! per member (array assignment of fckit_configuration drops the underlying + ! c_ptr in some compiler/version combos, so we materialize a fresh handle + ! each iteration -- same pattern as the single-state interface). + wi = 0 + do m = mb_start, mb_end + fld => fields_ptrs(m)%p + f_conf = fckit_configuration(c_confs(m)) + date_cols = .true. + if (f_conf%has("date colons")) then + call f_conf%get_or_die("date colons", date_cols) + end if + + ! Rotate the writer root PE by position WITHIN the batch so each batch + ! spreads its members one-per-PE across the comm (rather than by global + ! member index, which would pile a consecutive batch onto few PEs). + root_pe = soca_io_ensemble_root_pe(m - mb_start + 1, nbatch) + + do d = 1, size(domains) + ! count vars for this (member, domain); must stay in lockstep with the + ! pass-1 nwriters count so wi indexes the preallocated writers array. + n = 0 + do f = 1, size(fld%fields) + if (fld%fields(f)%metadata%io_file == domains(d)) n = n + 1 + end do + if (n == 0) cycle + + wi = wi + 1 + domain_filename = soca_genfilename(f_conf, max_string_length, vdates(m), date_cols, domains(d)) + call soca_fields_enqueue_domain(fld, domains(d), domain_filename, & + writers(wi), vars_arr(wi)%vars, root_pe=root_pe) + end do + end do + + ! Commit the batch through the ensemble orchestrator. Serialization for a + ! memory ceiling is the batch-size knob's job (batch 1 commits one member at + ! a time, matching a serial write); there is no separate sequential path. + call soca_io_writers_commit_ensemble(writers) + + ! Cleanup. Atlas fields use refcounted handles -> need explicit final(). + do wi = 1, nwriters + if (.not. allocated(vars_arr(wi)%vars)) cycle + do n = 1, size(vars_arr(wi)%vars) + if (allocated(vars_arr(wi)%vars(n)%data)) deallocate(vars_arr(wi)%vars(n)%data) + call vars_arr(wi)%vars(n)%afield%final() + end do + deallocate(vars_arr(wi)%vars) + end do + deallocate(writers, vars_arr) + end do +end subroutine soca_fields_write_ensemble + + +! ------------------------------------------------------------------------------ +!> Bulk read multiple states/increments. Mirrors soca_fields_write_ensemble: +!! prepare per-member read_plan_t with a rotated reader_pe, flatten all +!! per-(member, domain) readers into a single array, drive the bulk I/O via +!! soca_io_readers_commit_ensemble (strided or scatter, sync or async per the +!! geometry.io knobs), then finalize each member to copy scratch into atlas +!! fields and run post-read transforms. +!! +!! \param fields_ptrs polymorphic pointers to soca_state or soca_increment +!! \param c_confs per-member configurations (c_ptrs; materialized inline) +!! \param vdates per-member valid dates +!! \relates soca_fields_mod::soca_fields +subroutine soca_fields_read_ensemble(fields_ptrs, c_confs, vdates) + type(soca_fields_ptr_t), intent(inout) :: fields_ptrs(:) + type(c_ptr), intent(in) :: c_confs(:) + type(datetime), intent(inout) :: vdates(:) + + integer :: nmembers, m, ri, nreaders_total, root_pe, k + integer :: batch_size, mb_start, mb_end, nbatch, j + type(read_plan_t), allocatable, target :: plans(:) + type(soca_io_reader), allocatable :: flat_readers(:) + type(fckit_configuration) :: f_conf + + nmembers = size(fields_ptrs) + if (nmembers == 0) return + if (size(c_confs) /= nmembers .or. size(vdates) /= nmembers) then + call abor1_ftn("soca_fields_read_ensemble: input array size mismatch") + end if + + ! Resolve the batch size: read the ensemble in waves of at most batch_size + ! members so the staged global fields (one per member on each reader PE) plus + ! per-member scratch never exceed one wave's worth. 0 / >= nmembers means a + ! single batch (the default, byte-identical to the unbatched read). + batch_size = soca_io_ensemble_batch_size() + if (batch_size <= 0 .or. batch_size > nmembers) batch_size = nmembers + + do mb_start = 1, nmembers, batch_size + mb_end = min(mb_start + batch_size - 1, nmembers) + nbatch = mb_end - mb_start + 1 + + allocate(plans(nbatch)) + + ! Pass 1: prepare each member in the batch. The reader_pe is rotated by + ! position WITHIN the batch so each batch spreads its bulk reads one-per-PE + ! across the comm (hitting different files concurrently in the scatter path). + ! Construct fckit_configuration on the fly per member -- array-of-handles + ! through the c-binding has dropped the c_ptr in some compiler/version + ! combos; same pattern as soca_fields_write_ensemble. + do j = 1, nbatch + m = mb_start + j - 1 + f_conf = fckit_configuration(c_confs(m)) + root_pe = soca_io_ensemble_root_pe(j, nbatch) + call soca_fields_read_prepare(fields_ptrs(m)%p, f_conf, vdates(m), plans(j), & + reader_pe=root_pe) + end do + + ! Pass 2: flatten all per-domain readers into one array for the bulk commit. + ! Intrinsic assignment deep-copies the read_entry vars(:) but shallow-copies + ! their pointer components (data1d/2d/3d/4d -> plans(j)%domain_vars(...)%data + ! scratch buffers) and the domain pointer. The read_entry gbuf_* allocatables + ! are unallocated at copy time (only populated during reader_stage_read on + ! flat_readers), so the deep-copy doesn't duplicate any payload. The bulk + ! scatter delivers data straight into the buffers finalize reads from. + nreaders_total = 0 + do j = 1, nbatch + nreaders_total = nreaders_total + size(plans(j)%readers) + end do + if (nreaders_total > 0) then + allocate(flat_readers(nreaders_total)) + k = 0 + do j = 1, nbatch + do ri = 1, size(plans(j)%readers) + k = k + 1 + flat_readers(k) = plans(j)%readers(ri) + end do + end do + call soca_io_readers_commit_ensemble(flat_readers) + deallocate(flat_readers) + end if + + ! Pass 3: per-member finalize -- copy scratch -> atlas, run post-processing + ! (CICE category deferred read, ice/snow thickness, mid-layer depth, MLD, + ! vertical remap), release scratch. + do j = 1, nbatch + m = mb_start + j - 1 + call soca_fields_read_finalize(fields_ptrs(m)%p, plans(j)) + end do + + deallocate(plans) + end do +end subroutine soca_fields_read_ensemble + + ! ------------------------------------------------------------------------------ !> Interpolates from uv-points location to h-points. !! diff --git a/src/soca/Geometry/soca_geom_mod.F90 b/src/soca/Geometry/soca_geom_mod.F90 index 160e7dd5..2c021cb7 100644 --- a/src/soca/Geometry/soca_geom_mod.F90 +++ b/src/soca/Geometry/soca_geom_mod.F90 @@ -17,7 +17,7 @@ module soca_geom_mod ! mom6 / fms modules use fms_mod, only : fms_init, fms_end -use soca_io_mod, only : soca_io_writer, soca_io_reader +use soca_io_mod, only : soca_io_writer, soca_io_reader, soca_io_config_from_yaml use MOM, only : MOM_control_struct, initialize_MOM, MOM_end, get_MOM_state_elements use MOM_restart, only :MOM_restart_CS ! NOTE remove this when updating MOM6 use MOM_domains, only : MOM_domain_type, MOM_domains_init, MOM_infra_init, MOM_infra_end @@ -181,6 +181,10 @@ subroutine soca_geom_init(self, f_conf, f_comm, gen) ! MPI communicator self%f_comm = f_comm + ! Resolve the soca_io dispatch knobs (parallel/sequential, scatter/strided, + ! async mpi) from geometry.io once; values persist module-level for the run. + call soca_io_config_from_yaml(f_conf) + ! use MOM6 to setup domain decomposition call mpp_init(localcomm=f_comm%communicator()) call fms_init() diff --git a/src/soca/IO/soca_io_mod.F90 b/src/soca/IO/soca_io_mod.F90 index 6baa6569..fa14729b 100644 --- a/src/soca/IO/soca_io_mod.F90 +++ b/src/soca/IO/soca_io_mod.F90 @@ -24,12 +24,14 @@ module soca_io_mod use netcdf +use mpi use kinds, only: kind_real use mpp_mod, only: mpp_gather, mpp_scatter, mpp_broadcast, mpp_pe, mpp_root_pe, mpp_npes, & mpp_get_current_pelist, mpp_error, FATAL use mpp_domains_mod, only: domain2D, & mpp_get_compute_domain, mpp_get_global_domain, & mpp_get_data_domain +use fckit_configuration_module, only: fckit_configuration implicit none private @@ -37,8 +39,28 @@ module soca_io_mod public :: soca_io_writer public :: soca_io_reader public :: soca_io_file_exists, soca_io_var_exists +public :: soca_io_config_from_yaml +public :: soca_io_ensemble_batch_size +public :: soca_io_ensemble_root_pe +public :: soca_io_writers_commit_ensemble +public :: soca_io_readers_commit_ensemble integer, parameter :: MAX_NAME = 256 + +! Module-level I/O dispatch knobs, set once from geometry.io YAML at +! soca_geom_init via soca_io_config_from_yaml; persist for the run. Defaults are +! tuned for cluster-scale production (parallel write, scatter ensemble read, +! async MPI); single-state read stays strided to preserve direct-io behavior. +logical, save :: ensemble_read_scatter_ = .true. +logical, save :: single_state_read_scatter_ = .false. +logical, save :: async_mpi_enabled_ = .true. +! Ensemble I/O batch size: how many members are processed (built, committed, +! freed) in one bulk read/write pass. 0 (the default) means a single batch over +! all members -- maximum I/O overlap, peak memory ~ all members resident at +! once. A positive N caps the working set to N members per pass, trading I/O +! concurrency for a memory ceiling (set N ~ npes for one wave per PE). Consumed +! by the soca_fields ensemble drivers. +integer, save :: ensemble_batch_size_ = 0 integer, parameter :: MAX_FILE_NDIMS = 8 ! Reader is stateless across commits: every reader_commit nf90_opens the file, @@ -63,6 +85,14 @@ module soca_io_mod real(kind=kind_real), pointer :: data3d(:,:,:) => null() character(len=MAX_NAME) :: long_name = '' character(len=MAX_NAME) :: units = 'none' + ! Inter-stage gather buffers used by writer_stage_gather -> writer_stage_write. + ! Allocated lazily on the writer's root_pe in stage_gather (size nx_g x ny_g + ! for 2D, nx_g x ny_g x nlevels for 3D); freed in stage_close. Holding them + ! per-var (rather than per-writer reused) lets writer_stage_gather populate + ! every var before any disk write starts -- so in the ensemble path the + ! writes in stage_write can run concurrently across writer PEs with no MPI. + real(kind=kind_real), allocatable :: gbuf_2d(:,:) + real(kind=kind_real), allocatable :: gbuf_3d(:,:,:) end type var_entry ! Tracks one unique axis (X/Y/Z) by (size, domain_key) and remembers the netcdf @@ -85,6 +115,18 @@ module soca_io_mod integer :: nx_g, ny_g ! global x/y sizes type(var_entry), allocatable :: vars(:) integer :: nvars = 0 + ! Writer's root PE: -1 sentinel means writer_init defaults it to mpp_root_pe() + ! at use time (single-shot path). The ensemble orchestrator sets a rotated + ! root_pe before stage dispatch so different members write from different PEs. + integer :: root_pe = -1 + ! Inter-stage state populated by writer_stage_define and consumed by the later + ! stages. ncid is held open between define and close (only meaningful on root). + integer :: ncid = -1 + integer, allocatable :: varids(:) + ! Per-direction unique-axis tables built in stage_define; stage_close frees. + type(axis_entry), allocatable :: x_axes(:), y_axes(:), z_axes(:) + integer, allocatable :: var_x_idx(:), var_y_idx(:), var_z_idx(:) + integer :: nx_axes = 0, ny_axes = 0, nz_axes = 0 contains procedure :: init => writer_init procedure :: enqueue_1d => writer_enqueue_1d @@ -104,6 +146,13 @@ module soca_io_mod real(kind=kind_real), pointer :: data2d(:,:) => null() real(kind=kind_real), pointer :: data3d(:,:,:) => null() real(kind=kind_real), pointer :: data4d(:,:,:,:) => null() + ! Inter-stage staged buffers for the scatter path: reader_pe pulls the + ! WHOLE global field into these during reader_stage_read; reader_stage_distribute + ! ships the per-PE compute tile out via MPI_Scatterv and the buffers are freed + ! in reader_stage_close. Only allocated on reader_pe (1x1 dummies elsewhere). + real(kind=kind_real), allocatable :: gbuf_2d(:,:) + real(kind=kind_real), allocatable :: gbuf_3d(:,:,:) + real(kind=kind_real), allocatable :: gbuf_4d(:,:,:,:) end type read_entry type :: soca_io_reader @@ -115,6 +164,10 @@ module soca_io_mod integer :: nx_g, ny_g ! global x/y sizes type(read_entry), allocatable :: vars(:) integer :: nvars = 0 + ! Reader's source PE for the scatter path: -1 sentinel means reader_init + ! defaults it to mpp_root_pe(). The ensemble orchestrator sets a rotated + ! reader_pe per member so different files are read concurrently across PEs. + integer :: reader_pe = -1 contains procedure :: init => reader_init procedure :: enqueue_1d => reader_enqueue_1d @@ -129,12 +182,15 @@ module soca_io_mod !============================================================================== ! init: prepare a writer for a specific file. The domain pointer is stored, so -! the caller must keep the domain alive until commit returns. +! the caller must keep the domain alive until commit returns. Optional root_pe +! overrides the default mpp_root_pe() target for gather+write; the ensemble +! orchestrator uses this to stride writes across writer PEs. !============================================================================== -subroutine writer_init(self, domain, filename) +subroutine writer_init(self, domain, filename, root_pe) class(soca_io_writer), intent(inout) :: self type(domain2D), target, intent(in) :: domain character(len=*), intent(in) :: filename + integer, optional, intent(in) :: root_pe self%filename = filename self%domain => domain @@ -148,6 +204,16 @@ subroutine writer_init(self, domain, filename) if (allocated(self%vars)) deallocate(self%vars) allocate(self%vars(64)) self%nvars = 0 + + if (present(root_pe)) then + self%root_pe = root_pe + else + self%root_pe = mpp_root_pe() + end if + self%ncid = -1 + self%nx_axes = 0 + self%ny_axes = 0 + self%nz_axes = 0 end subroutine writer_init @@ -236,200 +302,542 @@ end subroutine grow_if_needed !============================================================================== -! commit: PE 0 creates the file structure, then each var is mpp_gather'd and -! PE 0 writes via nf90_put_var. Equivalent to FMS mpp_io threading=MPP_SINGLE -! -- the goal is a clean, debuggable, FMS-free baseline, not a speedup. +! commit: single-shot orchestrator. Sequences the four stages on one writer +! (define -> gather -> write -> close). The original monolithic writer_commit +! was split into stages so the ensemble path (soca_io_writers_commit_ensemble) +! can interleave them across multiple writers: define-all -> gather-all -> +! write-all -> close-all. Single-shot semantics are unchanged. !============================================================================== subroutine writer_commit(self) class(soca_io_writer), intent(inout) :: self - integer :: ncid, dimid_t, varid_t, v - integer, allocatable :: varids(:) - integer, allocatable :: var_x_idx(:), var_y_idx(:), var_z_idx(:) - type(axis_entry), allocatable :: x_axes(:), y_axes(:), z_axes(:) - integer :: nx_axes, ny_axes, nz_axes - logical :: is_root - real(kind=kind_real), allocatable :: gbuf2d(:,:), gbuf3d(:,:,:) - integer, allocatable :: pelist(:) - integer :: dom_key - integer, parameter :: MAX_AXES_PER_DIR = 32 - integer :: nx_c, ny_c, i_off, j_off, nlev - real(kind=kind_real), allocatable :: tile2(:,:), tile3(:,:,:) + call writer_stage_define(self) + call writer_stage_gather(self) + call writer_stage_write(self) + call writer_stage_close(self) +end subroutine writer_commit - is_root = (mpp_pe() == mpp_root_pe()) - call mpi_pelist(pelist) - nx_c = self%iec - self%isc + 1 - ny_c = self%jec - self%jsc + 1 - i_off = self%isc - self%isd + 1 - j_off = self%jsc - self%jsd + 1 +!============================================================================== +! Stage 1 (define): build axis tables, then on the writer's root_pe create the +! netCDF file, define every dim/coord-var/data-var, exit define mode, and write +! the Time + axis coord-var data. Across multiple writers with different +! root_pes (ensemble path), the file-create work runs concurrently because each +! writer's root is on a different rank and no MPI is involved. +!============================================================================== +subroutine writer_stage_define(self) + type(soca_io_writer), intent(inout) :: self - ! Build per-direction unique-axis tables (FMS algorithm): match each var's - ! dim against existing axes by (size, domain_key); reuse on match, append on - ! miss. - allocate(x_axes(MAX_AXES_PER_DIR), y_axes(MAX_AXES_PER_DIR), z_axes(MAX_AXES_PER_DIR)) - allocate(var_x_idx(self%nvars), var_y_idx(self%nvars), var_z_idx(self%nvars)) - var_x_idx = 0; var_y_idx = 0; var_z_idx = 0 - nx_axes = 0; ny_axes = 0; nz_axes = 0 + integer, parameter :: MAX_AXES_PER_DIR = 32 + integer :: v, dom_key, dimid_t, varid_t + logical :: is_root + + is_root = (mpp_pe() == self%root_pe) + + ! Allocate per-direction unique-axis tables on every PE (cheap; need them on + ! every PE because find_or_add_axis is invoked outside the is_root branch). + if (allocated(self%x_axes)) deallocate(self%x_axes) + if (allocated(self%y_axes)) deallocate(self%y_axes) + if (allocated(self%z_axes)) deallocate(self%z_axes) + if (allocated(self%var_x_idx)) deallocate(self%var_x_idx) + if (allocated(self%var_y_idx)) deallocate(self%var_y_idx) + if (allocated(self%var_z_idx)) deallocate(self%var_z_idx) + allocate(self%x_axes(MAX_AXES_PER_DIR), self%y_axes(MAX_AXES_PER_DIR), self%z_axes(MAX_AXES_PER_DIR)) + allocate(self%var_x_idx(self%nvars), self%var_y_idx(self%nvars), self%var_z_idx(self%nvars)) + self%var_x_idx = 0; self%var_y_idx = 0; self%var_z_idx = 0 + self%nx_axes = 0; self%ny_axes = 0; self%nz_axes = 0 do v = 1, self%nvars - ! 1D vars: no domain (global on every PE, key=0). 2D/3D vars share - ! self%domain (key=1). + ! 1D vars: no domain (key=0; global on every PE). 2D/3D vars share self%domain (key=1). dom_key = 1 if (self%vars(v)%ndims == 1) dom_key = 0 - - ! Every var contributes an X axis (Fortran first dim). For 2D/3D the local - ! buffer is only the compute slice, so use the writer's GLOBAL extents for - ! the axis size; 1D buffers are already global. select case (self%vars(v)%ndims) case (1) - call find_or_add_axis(x_axes, nx_axes, size(self%vars(v)%data1d), dom_key, var_x_idx(v)) + call find_or_add_axis(self%x_axes, self%nx_axes, & + size(self%vars(v)%data1d), dom_key, self%var_x_idx(v)) case (2) - call find_or_add_axis(x_axes, nx_axes, self%nx_g, dom_key, var_x_idx(v)) - call find_or_add_axis(y_axes, ny_axes, self%ny_g, dom_key, var_y_idx(v)) + call find_or_add_axis(self%x_axes, self%nx_axes, self%nx_g, dom_key, self%var_x_idx(v)) + call find_or_add_axis(self%y_axes, self%ny_axes, self%ny_g, dom_key, self%var_y_idx(v)) case (3) - call find_or_add_axis(x_axes, nx_axes, self%nx_g, dom_key, var_x_idx(v)) - call find_or_add_axis(y_axes, ny_axes, self%ny_g, dom_key, var_y_idx(v)) - call find_or_add_axis(z_axes, nz_axes, size(self%vars(v)%data3d, 3), dom_key, var_z_idx(v)) + call find_or_add_axis(self%x_axes, self%nx_axes, self%nx_g, dom_key, self%var_x_idx(v)) + call find_or_add_axis(self%y_axes, self%ny_axes, self%ny_g, dom_key, self%var_y_idx(v)) + call find_or_add_axis(self%z_axes, self%nz_axes, & + size(self%vars(v)%data3d, 3), dom_key, self%var_z_idx(v)) end select end do - ! Phase 1: PE 0 defines the file structure -- dims, coord vars, data vars. - if (is_root) then - allocate(varids(self%nvars)) + if (.not. is_root) return - call ncc(nf90_create(self%filename, & - ior(NF90_CLOBBER, ior(NF90_NETCDF4, NF90_CLASSIC_MODEL)), ncid), & - 'nf90_create '//trim(self%filename)) + if (allocated(self%varids)) deallocate(self%varids) + allocate(self%varids(self%nvars)) - call define_axis_dims_and_coords(ncid, x_axes, nx_axes, 'xaxis_', 'X') - call define_axis_dims_and_coords(ncid, y_axes, ny_axes, 'yaxis_', 'Y') - call define_axis_dims_and_coords(ncid, z_axes, nz_axes, 'zaxis_', 'Z') + call ncc(nf90_create(self%filename, & + ior(NF90_CLOBBER, ior(NF90_NETCDF4, NF90_CLASSIC_MODEL)), self%ncid), & + 'nf90_create '//trim(self%filename)) - call ncc(nf90_def_dim(ncid, 'Time', NF90_UNLIMITED, dimid_t), 'def_dim Time') - call ncc(nf90_def_var(ncid, 'Time', NF90_DOUBLE, [dimid_t], varid_t), 'def_var Time') - call ncc(nf90_put_att(ncid, varid_t, 'long_name', 'Time'), 'att Time:long_name') - call ncc(nf90_put_att(ncid, varid_t, 'units', 'time level'), 'att Time:units') - call ncc(nf90_put_att(ncid, varid_t, 'cartesian_axis','T'), 'att Time:cartesian_axis') + call define_axis_dims_and_coords(self%ncid, self%x_axes, self%nx_axes, 'xaxis_', 'X') + call define_axis_dims_and_coords(self%ncid, self%y_axes, self%ny_axes, 'yaxis_', 'Y') + call define_axis_dims_and_coords(self%ncid, self%z_axes, self%nz_axes, 'zaxis_', 'Z') - do v = 1, self%nvars - select case (self%vars(v)%ndims) - case (1) - ! Fortran dim list: [xaxis, Time] -> file order (Time, xaxis) - call ncc(nf90_def_var(ncid, trim(self%vars(v)%name), NF90_DOUBLE, & - [x_axes(var_x_idx(v))%dimid, dimid_t], varids(v)), & - 'def_var '//trim(self%vars(v)%name)) - case (2) - ! Fortran [xaxis, yaxis, Time] -> file (Time, yaxis, xaxis) - call ncc(nf90_def_var(ncid, trim(self%vars(v)%name), NF90_DOUBLE, & - [x_axes(var_x_idx(v))%dimid, y_axes(var_y_idx(v))%dimid, dimid_t], & - varids(v)), 'def_var '//trim(self%vars(v)%name)) - case (3) - ! Fortran [xaxis, yaxis, zaxis, Time] -> file (Time, zaxis, yaxis, xaxis) - call ncc(nf90_def_var(ncid, trim(self%vars(v)%name), NF90_DOUBLE, & - [x_axes(var_x_idx(v))%dimid, y_axes(var_y_idx(v))%dimid, & - z_axes(var_z_idx(v))%dimid, dimid_t], varids(v)), & - 'def_var '//trim(self%vars(v)%name)) - end select - call ncc(nf90_put_att(ncid, varids(v), 'long_name', trim(self%vars(v)%long_name)), & - 'att '//trim(self%vars(v)%name)//':long_name') - call ncc(nf90_put_att(ncid, varids(v), 'units', trim(self%vars(v)%units)), & - 'att '//trim(self%vars(v)%name)//':units') - end do + call ncc(nf90_def_dim(self%ncid, 'Time', NF90_UNLIMITED, dimid_t), 'def_dim Time') + call ncc(nf90_def_var(self%ncid, 'Time', NF90_DOUBLE, [dimid_t], varid_t), 'def_var Time') + call ncc(nf90_put_att(self%ncid, varid_t, 'long_name', 'Time'), 'att Time:long_name') + call ncc(nf90_put_att(self%ncid, varid_t, 'units', 'time level'), 'att Time:units') + call ncc(nf90_put_att(self%ncid, varid_t, 'cartesian_axis','T'), 'att Time:cartesian_axis') - call ncc(nf90_enddef(ncid), 'enddef') + do v = 1, self%nvars + select case (self%vars(v)%ndims) + case (1) + call ncc(nf90_def_var(self%ncid, trim(self%vars(v)%name), NF90_DOUBLE, & + [self%x_axes(self%var_x_idx(v))%dimid, dimid_t], self%varids(v)), & + 'def_var '//trim(self%vars(v)%name)) + case (2) + call ncc(nf90_def_var(self%ncid, trim(self%vars(v)%name), NF90_DOUBLE, & + [self%x_axes(self%var_x_idx(v))%dimid, & + self%y_axes(self%var_y_idx(v))%dimid, dimid_t], & + self%varids(v)), 'def_var '//trim(self%vars(v)%name)) + case (3) + call ncc(nf90_def_var(self%ncid, trim(self%vars(v)%name), NF90_DOUBLE, & + [self%x_axes(self%var_x_idx(v))%dimid, & + self%y_axes(self%var_y_idx(v))%dimid, & + self%z_axes(self%var_z_idx(v))%dimid, dimid_t], & + self%varids(v)), 'def_var '//trim(self%vars(v)%name)) + end select + call ncc(nf90_put_att(self%ncid, self%varids(v), 'long_name', trim(self%vars(v)%long_name)), & + 'att '//trim(self%vars(v)%name)//':long_name') + call ncc(nf90_put_att(self%ncid, self%varids(v), 'units', trim(self%vars(v)%units)), & + 'att '//trim(self%vars(v)%name)//':units') + end do - ! Coordinate-var data is just the index sequence 1..size, matching FMS. - call ncc(nf90_put_var(ncid, varid_t, [1.0_kind_real], start=[1], count=[1]), 'put Time') - call put_axis_coord_data(ncid, x_axes, nx_axes) - call put_axis_coord_data(ncid, y_axes, ny_axes) - call put_axis_coord_data(ncid, z_axes, nz_axes) - end if + call ncc(nf90_enddef(self%ncid), 'enddef') - ! Phase 2: gather and write each user variable. The compute-domain tile is - ! extracted from the caller's halo-inclusive buffer one var (and, for 3D, - ! one level) at a time, so peak local memory is a single (nx_c, ny_c) tile. - ! Non-root allocates a 1x1 dummy for gbuf2d so the actual argument to - ! mpp_gather is always allocated (assumed-shape dummy requires it). - if (is_root) then - allocate(gbuf2d(self%nx_g, self%ny_g)) - else - allocate(gbuf2d(1, 1)) - end if - allocate(tile2(nx_c, ny_c)) + call ncc(nf90_put_var(self%ncid, varid_t, [1.0_kind_real], start=[1], count=[1]), 'put Time') + call put_axis_coord_data(self%ncid, self%x_axes, self%nx_axes) + call put_axis_coord_data(self%ncid, self%y_axes, self%ny_axes) + call put_axis_coord_data(self%ncid, self%z_axes, self%nz_axes) +end subroutine writer_stage_define + + +!============================================================================== +! Stage 2 (gather): for each 2D/3D var, extract the compute-domain tile from +! the caller's halo-inclusive buffer and mpp_gather it to self%root_pe; the +! gathered global field lands in self%vars(v)%gbuf_2d / gbuf_3d (held until +! stage_write). 1D vars are global on every PE and stage_write reads them +! directly; this stage is a no-op for 1D. +! +! Non-root PEs allocate 1x1 dummies so the recv buffer argument to mpp_gather +! is allocated (Fortran assumed-shape contract). +!============================================================================== +subroutine writer_stage_gather(self) + type(soca_io_writer), intent(inout) :: self + + integer, allocatable :: pelist(:) + integer :: v, nx_c, ny_c, i_off, j_off, nlev + real(kind=kind_real), allocatable :: tile2(:,:), tile3(:,:,:) + logical :: is_root + + is_root = (mpp_pe() == self%root_pe) + call mpi_pelist(pelist) + + nx_c = self%iec - self%isc + 1 + ny_c = self%jec - self%jsc + 1 + i_off = self%isc - self%isd + 1 + j_off = self%jsc - self%jsd + 1 do v = 1, self%nvars - if (self%vars(v)%ndims == 1) then + select case (self%vars(v)%ndims) + case (1) + ! Nothing to do; data1d is global on every PE. + case (2) + if (.not. allocated(tile2)) allocate(tile2(nx_c, ny_c)) if (is_root) then - call ncc(nf90_put_var(ncid, varids(v), self%vars(v)%data1d, & - start=[1, 1], count=[size(self%vars(v)%data1d), 1]), & - 'put '//trim(self%vars(v)%name)) + if (.not. allocated(self%vars(v)%gbuf_2d)) & + allocate(self%vars(v)%gbuf_2d(self%nx_g, self%ny_g)) + else + if (.not. allocated(self%vars(v)%gbuf_2d)) allocate(self%vars(v)%gbuf_2d(1, 1)) end if - else if (self%vars(v)%ndims == 2) then tile2 = self%vars(v)%data2d(i_off : i_off + nx_c - 1, & j_off : j_off + ny_c - 1) call mpp_gather(self%isc, self%iec, self%jsc, self%jec, pelist, & - tile2, gbuf2d, is_root) - if (is_root) then - call ncc(nf90_put_var(ncid, varids(v), gbuf2d, & - start=[1, 1, 1], count=[self%nx_g, self%ny_g, 1]), & - 'put '//trim(self%vars(v)%name)) - end if - else - ! Single 3D mpp_gather per 3D var: one collective replaces nlevels 2D - ! gathers, and root receives the assembled global field directly into - ! gbuf3d (no per-level gbuf2d->gbuf3d memcpy). Reuse tile3/gbuf3d across - ! 3D vars when nlevels matches (typical for ocean state). + tile2, self%vars(v)%gbuf_2d, is_root) + case (3) nlev = self%vars(v)%nlevels - if (is_root) then - if (allocated(gbuf3d)) then - if (size(gbuf3d, 3) /= nlev) deallocate(gbuf3d) - end if - if (.not. allocated(gbuf3d)) allocate(gbuf3d(self%nx_g, self%ny_g, nlev)) - else - ! Non-root dummy so the actual argument to mpp_gather is allocated. - if (.not. allocated(gbuf3d)) allocate(gbuf3d(1, 1, 1)) - end if if (allocated(tile3)) then if (size(tile3, 3) /= nlev) deallocate(tile3) end if if (.not. allocated(tile3)) allocate(tile3(nx_c, ny_c, nlev)) + if (is_root) then + if (.not. allocated(self%vars(v)%gbuf_3d)) & + allocate(self%vars(v)%gbuf_3d(self%nx_g, self%ny_g, nlev)) + else + if (.not. allocated(self%vars(v)%gbuf_3d)) allocate(self%vars(v)%gbuf_3d(1, 1, 1)) + end if tile3 = self%vars(v)%data3d(i_off : i_off + nx_c - 1, & j_off : j_off + ny_c - 1, :) + ! Single 3D mpp_gather per 3D var: one collective replaces nlevels 2D + ! gathers, and root receives the assembled global field directly. call mpp_gather(self%isc, self%iec, self%jsc, self%jec, nlev, pelist, & - tile3, gbuf3d, is_root) - if (is_root) then - call ncc(nf90_put_var(ncid, varids(v), gbuf3d, & - start=[1, 1, 1, 1], count=[self%nx_g, self%ny_g, nlev, 1]), & - 'put '//trim(self%vars(v)%name)) + tile3, self%vars(v)%gbuf_3d, is_root) + end select + end do + + if (allocated(tile2)) deallocate(tile2) + if (allocated(tile3)) deallocate(tile3) + if (allocated(pelist)) deallocate(pelist) +end subroutine writer_stage_gather + + +!============================================================================== +! Stage 3 (write): on self%root_pe, dump each var to disk via nf90_put_var. +! For 1D vars the source is data1d directly (global on every PE; root just +! reads its own copy). For 2D/3D the source is the gbuf_2d / gbuf_3d that +! stage_gather assembled. No MPI -> across multiple writers with different +! root_pes (ensemble path) this stage runs concurrently per PE. +!============================================================================== +subroutine writer_stage_write(self) + type(soca_io_writer), intent(inout) :: self + + integer :: v, nlev + logical :: is_root + + is_root = (mpp_pe() == self%root_pe) + if (.not. is_root) return + + do v = 1, self%nvars + select case (self%vars(v)%ndims) + case (1) + call ncc(nf90_put_var(self%ncid, self%varids(v), self%vars(v)%data1d, & + start=[1, 1], count=[size(self%vars(v)%data1d), 1]), & + 'put '//trim(self%vars(v)%name)) + case (2) + call ncc(nf90_put_var(self%ncid, self%varids(v), self%vars(v)%gbuf_2d, & + start=[1, 1, 1], count=[self%nx_g, self%ny_g, 1]), & + 'put '//trim(self%vars(v)%name)) + case (3) + nlev = self%vars(v)%nlevels + call ncc(nf90_put_var(self%ncid, self%varids(v), self%vars(v)%gbuf_3d, & + start=[1, 1, 1, 1], count=[self%nx_g, self%ny_g, nlev, 1]), & + 'put '//trim(self%vars(v)%name)) + end select + end do +end subroutine writer_stage_write + + +!============================================================================== +! Stage 4 (close): on self%root_pe close the file; everyone deallocates the +! inter-stage buffers and the writer's working set. No MPI -> concurrent across +! writers in the ensemble path. +!============================================================================== +subroutine writer_stage_close(self) + type(soca_io_writer), intent(inout) :: self + + integer :: v + logical :: is_root + + is_root = (mpp_pe() == self%root_pe) + + if (is_root .and. self%ncid >= 0) then + call ncc(nf90_close(self%ncid), 'nf90_close') + end if + self%ncid = -1 + + if (allocated(self%varids)) deallocate(self%varids) + if (allocated(self%x_axes)) deallocate(self%x_axes) + if (allocated(self%y_axes)) deallocate(self%y_axes) + if (allocated(self%z_axes)) deallocate(self%z_axes) + if (allocated(self%var_x_idx)) deallocate(self%var_x_idx) + if (allocated(self%var_y_idx)) deallocate(self%var_y_idx) + if (allocated(self%var_z_idx)) deallocate(self%var_z_idx) + self%nx_axes = 0; self%ny_axes = 0; self%nz_axes = 0 + + if (allocated(self%vars)) then + do v = 1, self%nvars + if (allocated(self%vars(v)%gbuf_2d)) deallocate(self%vars(v)%gbuf_2d) + if (allocated(self%vars(v)%gbuf_3d)) deallocate(self%vars(v)%gbuf_3d) + end do + deallocate(self%vars) + end if + self%nvars = 0 +end subroutine writer_stage_close + + +!============================================================================== +! Bulk-write orchestrator: take an array of pre-configured writers (each one +! init'd against the same domain with its own filename and rotated root_pe) +! and run all four stages across the whole set. +! +! Phase ordering: +! 1. define-all (each writer's root_pe creates its file; concurrent, no MPI) +! 2. gather-all (sync: per-writer per-var mpp_gather -- world collectives +! sequenced; async: all (writer, var) MPI_Igatherv batched +! into one wave + MPI_Waitall so the runtime can overlap) +! 3. write-all (each writer's root_pe nf90_put_var's its data; concurrent) +! 4. close-all (concurrent nf90_close) +! +! No inter-phase barriers are needed: the gather phase is self-synchronizing +! (collective + MPI_Waitall), and define/write/close are independent per +! writer root_pe, so each rank runs them at its own pace. +! +! Memory: each writer's root_pe holds one member's worth of every var's +! gbuf_2d/gbuf_3d between phases 2 and 3 -- peak ~global state size per writer. +!============================================================================== +subroutine soca_io_writers_commit_ensemble(writers) + type(soca_io_writer), intent(inout) :: writers(:) + + integer :: m + + if (size(writers) == 0) return + + ! All writers in a batch must share one FMS domain: the async gather + ! allgathers writers(1)'s compute-domain bounds once and reuses them to unpack + ! every member's tiles. A mismatched domain would silently corrupt the field. + do m = 2, size(writers) + if (.not. associated(writers(m)%domain, writers(1)%domain)) & + call mpp_error(FATAL, & + 'soca_io_writers_commit_ensemble: all writers in a batch must share one domain') + end do + + ! Phase 1: file create + dim/var defs. Concurrent across root_pes; no MPI. + do m = 1, size(writers) + call writer_stage_define(writers(m)) + end do + + ! Phase 2: gather. Sync path issues per-writer per-var mpp_gather (each is a + ! world-comm collective, so they serialize across writers). Async path + ! batches every (writer, var) MPI_Igatherv into one wave + MPI_Waitall. + if (async_mpi_enabled_) then + call writer_stage_gather_all_async(writers) + else + do m = 1, size(writers) + call writer_stage_gather(writers(m)) + end do + end if + + ! Phase 3: each root_pe writes its writer's data. No MPI -> concurrent disk. + do m = 1, size(writers) + call writer_stage_write(writers(m)) + end do + + ! Phase 4: each root_pe closes its file. No MPI -> concurrent. + do m = 1, size(writers) + call writer_stage_close(writers(m)) + end do +end subroutine soca_io_writers_commit_ensemble + + +!============================================================================== +! Async cross-writer gather: pack every PE's compute-domain tile for every +! (writer, 2D-or-3D var) into a per-(request) send buffer, post one +! MPI_Igatherv per (writer, var) targeted at writers(m)%root_pe, MPI_Waitall, +! then each root unpacks its received tiles into the var's gbuf. The runtime +! can overlap Igathers whose roots differ; sync mpp_gather would force one +! world-comm collective per gather, serializing the whole batch. +! +! Each 3D var ships as a single block (i fastest, then j, then k); 2D vars get +! one Igatherv each. Per-sender memory cost is the sum of (writers, vars) of +! compute-tile-size * nz_var bytes -- the whole batch in flight at once. +!============================================================================== +subroutine writer_stage_gather_all_async(writers) + type(soca_io_writer), intent(inout) :: writers(:) + + integer :: nprocs, ierr, r, m, v, idx, rq_count, max_reqs, req_idx, idx_done + integer :: my_isc, my_iec, my_jsc, my_jec, my_count2d + integer :: nz_v, full_count, isg, jsg + integer :: i, j, k, ig, jg, kk, mp_comm + integer, allocatable :: all_isc(:), all_iec(:), all_jsc(:), all_jec(:) + integer, allocatable :: tile_counts2d(:) + integer, allocatable :: reqs(:) + integer, allocatable :: req_writer(:), req_var(:) + ! Per-request count/displacement arrays. MPI_Igatherv is nonblocking, so the + ! count/displ arguments must stay valid (unmutated) until MPI_Waitall. Each + ! request has its own column rc(:,rq)/disp(:,rq) -- a single shared rc(:) + ! recomputed each loop iteration would be aliased across in-flight gathers. + integer, allocatable :: rc(:,:), disp(:,:) + type :: gather_buf_t + real(kind=kind_real), allocatable :: buf(:) + end type + type(gather_buf_t), allocatable :: sends(:), recvs(:) + logical :: is_root_for_this + + if (size(writers) == 0) return + call current_mpi_comm(mp_comm) + nprocs = mpp_npes() + + ! All writers in the ensemble share the same FMS domain (the geometry's), + ! so gather the compute-domain bounds once and reuse for every Igatherv. + call mpp_get_compute_domain(writers(1)%domain, my_isc, my_iec, my_jsc, my_jec) + call gather_compute_domains(writers(1)%domain, mp_comm, all_isc, all_iec, all_jsc, all_jec) + + allocate(tile_counts2d(nprocs)) + do r = 1, nprocs + tile_counts2d(r) = (all_iec(r) - all_isc(r) + 1) * (all_jec(r) - all_jsc(r) + 1) + end do + my_count2d = (my_iec - my_isc + 1) * (my_jec - my_jsc + 1) + + ! Allocate the gbufs on each writer's root_pe up front (one per 2D/3D var). + ! These buffers persist through stage_write and are freed in stage_close. + do m = 1, size(writers) + is_root_for_this = (mpp_pe() == writers(m)%root_pe) + if (.not. is_root_for_this) cycle + do v = 1, writers(m)%nvars + select case (writers(m)%vars(v)%ndims) + case (2) + if (.not. allocated(writers(m)%vars(v)%gbuf_2d)) & + allocate(writers(m)%vars(v)%gbuf_2d(writers(m)%nx_g, writers(m)%ny_g)) + case (3) + if (.not. allocated(writers(m)%vars(v)%gbuf_3d)) & + allocate(writers(m)%vars(v)%gbuf_3d( & + writers(m)%nx_g, writers(m)%ny_g, writers(m)%vars(v)%nlevels)) + end select + end do + end do + + ! Count requests: one per (writer, gatherable var). 1D vars contribute no + ! request (data1d is global on every PE). + max_reqs = 0 + do m = 1, size(writers) + do v = 1, writers(m)%nvars + if (writers(m)%vars(v)%ndims >= 2) max_reqs = max_reqs + 1 + end do + end do + if (max_reqs == 0) goto 9999 + + allocate(reqs(max_reqs)) + allocate(req_writer(max_reqs), req_var(max_reqs)) + allocate(sends(max_reqs), recvs(max_reqs)) + allocate(rc(nprocs, max_reqs), disp(nprocs, max_reqs)) + rq_count = 0 + + ! Phase A/B: pack send + post Igatherv per (writer, var). + do m = 1, size(writers) + is_root_for_this = (mpp_pe() == writers(m)%root_pe) + do v = 1, writers(m)%nvars + if (writers(m)%vars(v)%ndims < 2) cycle + rq_count = rq_count + 1 + req_writer(rq_count) = m + req_var(rq_count) = v + + if (writers(m)%vars(v)%ndims == 2) then + nz_v = 1 + else + nz_v = writers(m)%vars(v)%nlevels + end if + + do r = 1, nprocs + rc(r, rq_count) = tile_counts2d(r) * nz_v + end do + disp(1, rq_count) = 0 + do r = 2, nprocs + disp(r, rq_count) = disp(r-1, rq_count) + rc(r-1, rq_count) + end do + + allocate(sends(rq_count)%buf(my_count2d * nz_v)) + idx = 0 + if (writers(m)%vars(v)%ndims == 2) then + do j = 1, my_jec - my_jsc + 1 + do i = 1, my_iec - my_isc + 1 + idx = idx + 1 + sends(rq_count)%buf(idx) = writers(m)%vars(v)%data2d( & + writers(m)%isc - writers(m)%isd + i, & + writers(m)%jsc - writers(m)%jsd + j) + end do + end do + else + do k = 1, nz_v + do j = 1, my_jec - my_jsc + 1 + do i = 1, my_iec - my_isc + 1 + idx = idx + 1 + sends(rq_count)%buf(idx) = writers(m)%vars(v)%data3d( & + writers(m)%isc - writers(m)%isd + i, & + writers(m)%jsc - writers(m)%jsd + j, k) + end do + end do + end do + end if + + if (is_root_for_this) then + full_count = 0 + do r = 1, nprocs + full_count = full_count + rc(r, rq_count) + end do + allocate(recvs(rq_count)%buf(full_count)) + else + ! Non-root recvbuf is not significant to MPI_Igatherv, but the actual + ! argument must still be an allocated array (1-elem dummy mirrors the + ! reader scatter path's sendbuf handling). + allocate(recvs(rq_count)%buf(1)) + end if + + call MPI_Igatherv(sends(rq_count)%buf, my_count2d * nz_v, MPI_DOUBLE_PRECISION, & + recvs(rq_count)%buf, rc(:, rq_count), disp(:, rq_count), MPI_DOUBLE_PRECISION, & + writers(m)%root_pe, mp_comm, reqs(rq_count), ierr) + end do + end do + + ! Phase C/D: complete the gathers one at a time. As each Igatherv finishes, + ! its send buffer is done with, and on the target root the recv buffer can be + ! unpacked into the per-var gbuf and freed immediately. So the transient + ! buffer high-water is the in-flight set, not all rq_count requests at once + ! (the full recv buffers are global-field sized -- holding every member's at + ! once doubled peak memory). All requests are still posted up front, so the + ! runtime overlap is unchanged. The gbufs persist to stage_write. + do req_idx = 1, rq_count + call MPI_Waitany(rq_count, reqs, idx_done, MPI_STATUS_IGNORE, ierr) + if (idx_done == MPI_UNDEFINED) exit + if (allocated(sends(idx_done)%buf)) deallocate(sends(idx_done)%buf) + m = req_writer(idx_done) + v = req_var(idx_done) + if (mpp_pe() == writers(m)%root_pe) then + isg = writers(m)%isg + jsg = writers(m)%jsg + if (writers(m)%vars(v)%ndims == 2) then + idx = 0 + do r = 1, nprocs + do jg = all_jsc(r), all_jec(r) + do ig = all_isc(r), all_iec(r) + idx = idx + 1 + writers(m)%vars(v)%gbuf_2d(ig - isg + 1, jg - jsg + 1) = recvs(idx_done)%buf(idx) + end do + end do + end do + else + idx = 0 + do r = 1, nprocs + do kk = 1, writers(m)%vars(v)%nlevels + do jg = all_jsc(r), all_jec(r) + do ig = all_isc(r), all_iec(r) + idx = idx + 1 + writers(m)%vars(v)%gbuf_3d(ig - isg + 1, jg - jsg + 1, kk) = recvs(idx_done)%buf(idx) + end do + end do + end do + end do end if end if + if (allocated(recvs(idx_done)%buf)) deallocate(recvs(idx_done)%buf) end do - ! Phase 3: close. - if (is_root) then - call ncc(nf90_close(ncid), 'nf90_close') - deallocate(varids) - end if - deallocate(gbuf2d, tile2, pelist, x_axes, y_axes, z_axes, var_x_idx, var_y_idx, var_z_idx) - if (allocated(gbuf3d)) deallocate(gbuf3d) - if (allocated(tile3)) deallocate(tile3) + deallocate(reqs, req_writer, req_var, sends, recvs, rc, disp) - ! drop pointer entries; caller's data is unaffected - if (allocated(self%vars)) deallocate(self%vars) - self%nvars = 0 -end subroutine writer_commit +9999 continue + deallocate(all_isc, all_iec, all_jsc, all_jec, tile_counts2d) +end subroutine writer_stage_gather_all_async !============================================================================== ! Reader: init / enqueue_* / commit. Caller buffer stays in place; enqueue ! records a pointer, commit fills the compute-domain interior. Halos are left ! untouched -- the caller refreshes them via mpp_update_domains (same as FMS). +! +! Optional reader_pe sets the source PE for the scatter path (defaults to +! mpp_root_pe() in single-state mode); the ensemble orchestrator pre-sets a +! rotated reader_pe per member so different files are read concurrently. !============================================================================== -subroutine reader_init(self, domain, filename) +subroutine reader_init(self, domain, filename, reader_pe) class(soca_io_reader), intent(inout) :: self type(domain2D), target, intent(in) :: domain character(len=*), intent(in) :: filename + integer, optional, intent(in) :: reader_pe self%filename = filename self%domain => domain @@ -443,6 +851,12 @@ subroutine reader_init(self, domain, filename) if (allocated(self%vars)) deallocate(self%vars) allocate(self%vars(64)) self%nvars = 0 + + if (present(reader_pe)) then + self%reader_pe = reader_pe + else + self%reader_pe = mpp_root_pe() + end if end subroutine reader_init @@ -528,20 +942,27 @@ end subroutine grow_reader_if_needed !============================================================================== -! Read all enqueued vars. Every PE opens the file NF90_NOWRITE and pulls only -! its compute-domain tile via nf90_get_var(start, count) -- mirrors FMS's -! MPP_READ_2DDECOMP: no PE-0 bottleneck, no mpp_broadcast, N parallel reads. -! Classic / 64-bit-offset netcdf allows concurrent read-only opens; library -! state is process-local. 1D vars also read independently on every PE. +! Read all enqueued vars. Single-state path dispatches on the YAML knob +! geometry.io.'single state read': 'strided' (default, every PE reads its own +! tile direct) or 'scatter' (reader_pe reads the whole field then MPI_Scatterv +! to tile owners). The strided path mirrors FMS MPP_READ_2DDECOMP -- no PE-0 +! bottleneck, N parallel reads. The scatter path concentrates the disk read +! on one PE: useful when 200-500 PEs concurrently striding the same file +! stresses the filesystem. !============================================================================== subroutine reader_commit(self) class(soca_io_reader), intent(inout) :: self - call commit_reader_strided(self) - - ! release pointers; caller's buffers untouched (they hold the read data) - if (allocated(self%vars)) deallocate(self%vars) - self%nvars = 0 + if (single_state_read_scatter_) then + call reader_stage_read(self) + call reader_stage_distribute(self) + call reader_stage_close(self) + else + call commit_reader_strided(self) + ! release pointers; caller's buffers untouched (they hold the read data) + if (allocated(self%vars)) deallocate(self%vars) + self%nvars = 0 + end if end subroutine reader_commit @@ -616,119 +1037,448 @@ end subroutine commit_reader_strided !============================================================================== -! PE-0 read + per-PE scatter. PE 0 nf90_get_var's the global field; mpp_scatter -! sends each PE its compute-domain slice (1D vars are broadcast). Mirrors FMS -! 2024.02 fms_netcdf_domain_io.F90:domain_read_3d. +! Stage 1 (read): on self%reader_pe, pull the whole global field for every +! enqueued var into a per-var staged buffer. No MPI -> across multiple readers +! with different reader_pes (ensemble path) this stage runs concurrently per +! PE (independent files, independent disk I/O). ! -! TODO: currently unused -- single-state reads use commit_reader_strided. Will -! be exercised by parallel-ensemble I/O, where one reader PE per member -! scatters that member to its compute-PE group. +! 1D vars: read directly into data1d on reader_pe (reader_stage_distribute +! broadcasts it). 2D/3D/4D: alloc gbuf_{2,3,4}d on reader_pe; non-reader PEs +! hold nothing (the scatter recvbuf is allocated on the fly in distribute). !============================================================================== -subroutine commit_reader_scatter(self) - class(soca_io_reader), intent(inout) :: self +subroutine reader_stage_read(self) + type(soca_io_reader), intent(inout) :: self - integer :: ncid, v, n, n3, n4, k4 - integer :: nx_c, ny_c, i_off, j_off - integer :: is_f, ie_f, js_f, je_f ! PE tile in 1-based file-space indices - integer, allocatable :: pelist(:) - real(kind=kind_real), allocatable :: gbuf2(:,:), gbuf3(:,:,:), gbuf4(:,:,:,:) - real(kind=kind_real), allocatable :: tile2(:,:), tile3(:,:,:) + integer :: ncid, v, n, n3, n4 logical :: is_root - is_root = (mpp_pe() == mpp_root_pe()) - call mpi_pelist(pelist) + is_root = (mpp_pe() == self%reader_pe) + if (.not. is_root) return - ncid = -1 - if (is_root) call ncc(nf90_open(self%filename, NF90_NOWRITE, ncid), & + call ncc(nf90_open(self%filename, NF90_NOWRITE, ncid), & 'nf90_open '//trim(self%filename)) - nx_c = self%iec - self%isc + 1 - ny_c = self%jec - self%jsc + 1 - i_off = self%isc - self%isd + 1 - j_off = self%jsc - self%jsd + 1 - ! Map PE compute indices (isg-based) to 1-based file/global-buffer indices. - ! FMS 2025.02 removed the ishift/jshift optional args from mpp_scatter, so - ! we pre-apply the shift in the indices we pass. - is_f = self%isc - self%isg + 1 - ie_f = self%iec - self%isg + 1 - js_f = self%jsc - self%jsg + 1 - je_f = self%jec - self%jsg + 1 - do v = 1, self%nvars select case (self%vars(v)%ndims) case (1) - ! Global on every PE: PE 0 reads, broadcasts. n = size(self%vars(v)%data1d) - if (is_root) call read_var_strided(ncid, self%vars(v)%name, & + call read_var_strided(ncid, self%vars(v)%name, & 1, 1, n, 1, dst1=self%vars(v)%data1d) - call mpp_broadcast(self%vars(v)%data1d, n, mpp_root_pe()) - case (2) - allocate(tile2(nx_c, ny_c)) - if (is_root) then - allocate(gbuf2(self%nx_g, self%ny_g)) - call read_var_strided(ncid, self%vars(v)%name, & - 1, 1, self%nx_g, self%ny_g, dst2=gbuf2) - else - allocate(gbuf2(1, 1)) ! dummy: mpp_scatter only reads input_data on root - end if - call mpp_scatter(is_f, ie_f, js_f, je_f, & - pelist, tile2, gbuf2, is_root) - deallocate(gbuf2) - self%vars(v)%data2d(i_off : i_off + nx_c - 1, & - j_off : j_off + ny_c - 1) = tile2 - deallocate(tile2) - + if (.not. allocated(self%vars(v)%gbuf_2d)) & + allocate(self%vars(v)%gbuf_2d(self%nx_g, self%ny_g)) + call read_var_strided(ncid, self%vars(v)%name, & + 1, 1, self%nx_g, self%ny_g, dst2=self%vars(v)%gbuf_2d) case (3) n3 = size(self%vars(v)%data3d, 3) - allocate(tile3(nx_c, ny_c, n3)) - if (is_root) then - allocate(gbuf3(self%nx_g, self%ny_g, n3)) - call read_var_strided(ncid, self%vars(v)%name, & - 1, 1, self%nx_g, self%ny_g, dst3=gbuf3) - else - allocate(gbuf3(1, 1, 1)) - end if - call mpp_scatter(is_f, ie_f, js_f, je_f, n3, & - pelist, tile3, gbuf3, is_root) - deallocate(gbuf3) - self%vars(v)%data3d(i_off : i_off + nx_c - 1, & - j_off : j_off + ny_c - 1, :) = tile3 - deallocate(tile3) - + if (.not. allocated(self%vars(v)%gbuf_3d)) & + allocate(self%vars(v)%gbuf_3d(self%nx_g, self%ny_g, n3)) + call read_var_strided(ncid, self%vars(v)%name, & + 1, 1, self%nx_g, self%ny_g, dst3=self%vars(v)%gbuf_3d) case (4) - ! mpp_scatter is 2D/3D only; loop the outer (4th) dim and call 3D scatter. n3 = size(self%vars(v)%data4d, 3) n4 = size(self%vars(v)%data4d, 4) - allocate(tile3(nx_c, ny_c, n3)) + if (.not. allocated(self%vars(v)%gbuf_4d)) & + allocate(self%vars(v)%gbuf_4d(self%nx_g, self%ny_g, n3, n4)) + call read_var_strided(ncid, self%vars(v)%name, & + 1, 1, self%nx_g, self%ny_g, dst4=self%vars(v)%gbuf_4d) + end select + end do + + call ncc(nf90_close(ncid), 'nf90_close '//trim(self%filename)) +end subroutine reader_stage_read + + +!============================================================================== +! Stage 2 (distribute, sync): per-var send the relevant slice from reader_pe +! to every PE in the world communicator. 1D vars are MPI_Bcast. 2D/3D/4D vars +! are packed rank-ordered on reader_pe and shipped via MPI_Scatterv; each +! receiver unpacks its compute tile into data2d/data3d/data4d. +! +! Per-var Allgather of the compute-domain bounds: needed because the pack +! order on reader_pe must match the unpack order on every receiver. The +! Allgather happens once per (reader, var); for the async batched path the +! same bounds are reused across all vars (see reader_stage_distribute_all_async). +!============================================================================== +subroutine reader_stage_distribute(self) + type(soca_io_reader), intent(inout) :: self + + integer :: nprocs, ierr, v, r, i, j, k, k4, n3, n4, idx, full_count + integer :: my_isc, my_iec, my_jsc, my_jec, my_count2d + integer :: nx_c, ny_c, i_off, j_off, mp_comm + integer, allocatable :: all_isc(:), all_iec(:), all_jsc(:), all_jec(:) + integer, allocatable :: tile_counts2d(:), rc(:), disp(:) + real(kind=kind_real), allocatable :: sendbuf(:), recvbuf(:) + logical :: is_root + + is_root = (mpp_pe() == self%reader_pe) + nprocs = mpp_npes() + call current_mpi_comm(mp_comm) + + call mpp_get_compute_domain(self%domain, my_isc, my_iec, my_jsc, my_jec) + call gather_compute_domains(self%domain, mp_comm, all_isc, all_iec, all_jsc, all_jec) + + allocate(tile_counts2d(nprocs)) + do r = 1, nprocs + tile_counts2d(r) = (all_iec(r) - all_isc(r) + 1) * (all_jec(r) - all_jsc(r) + 1) + end do + my_count2d = (my_iec - my_isc + 1) * (my_jec - my_jsc + 1) + + nx_c = self%iec - self%isc + 1 + ny_c = self%jec - self%jsc + 1 + i_off = self%isc - self%isd + 1 + j_off = self%jsc - self%jsd + 1 + + allocate(rc(nprocs), disp(nprocs)) + + do v = 1, self%nvars + select case (self%vars(v)%ndims) + case (1) + ! 1D: every PE wants the same global vector. Bcast directly into data1d. + call MPI_Bcast(self%vars(v)%data1d, size(self%vars(v)%data1d), & + MPI_DOUBLE_PRECISION, self%reader_pe, mp_comm, ierr) + + case (2, 3, 4) + if (self%vars(v)%ndims == 2) then + n3 = 1; n4 = 1 + else if (self%vars(v)%ndims == 3) then + n3 = size(self%vars(v)%data3d, 3); n4 = 1 + else + n3 = size(self%vars(v)%data4d, 3); n4 = size(self%vars(v)%data4d, 4) + end if + + do r = 1, nprocs + rc(r) = tile_counts2d(r) * n3 * n4 + end do + disp(1) = 0 + do r = 2, nprocs + disp(r) = disp(r-1) + rc(r-1) + end do + if (is_root) then - allocate(gbuf4(self%nx_g, self%ny_g, n3, n4)) - call read_var_strided(ncid, self%vars(v)%name, & - 1, 1, self%nx_g, self%ny_g, dst4=gbuf4) + full_count = 0 + do r = 1, nprocs + full_count = full_count + rc(r) + end do + allocate(sendbuf(full_count)) + idx = 0 + ! Pack rank-by-rank, k4 fastest-outer, then k, then j, then i (matches unpack) + do r = 1, nprocs + do k4 = 1, n4 + do k = 1, n3 + do j = all_jsc(r), all_jec(r) + do i = all_isc(r), all_iec(r) + idx = idx + 1 + select case (self%vars(v)%ndims) + case (2) + sendbuf(idx) = self%vars(v)%gbuf_2d(i - self%isg + 1, j - self%jsg + 1) + case (3) + sendbuf(idx) = self%vars(v)%gbuf_3d(i - self%isg + 1, j - self%jsg + 1, k) + case (4) + sendbuf(idx) = self%vars(v)%gbuf_4d(i - self%isg + 1, j - self%jsg + 1, k, k4) + end select + end do + end do + end do + end do + end do else - allocate(gbuf3(1, 1, 1)) ! dummy for non-root in the 3D scatter call + allocate(sendbuf(1)) end if + + allocate(recvbuf(my_count2d * n3 * n4)) + call MPI_Scatterv(sendbuf, rc, disp, MPI_DOUBLE_PRECISION, & + recvbuf, my_count2d * n3 * n4, MPI_DOUBLE_PRECISION, & + self%reader_pe, mp_comm, ierr) + + ! Unpack recvbuf into the caller's data buffer compute slice. Halos + ! untouched (caller refreshes via mpp_update_domains). + idx = 0 do k4 = 1, n4 - if (is_root) then - call mpp_scatter(is_f, ie_f, js_f, je_f, n3, & - pelist, tile3, gbuf4(:,:,:,k4), is_root) - else - call mpp_scatter(is_f, ie_f, js_f, je_f, n3, & - pelist, tile3, gbuf3, is_root) - end if - self%vars(v)%data4d(i_off : i_off + nx_c - 1, & - j_off : j_off + ny_c - 1, :, k4) = tile3 + do k = 1, n3 + do j = 1, ny_c + do i = 1, nx_c + idx = idx + 1 + select case (self%vars(v)%ndims) + case (2) + self%vars(v)%data2d(i_off + i - 1, j_off + j - 1) = recvbuf(idx) + case (3) + self%vars(v)%data3d(i_off + i - 1, j_off + j - 1, k) = recvbuf(idx) + case (4) + self%vars(v)%data4d(i_off + i - 1, j_off + j - 1, k, k4) = recvbuf(idx) + end select + end do + end do + end do end do - if (is_root) deallocate(gbuf4) - if (allocated(gbuf3)) deallocate(gbuf3) - deallocate(tile3) + + deallocate(sendbuf, recvbuf) end select end do - if (is_root) call ncc(nf90_close(ncid), 'nf90_close '//trim(self%filename)) + deallocate(all_isc, all_iec, all_jsc, all_jec, tile_counts2d, rc, disp) +end subroutine reader_stage_distribute - if (allocated(pelist)) deallocate(pelist) -end subroutine commit_reader_scatter + +!============================================================================== +! Stage 3 (close): drop the staged global buffers and the reader's working set. +!============================================================================== +subroutine reader_stage_close(self) + type(soca_io_reader), intent(inout) :: self + integer :: v + if (allocated(self%vars)) then + do v = 1, self%nvars + if (allocated(self%vars(v)%gbuf_2d)) deallocate(self%vars(v)%gbuf_2d) + if (allocated(self%vars(v)%gbuf_3d)) deallocate(self%vars(v)%gbuf_3d) + if (allocated(self%vars(v)%gbuf_4d)) deallocate(self%vars(v)%gbuf_4d) + end do + deallocate(self%vars) + end if + self%nvars = 0 +end subroutine reader_stage_close + + +!============================================================================== +! Bulk-read orchestrator. Two modes governed by the geometry.io knob +! 'ensemble read': +! - strided: loop readers, each one does commit_reader_strided (every PE +! reads its compute-domain tile direct from that member's file). All +! members done sequentially; per-member read is parallel across PEs. +! - scatter: read-all -> distribute-all. Phase 1: every reader_pe reads its +! member's whole global field into staged buffers, no MPI -> reader PEs +! hit different files concurrently. Phase 2: per-reader MPI_Scatterv from +! reader_pe (sync: per-reader collectives serialize; async: all +! (reader, var) Iscatterv + Waitall batched). +!============================================================================== +subroutine soca_io_readers_commit_ensemble(readers) + type(soca_io_reader), intent(inout) :: readers(:) + + integer :: m + + if (size(readers) == 0) return + + ! All readers in a batch must share one FMS domain (see the writer + ! orchestrator note): the async scatter allgathers readers(1)'s + ! compute-domain bounds once and reuses them for every member. + do m = 2, size(readers) + if (.not. associated(readers(m)%domain, readers(1)%domain)) & + call mpp_error(FATAL, & + 'soca_io_readers_commit_ensemble: all readers in a batch must share one domain') + end do + + if (.not. ensemble_read_scatter_) then + ! Strided ensemble path: loop members, run the existing per-PE strided + ! reader for each. Members are sequenced; within each member every PE + ! reads its own tile. + do m = 1, size(readers) + call commit_reader_strided(readers(m)) + if (allocated(readers(m)%vars)) deallocate(readers(m)%vars) + readers(m)%nvars = 0 + end do + return + end if + + ! Scatter ensemble path. + ! Phase 1: each reader_pe reads its file. No MPI -> reads concurrent. + do m = 1, size(readers) + call reader_stage_read(readers(m)) + end do + + ! Phase 2: scatter. Sync per-reader serializes; async batches all + ! (reader, var) Iscatterv/Ibcast + Waitall. + if (async_mpi_enabled_) then + call reader_stage_distribute_all_async(readers) + else + do m = 1, size(readers) + call reader_stage_distribute(readers(m)) + end do + end if + + ! Phase 3: free staged buffers + drop working set. + do m = 1, size(readers) + call reader_stage_close(readers(m)) + end do +end subroutine soca_io_readers_commit_ensemble + + +!============================================================================== +! Async cross-reader scatter: per (reader, var) pack the relevant rank-ordered +! slices on the reader_pe, post one MPI_Iscatterv per var (MPI_Ibcast for 1D), +! MPI_Waitall, then each receiver unpacks its recvbuf into the caller buffer +! compute slice. Sync mpp/MPI_Scatterv serializes across readers because each +! is a world-comm collective; async batches them so the runtime can overlap +! scatters whose roots differ. +!============================================================================== +subroutine reader_stage_distribute_all_async(readers) + type(soca_io_reader), intent(inout) :: readers(:) + + integer :: nprocs, ierr, r, m, v, idx, full_count, idx_done + integer :: my_isc, my_iec, my_jsc, my_jec, my_count2d, mp_comm + integer :: n3, n4, i, j, k, k4, nx_c, ny_c, i_off, j_off + integer, allocatable :: all_isc(:), all_iec(:), all_jsc(:), all_jec(:) + integer, allocatable :: tile_counts2d(:) + ! Per-request count/displacement columns: MPI_Iscatterv is nonblocking, so + ! these must stay valid until MPI_Waitall (one column per in-flight request). + integer, allocatable :: rc(:,:), disp(:,:) + integer :: rq_count, max_reqs, ux + integer, allocatable :: reqs(:), req_reader(:), req_var(:) + type :: scatter_buf_t + real(kind=kind_real), allocatable :: buf(:) + end type + type(scatter_buf_t), allocatable :: sends(:), recvs(:) + logical :: is_root_for_this + + if (size(readers) == 0) return + call current_mpi_comm(mp_comm) + nprocs = mpp_npes() + + call mpp_get_compute_domain(readers(1)%domain, my_isc, my_iec, my_jsc, my_jec) + call gather_compute_domains(readers(1)%domain, mp_comm, all_isc, all_iec, all_jsc, all_jec) + + allocate(tile_counts2d(nprocs)) + do r = 1, nprocs + tile_counts2d(r) = (all_iec(r) - all_isc(r) + 1) * (all_jec(r) - all_jsc(r) + 1) + end do + my_count2d = (my_iec - my_isc + 1) * (my_jec - my_jsc + 1) + + ! Count requests: one per (reader, var). 1D and 2D/3D/4D all count. + max_reqs = 0 + do m = 1, size(readers) + max_reqs = max_reqs + readers(m)%nvars + end do + if (max_reqs == 0) goto 9999 + + allocate(reqs(max_reqs), req_reader(max_reqs), req_var(max_reqs)) + allocate(sends(max_reqs), recvs(max_reqs)) + allocate(rc(nprocs, max_reqs), disp(nprocs, max_reqs)) + rq_count = 0 + + ! Phase A: pack send buffers + post Iscatterv / Ibcast. + do m = 1, size(readers) + is_root_for_this = (mpp_pe() == readers(m)%reader_pe) + do v = 1, readers(m)%nvars + rq_count = rq_count + 1 + req_reader(rq_count) = m + req_var(rq_count) = v + + if (readers(m)%vars(v)%ndims == 1) then + ! 1D: post a single MPI_Ibcast of data1d (every PE wants the same). + call MPI_Ibcast(readers(m)%vars(v)%data1d, size(readers(m)%vars(v)%data1d), & + MPI_DOUBLE_PRECISION, readers(m)%reader_pe, mp_comm, reqs(rq_count), ierr) + cycle + end if + + if (readers(m)%vars(v)%ndims == 2) then + n3 = 1; n4 = 1 + else if (readers(m)%vars(v)%ndims == 3) then + n3 = size(readers(m)%vars(v)%data3d, 3); n4 = 1 + else + n3 = size(readers(m)%vars(v)%data4d, 3); n4 = size(readers(m)%vars(v)%data4d, 4) + end if + + do r = 1, nprocs + rc(r, rq_count) = tile_counts2d(r) * n3 * n4 + end do + disp(1, rq_count) = 0 + do r = 2, nprocs + disp(r, rq_count) = disp(r-1, rq_count) + rc(r-1, rq_count) + end do + + if (is_root_for_this) then + full_count = 0 + do r = 1, nprocs + full_count = full_count + rc(r, rq_count) + end do + allocate(sends(rq_count)%buf(full_count)) + idx = 0 + do r = 1, nprocs + do k4 = 1, n4 + do k = 1, n3 + do j = all_jsc(r), all_jec(r) + do i = all_isc(r), all_iec(r) + idx = idx + 1 + select case (readers(m)%vars(v)%ndims) + case (2) + sends(rq_count)%buf(idx) = readers(m)%vars(v)%gbuf_2d( & + i - readers(m)%isg + 1, j - readers(m)%jsg + 1) + case (3) + sends(rq_count)%buf(idx) = readers(m)%vars(v)%gbuf_3d( & + i - readers(m)%isg + 1, j - readers(m)%jsg + 1, k) + case (4) + sends(rq_count)%buf(idx) = readers(m)%vars(v)%gbuf_4d( & + i - readers(m)%isg + 1, j - readers(m)%jsg + 1, k, k4) + end select + end do + end do + end do + end do + end do + else + allocate(sends(rq_count)%buf(1)) + end if + + allocate(recvs(rq_count)%buf(my_count2d * n3 * n4)) + call MPI_Iscatterv(sends(rq_count)%buf, rc(:, rq_count), disp(:, rq_count), MPI_DOUBLE_PRECISION, & + recvs(rq_count)%buf, my_count2d * n3 * n4, MPI_DOUBLE_PRECISION, & + readers(m)%reader_pe, mp_comm, reqs(rq_count), ierr) + end do + end do + + ! Phase B/C: complete the scatters one at a time. As each Iscatterv (or the + ! 1D Ibcast) finishes, the root's send buffer is done with and is freed, and + ! the receiver unpacks its recv buffer into the caller compute slice and frees + ! it. The send buffers are global-field sized on the root (packed from gbuf), + ! so freeing per-completion instead of holding all of them until one Waitall + ! keeps the transient high-water at the in-flight set, not every member at + ! once. Requests are still all posted up front, so the overlap is unchanged. + do idx = 1, rq_count + call MPI_Waitany(rq_count, reqs, idx_done, MPI_STATUS_IGNORE, ierr) + if (idx_done == MPI_UNDEFINED) exit + ! 1D vars were Ibcast in-place with no send/recv buffers; the guards skip. + if (allocated(sends(idx_done)%buf)) deallocate(sends(idx_done)%buf) + m = req_reader(idx_done) + v = req_var(idx_done) + if (readers(m)%vars(v)%ndims == 1) cycle + + if (readers(m)%vars(v)%ndims == 2) then + n3 = 1; n4 = 1 + else if (readers(m)%vars(v)%ndims == 3) then + n3 = size(readers(m)%vars(v)%data3d, 3); n4 = 1 + else + n3 = size(readers(m)%vars(v)%data4d, 3); n4 = size(readers(m)%vars(v)%data4d, 4) + end if + + nx_c = readers(m)%iec - readers(m)%isc + 1 + ny_c = readers(m)%jec - readers(m)%jsc + 1 + i_off = readers(m)%isc - readers(m)%isd + 1 + j_off = readers(m)%jsc - readers(m)%jsd + 1 + + ux = 0 + do k4 = 1, n4 + do k = 1, n3 + do j = 1, ny_c + do i = 1, nx_c + ux = ux + 1 + select case (readers(m)%vars(v)%ndims) + case (2) + readers(m)%vars(v)%data2d(i_off + i - 1, j_off + j - 1) = recvs(idx_done)%buf(ux) + case (3) + readers(m)%vars(v)%data3d(i_off + i - 1, j_off + j - 1, k) = recvs(idx_done)%buf(ux) + case (4) + readers(m)%vars(v)%data4d(i_off + i - 1, j_off + j - 1, k, k4) = recvs(idx_done)%buf(ux) + end select + end do + end do + end do + end do + if (allocated(recvs(idx_done)%buf)) deallocate(recvs(idx_done)%buf) + end do + + deallocate(reqs, req_reader, req_var, sends, recvs, rc, disp) + +9999 continue + deallocate(all_isc, all_iec, all_jsc, all_jec, tile_counts2d) +end subroutine reader_stage_distribute_all_async @@ -959,6 +1709,53 @@ subroutine mpi_pelist(pelist) end subroutine mpi_pelist +!============================================================================== +! Fetch the MPI communicator handle that the current mpp pelist sits on. Used +! by the async ensemble I/O paths so MPI_I{gatherv,scatterv,bcast} run on the +! same comm as mpp's collectives (== the geometry's f_comm). Using +! MPI_COMM_WORLD here would deadlock in the LETKF "nens per MPI task" split +! where each task has its own size-1 mpp world. +!============================================================================== +subroutine current_mpi_comm(mp_comm) + integer, intent(out) :: mp_comm + integer, allocatable :: pelist(:) + character(len=128) :: name + ! mpp_get_current_pelist requires a pelist out-arg sized mpp_npes(); the async + ! I/O callers only need the communicator handle, so it stays local here. + allocate(pelist(mpp_npes())) + call mpp_get_current_pelist(pelist, name, mp_comm) + deallocate(pelist) +end subroutine current_mpi_comm + + +!============================================================================== +! Allgather every PE's compute-domain bounds (isc/iec/jsc/jec) in mp_comm rank +! order. One packed collective replaces the four separate Allgathers the bulk +! read/write paths used to issue; results are indexed by rank, so all_isc(r) is +! rank (r-1)'s compute-domain start -- matching the rank-ordered Igatherv / +! Iscatterv buffer layout. The bounds are static, but the gather is cheap and +! avoids any cross-call caching hazard with multiple geometries. +!============================================================================== +subroutine gather_compute_domains(domain, mp_comm, all_isc, all_iec, all_jsc, all_jec) + type(domain2D), intent(in) :: domain + integer, intent(in) :: mp_comm + integer, allocatable, intent(out) :: all_isc(:), all_iec(:), all_jsc(:), all_jec(:) + integer :: nprocs, ierr, my(4) + integer, allocatable :: gathered(:) + + nprocs = mpp_npes() + call mpp_get_compute_domain(domain, my(1), my(2), my(3), my(4)) + allocate(gathered(4 * nprocs)) + call MPI_Allgather(my, 4, MPI_INTEGER, gathered, 4, MPI_INTEGER, mp_comm, ierr) + allocate(all_isc(nprocs), all_iec(nprocs), all_jsc(nprocs), all_jec(nprocs)) + all_isc = gathered(1::4) + all_iec = gathered(2::4) + all_jsc = gathered(3::4) + all_jec = gathered(4::4) + deallocate(gathered) +end subroutine gather_compute_domains + + !============================================================================== ! Netcdf error check. Aborts on error with mpp_error(FATAL, ...). !============================================================================== @@ -971,6 +1768,103 @@ subroutine ncc(status, where) end subroutine ncc +!============================================================================== +! Resolve the module-level I/O dispatch knobs from the geometry YAML. Looked-up +! keys (all optional, defaults retained on miss): +! io: +! ensemble read: scatter | strided (default scatter) +! single state read: strided | scatter (default strided) +! async mpi: true | false (default true) +! ensemble batch size: N (default 0 = single batch) +! Called once from soca_geom_init with the geometry's fckit_configuration; the +! resolved values persist module-level for the rest of the run. +!============================================================================== +subroutine soca_io_config_from_yaml(f_conf) + type(fckit_configuration), intent(in) :: f_conf + type(fckit_configuration) :: io_conf + character(len=:), allocatable :: sval + logical :: lval, ok + integer :: ival + + if (.not. f_conf%has("io")) return + ok = f_conf%get("io", io_conf) + if (.not. ok) return + + if (io_conf%has("ensemble read")) then + ok = io_conf%get("ensemble read", sval) + if (ok) then + select case (trim(sval)) + case ("scatter"); ensemble_read_scatter_ = .true. + case ("strided"); ensemble_read_scatter_ = .false. + case default + call mpp_error(FATAL, "soca_io_mod: geometry.io.'ensemble read' must be " // & + "'scatter' or 'strided', got '" // trim(sval) // "'") + end select + end if + end if + + if (io_conf%has("single state read")) then + ok = io_conf%get("single state read", sval) + if (ok) then + select case (trim(sval)) + case ("scatter"); single_state_read_scatter_ = .true. + case ("strided"); single_state_read_scatter_ = .false. + case default + call mpp_error(FATAL, "soca_io_mod: geometry.io.'single state read' must be " // & + "'scatter' or 'strided', got '" // trim(sval) // "'") + end select + end if + end if + + if (io_conf%has("async mpi")) then + ok = io_conf%get("async mpi", lval) + if (ok) async_mpi_enabled_ = lval + end if + + if (io_conf%has("ensemble batch size")) then + ok = io_conf%get("ensemble batch size", ival) + if (ok) then + if (ival < 0) then + call mpp_error(FATAL, "soca_io_mod: geometry.io.'ensemble batch size' must be " // & + ">= 0 (0 = single batch over all members)") + end if + ensemble_batch_size_ = ival + end if + end if +end subroutine soca_io_config_from_yaml + + +!============================================================================== +! Public getter for the ensemble I/O batch-size knob (0 = single batch over all +! members). Consumed by fields_mod to wave the ensemble read/write in batches. +! (ensemble_read_scatter_ / single_state_read_scatter_ / async_mpi_enabled_ are +! read directly within this module, so they need no public getters.) +!============================================================================== +integer function soca_io_ensemble_batch_size() + soca_io_ensemble_batch_size = ensemble_batch_size_ +end function soca_io_ensemble_batch_size + + +!============================================================================== +! Strided writer/reader-PE rotation: map a 1-based member index m to a rank in +! [0, mpp_npes()-1] using floor((m-1) * npes / nmembers). Spreads ensemble +! members across the world communicator so concurrent disk I/O does not +! collide on PE 0; when nmembers > npes the surplus members wrap onto earlier +! PEs in stride. Reused by the writer and reader ensemble orchestrators. +! +! Examples (npes=8): +! nmembers=4 -> roots [0, 2, 4, 6] +! nmembers=8 -> roots [0..7] +! nmembers=20 -> roots [0,0,0,1,1,2,2,2,3,3,4,4,4,5,5,6,6,6,7,7] +!============================================================================== +function soca_io_ensemble_root_pe(m, nmembers) result(pe) + integer, intent(in) :: m + integer, intent(in) :: nmembers + integer :: pe + pe = (m - 1) * mpp_npes() / nmembers +end function soca_io_ensemble_root_pe + + !============================================================================== ! Defensive buffer-shape checks, called from each enqueue. Catches an ! unallocated actual (size 0) and a caller that mis-sized the data-domain diff --git a/src/soca/Increment/Increment.cc b/src/soca/Increment/Increment.cc index c510217e..ec4a059a 100755 --- a/src/soca/Increment/Increment.cc +++ b/src/soca/Increment/Increment.cc @@ -19,6 +19,7 @@ #include "soca/Increment/IncrementFortran.h" #include "soca/State/State.h" +#include "eckit/config/LocalConfiguration.h" #include "eckit/exception/Exceptions.h" #include "oops/base/GeometryData.h" @@ -373,6 +374,55 @@ namespace soca { // ----------------------------------------------------------------------------- + void Increment::writeEnsemble(const std::vector & increments, + const std::vector & configs) { + Log::trace() << "soca::Increment::writeEnsemble starting (" << increments.size() + << " members)" << std::endl; + ASSERT(increments.size() == configs.size()); + if (increments.empty()) return; + + const size_t n = increments.size(); + std::vector keys(n); + std::vector confPtrs(n); + std::vector dtPtrs(n); + for (size_t i = 0; i < n; ++i) { + keys[i] = increments[i]->keyFlds_; + confPtrs[i] = &configs[i]; + dtPtrs[i] = &increments[i]->time_; + } + const int nm = static_cast(n); + soca_increment_write_ensemble_f90(nm, keys.data(), confPtrs.data(), dtPtrs.data()); + Log::trace() << "soca::Increment::writeEnsemble done" << std::endl; + } + + // ----------------------------------------------------------------------------- + + void Increment::readEnsemble(const std::vector & increments, + const std::vector & configs) { + Log::trace() << "soca::Increment::readEnsemble starting (" << increments.size() + << " members)" << std::endl; + ASSERT(increments.size() == configs.size()); + if (increments.empty()) return; + + const size_t n = increments.size(); + std::vector keys(n); + std::vector confPtrs(n); + std::vector dtPtrs(n); + for (size_t i = 0; i < n; ++i) { + keys[i] = increments[i]->keyFlds_; + confPtrs[i] = &configs[i]; + dtPtrs[i] = &increments[i]->time_; + } + const int nm = static_cast(n); + soca_increment_read_ensemble_f90(nm, keys.data(), confPtrs.data(), dtPtrs.data()); + + for (auto* x : increments) x->fieldSet_.set_dirty(); + + Log::trace() << "soca::Increment::readEnsemble done" << std::endl; + } + + // ----------------------------------------------------------------------------- + void Increment::horiz_scales(const eckit::Configuration & config) { soca_increment_horiz_scales_f90(toFortran(), &config); Log::trace() << "Horiz decorrelation length scales computed." << std::endl; diff --git a/src/soca/Increment/Increment.h b/src/soca/Increment/Increment.h index eb4c29d8..a92f0056 100644 --- a/src/soca/Increment/Increment.h +++ b/src/soca/Increment/Increment.h @@ -94,6 +94,18 @@ namespace soca { /// I/O and diagnostics void read(const eckit::Configuration &); void write(const eckit::Configuration &) const; + + /// Bulk parallel write across ensemble members. Mirrors State::writeEnsemble + /// for ensemble-of-increments output (e.g. LETKF posterior ensemble + /// increments). Each member is gathered to a strided writer PE and + /// written concurrently. + static void writeEnsemble(const std::vector &, + const std::vector &); + + /// Bulk read across ensemble members. Mirrors State::readEnsemble. + static void readEnsemble(const std::vector &, + const std::vector &); + void horiz_scales(const eckit::Configuration &); void vert_scales(const double &); std::vector rmsByLevel(const std::string &) const; diff --git a/src/soca/Increment/IncrementFortran.h b/src/soca/Increment/IncrementFortran.h index 7d802c61..f7003701 100644 --- a/src/soca/Increment/IncrementFortran.h +++ b/src/soca/Increment/IncrementFortran.h @@ -46,6 +46,14 @@ namespace soca { void soca_increment_horiz_scales_f90(F90flds &, const eckit::Configuration * const &); void soca_increment_vert_scales_f90(F90flds &, const double &); + void soca_increment_write_ensemble_f90(const int &, + const F90flds * const, + const eckit::Configuration * const *, + const util::DateTime * const *); + void soca_increment_read_ensemble_f90(const int &, + const F90flds * const, + const eckit::Configuration * const *, + util::DateTime * const *); } } // namespace soca #endif // SOCA_INCREMENT_INCREMENTFORTRAN_H_ diff --git a/src/soca/Increment/soca_increment.interface.F90 b/src/soca/Increment/soca_increment.interface.F90 index 5640c413..b794bf70 100644 --- a/src/soca/Increment/soca_increment.interface.F90 +++ b/src/soca/Increment/soca_increment.interface.F90 @@ -14,6 +14,8 @@ module soca_increment_mod_c use oops_variables_mod, only : oops_variables ! soca modules +use soca_fields_mod, only: soca_fields, soca_fields_ptr_t, & + soca_fields_write_ensemble, soca_fields_read_ensemble use soca_geom_mod_c, only: soca_geom_registry use soca_geom_mod, only: soca_geom use soca_increment_mod, only : soca_increment @@ -197,5 +199,64 @@ subroutine soca_increment_vert_scales_c(c_key_self, c_vert) bind(c,name='soca_in end subroutine soca_increment_vert_scales_c + +! ------------------------------------------------------------------------------ +!> C++ interface: bulk write multiple increments in one ensemble pass. Each +!! member is gathered to a rotated writer PE and written concurrently with the +!! others. +subroutine soca_increment_write_ensemble_c(c_n, c_keys, c_confs, c_dts) & + bind(c, name='soca_increment_write_ensemble_f90') + integer(c_int), intent(in) :: c_n + integer(c_int), intent(in) :: c_keys(*) + type(c_ptr), intent(in) :: c_confs(*) + type(c_ptr), intent(in) :: c_dts(*) + + type(soca_fields_ptr_t), allocatable :: fptrs(:) + type(c_ptr), allocatable :: confs_local(:) + type(datetime), allocatable :: fdates(:) + type(soca_increment), pointer :: fld + integer :: i + + if (c_n <= 0) return + allocate(fptrs(c_n), confs_local(c_n), fdates(c_n)) + do i = 1, c_n + call soca_increment_registry%get(c_keys(i), fld) + fptrs(i)%p => fld + confs_local(i) = c_confs(i) + call c_f_datetime(c_dts(i), fdates(i)) + end do + + call soca_fields_write_ensemble(fptrs, confs_local, fdates) +end subroutine soca_increment_write_ensemble_c + + +! ------------------------------------------------------------------------------ +!> C++ interface: bulk read multiple increments in one ensemble pass. +subroutine soca_increment_read_ensemble_c(c_n, c_keys, c_confs, c_dts) & + bind(c, name='soca_increment_read_ensemble_f90') + integer(c_int), intent(in) :: c_n + integer(c_int), intent(in) :: c_keys(*) + type(c_ptr), intent(in) :: c_confs(*) + type(c_ptr), intent(inout) :: c_dts(*) + + type(soca_fields_ptr_t), allocatable :: fptrs(:) + type(c_ptr), allocatable :: confs_local(:) + type(datetime), allocatable :: fdates(:) + type(soca_increment), pointer :: fld + integer :: i + + if (c_n <= 0) return + allocate(fptrs(c_n), confs_local(c_n), fdates(c_n)) + do i = 1, c_n + call soca_increment_registry%get(c_keys(i), fld) + fptrs(i)%p => fld + confs_local(i) = c_confs(i) + call c_f_datetime(c_dts(i), fdates(i)) + end do + + call soca_fields_read_ensemble(fptrs, confs_local, fdates) +end subroutine soca_increment_read_ensemble_c + + ! ------------------------------------------------------------------------------ end module diff --git a/src/soca/State/State.cc b/src/soca/State/State.cc index 4e19c50f..08f197f1 100644 --- a/src/soca/State/State.cc +++ b/src/soca/State/State.cc @@ -289,6 +289,59 @@ namespace soca { // ----------------------------------------------------------------------------- + void State::writeEnsemble(const std::vector & states, + const std::vector & configs) { + Log::trace() << "soca::State::writeEnsemble starting (" << states.size() + << " members)" << std::endl; + ASSERT(states.size() == configs.size()); + if (states.empty()) return; + + const size_t n = states.size(); + std::vector keys(n); + std::vector confPtrs(n); + std::vector dtPtrs(n); + for (size_t i = 0; i < n; ++i) { + keys[i] = states[i]->keyFlds_; + confPtrs[i] = &configs[i]; + dtPtrs[i] = &states[i]->time_; + } + const int nm = static_cast(n); + soca_state_write_ensemble_f90(nm, keys.data(), confPtrs.data(), dtPtrs.data()); + Log::trace() << "soca::State::writeEnsemble done" << std::endl; + } + + // ----------------------------------------------------------------------------- + + void State::readEnsemble(const std::vector & states, + const std::vector & configs) { + Log::trace() << "soca::State::readEnsemble starting (" << states.size() + << " members)" << std::endl; + ASSERT(states.size() == configs.size()); + if (states.empty()) return; + + const size_t n = states.size(); + std::vector keys(n); + std::vector confPtrs(n); + std::vector dtPtrs(n); + for (size_t i = 0; i < n; ++i) { + keys[i] = states[i]->keyFlds_; + confPtrs[i] = &configs[i]; + dtPtrs[i] = &states[i]->time_; + } + const int nm = static_cast(n); + soca_state_read_ensemble_f90(nm, keys.data(), confPtrs.data(), dtPtrs.data()); + + // Fortran-side soca_fields_read_finalize marks each atlas::Field dirty, + // but the C++ atlas::FieldSet wrapper keeps its own dirty-bit cache that + // doesn't auto-update from the underlying fields. Mirror the writeEnsemble + // pattern and re-mark here so downstream FieldSet consumers refresh. + for (auto* x : states) x->fieldSet_.set_dirty(); + + Log::trace() << "soca::State::readEnsemble done" << std::endl; + } + + // ----------------------------------------------------------------------------- + void State::updateFields(const oops::Variables & vars) { // remove fields from the fieldset that are no longer in vars atlas::FieldSet orig = util::shareFields(fieldSet_); diff --git a/src/soca/State/State.h b/src/soca/State/State.h index 57152ef7..a123946f 100644 --- a/src/soca/State/State.h +++ b/src/soca/State/State.h @@ -83,6 +83,21 @@ namespace soca { void read(const eckit::Configuration &); void write(const eckit::Configuration &) const; + /// Bulk parallel write across ensemble members. Each member is gathered + /// to a strided writer PE (assigned by soca_io_ensemble_root_pe in + /// soca_io_mod) and written in phase 2 with no MPI, so per-member + /// netCDF writes happen concurrently across writer PEs. See + /// soca_io_writers_commit_ensemble for the staging. + static void writeEnsemble(const std::vector &, + const std::vector &); + + /// Bulk read across ensemble members via soca_io_readers_commit_ensemble. + /// Members are read in parallel (rotated reader PEs), honoring the + /// geometry.io.'ensemble read' (scatter vs strided) and 'async mpi' + /// toggles. See soca_io_readers_commit_ensemble for the staging. + static void readEnsemble(const std::vector &, + const std::vector &); + int & toFortran() {return keyFlds_;} const int & toFortran() const {return keyFlds_;} diff --git a/src/soca/State/StateFortran.h b/src/soca/State/StateFortran.h index f1b83700..888bad2b 100644 --- a/src/soca/State/StateFortran.h +++ b/src/soca/State/StateFortran.h @@ -39,6 +39,14 @@ namespace soca { const eckit::Configuration * const &, util::DateTime * const *); void soca_state_update_fields_f90(const F90flds &, const oops::Variables &); + void soca_state_write_ensemble_f90(const int &, + const F90flds * const, + const eckit::Configuration * const *, + const util::DateTime * const *); + void soca_state_read_ensemble_f90(const int &, + const F90flds * const, + const eckit::Configuration * const *, + util::DateTime * const *); } } // namespace soca #endif // SOCA_STATE_STATEFORTRAN_H_ diff --git a/src/soca/State/soca_state.interface.F90 b/src/soca/State/soca_state.interface.F90 index a9d2c5e3..98761848 100644 --- a/src/soca/State/soca_state.interface.F90 +++ b/src/soca/State/soca_state.interface.F90 @@ -14,7 +14,8 @@ module soca_state_mod_c use oops_variables_mod, only: oops_variables ! soca modules -use soca_fields_mod, only: soca_field +use soca_fields_mod, only: soca_field, soca_fields_ptr_t, & + soca_fields_write_ensemble, soca_fields_read_ensemble use soca_geom_mod_c, only: soca_geom_registry use soca_geom_mod, only: soca_geom use soca_increment_mod, only: soca_increment @@ -167,4 +168,62 @@ subroutine soca_state_update_fields_c(c_key_self, c_vars) & end subroutine soca_state_update_fields_c + +! ------------------------------------------------------------------------------ +!> C++ interface: bulk write multiple states in one ensemble pass. Each member +!! is gathered to a rotated writer PE and written concurrently with the others. +subroutine soca_state_write_ensemble_c(c_n, c_keys, c_confs, c_dts) & + bind(c, name='soca_state_write_ensemble_f90') + integer(c_int), intent(in) :: c_n + integer(c_int), intent(in) :: c_keys(*) + type(c_ptr), intent(in) :: c_confs(*) + type(c_ptr), intent(in) :: c_dts(*) + + type(soca_fields_ptr_t), allocatable :: fptrs(:) + type(c_ptr), allocatable :: confs_local(:) + type(datetime), allocatable :: fdates(:) + type(soca_state), pointer :: fld + integer :: i + + if (c_n <= 0) return + allocate(fptrs(c_n), confs_local(c_n), fdates(c_n)) + do i = 1, c_n + call soca_state_registry%get(c_keys(i), fld) + fptrs(i)%p => fld + confs_local(i) = c_confs(i) + call c_f_datetime(c_dts(i), fdates(i)) + end do + + call soca_fields_write_ensemble(fptrs, confs_local, fdates) +end subroutine soca_state_write_ensemble_c + + +! ------------------------------------------------------------------------------ +!> C++ interface: bulk read multiple states in one ensemble pass. +subroutine soca_state_read_ensemble_c(c_n, c_keys, c_confs, c_dts) & + bind(c, name='soca_state_read_ensemble_f90') + integer(c_int), intent(in) :: c_n + integer(c_int), intent(in) :: c_keys(*) + type(c_ptr), intent(in) :: c_confs(*) + type(c_ptr), intent(inout) :: c_dts(*) + + type(soca_fields_ptr_t), allocatable :: fptrs(:) + type(c_ptr), allocatable :: confs_local(:) + type(datetime), allocatable :: fdates(:) + type(soca_state), pointer :: fld + integer :: i + + if (c_n <= 0) return + allocate(fptrs(c_n), confs_local(c_n), fdates(c_n)) + do i = 1, c_n + call soca_state_registry%get(c_keys(i), fld) + fptrs(i)%p => fld + confs_local(i) = c_confs(i) + call c_f_datetime(c_dts(i), fdates(i)) + end do + + call soca_fields_read_ensemble(fptrs, confs_local, fdates) +end subroutine soca_state_read_ensemble_c + + end module soca_state_mod_c