module m_omp_backend use m_mpi, only: MPI_COMM_WORLD, MPI_IN_PLACE, MPI_MAX, MPI_SUM, & MPI_Allreduce use m_allocator, only: allocator_t use m_base_backend, only: base_backend_t, source_index_boxes, & check_point_fields, check_labelled_fields, & abl_wall_upper_plane use m_common, only: dp, MPI_X3D2_DP, get_dirs_from_rdr, move_data_loc, & DIR_X, DIR_Y, DIR_Z, DIR_C, NULL_LOC, & X_FACE, Y_FACE, Z_FACE, VERT, BC_DIRICHLET use m_field, only: field_t use m_mesh, only: mesh_t use m_ordering, only: get_index_reordering use m_tdsops, only: dirps_t, tdsops_t use m_omp_common, only: SZ use m_omp_exec_dist, only: exec_dist_tds_compact, exec_dist_transeq_compact use m_exec_thom, only: exec_thom_tds_compact use m_omp_sendrecv, only: sendrecv_fields use m_logger, only: abort_run implicit none private :: transeq_halo_exchange, transeq_dist_component type, extends(base_backend_t) :: omp_backend_t !character(len=*), parameter :: name = 'omp' real(dp), allocatable, dimension(:, :, :) :: & u_recv_s, u_recv_e, u_send_s, u_send_e, & v_recv_s, v_recv_e, v_send_s, v_send_e, & w_recv_s, w_recv_e, w_send_s, w_send_e, & du_send_s, du_send_e, du_recv_s, du_recv_e, & dud_send_s, dud_send_e, dud_recv_s, dud_recv_e, & d2u_send_s, d2u_send_e, d2u_recv_s, d2u_recv_e contains procedure :: alloc_tdsops => alloc_omp_tdsops procedure :: transeq_x => transeq_x_omp procedure :: transeq_y => transeq_y_omp procedure :: transeq_z => transeq_z_omp procedure :: transeq_species => transeq_species_omp procedure :: tds_solve => tds_solve_omp procedure :: thom_solve => thom_solve_omp procedure :: reorder => reorder_omp procedure :: sum_yintox => sum_yintox_omp procedure :: sum_zintox => sum_zintox_omp procedure :: veccopy => veccopy_omp procedure :: vecadd => vecadd_omp procedure :: vecmult => vecmult_omp procedure :: scalar_product => scalar_product_omp procedure :: vector_norm_squared => vector_norm_squared_omp procedure :: field_labelled_sums => field_labelled_sums_omp procedure :: field_labelled_axpy => field_labelled_axpy_omp procedure :: field_sample_points => field_sample_points_omp procedure :: field_add_gaussian_sources => field_add_gaussian_sources_omp procedure :: field_max_mean => field_max_mean_omp procedure :: slice_max_sum => slice_max_sum_omp procedure :: x_slice_weighted_sum => x_slice_weighted_sum_omp procedure :: field_scale => field_scale_omp procedure :: field_shift => field_shift_omp procedure :: field_set_face => field_set_face_omp procedure :: field_set_y_plane => field_set_y_plane_omp procedure :: field_plane_sums => field_plane_sums_omp procedure :: field_set_abl_wall_stress => field_set_abl_wall_stress_omp procedure :: field_set_face_from_field => field_set_face_from_field_omp procedure :: set_field_face_data => set_field_face_data_omp procedure :: field_add_const_x_face => field_add_const_x_face_omp procedure :: field_add_face_from_field => field_add_face_from_field_omp procedure :: compute_vorticity => compute_vorticity_omp procedure :: compute_qcriterion => compute_qcriterion_omp procedure :: compute_smagorinsky_nut => compute_smagorinsky_nut_omp procedure :: compute_sgs_stress => compute_sgs_stress_omp procedure :: field_volume_integral => field_volume_integral_omp procedure :: copy_data_to_f => copy_data_to_f_omp procedure :: copy_f_to_data => copy_f_to_data_omp procedure :: init_poisson_fft => init_omp_poisson_fft procedure :: sync => sync_omp procedure :: get_device_bw_info => get_device_bw_info_omp procedure :: transeq_omp_dist end type omp_backend_t interface omp_backend_t module procedure init end interface omp_backend_t contains function init(mesh, allocator) result(backend) implicit none type(mesh_t), target, intent(inout) :: mesh class(allocator_t), target, intent(inout) :: allocator type(omp_backend_t) :: backend integer :: n_groups call backend%base_init() select type (allocator) class is (allocator_t) ! class level access to the allocator backend%allocator => allocator end select backend%mesh => mesh n_groups = maxval([backend%allocator%get_n_groups(DIR_X), & backend%allocator%get_n_groups(DIR_Y), & backend%allocator%get_n_groups(DIR_Z)]) allocate (backend%u_send_s(SZ, backend%n_halo, n_groups)) allocate (backend%u_send_e(SZ, backend%n_halo, n_groups)) allocate (backend%u_recv_s(SZ, backend%n_halo, n_groups)) allocate (backend%u_recv_e(SZ, backend%n_halo, n_groups)) allocate (backend%v_send_s(SZ, backend%n_halo, n_groups)) allocate (backend%v_send_e(SZ, backend%n_halo, n_groups)) allocate (backend%v_recv_s(SZ, backend%n_halo, n_groups)) allocate (backend%v_recv_e(SZ, backend%n_halo, n_groups)) allocate (backend%w_send_s(SZ, backend%n_halo, n_groups)) allocate (backend%w_send_e(SZ, backend%n_halo, n_groups)) allocate (backend%w_recv_s(SZ, backend%n_halo, n_groups)) allocate (backend%w_recv_e(SZ, backend%n_halo, n_groups)) allocate (backend%du_send_s(SZ, 1, n_groups)) allocate (backend%du_send_e(SZ, 1, n_groups)) allocate (backend%du_recv_s(SZ, 1, n_groups)) allocate (backend%du_recv_e(SZ, 1, n_groups)) allocate (backend%dud_send_s(SZ, 1, n_groups)) allocate (backend%dud_send_e(SZ, 1, n_groups)) allocate (backend%dud_recv_s(SZ, 1, n_groups)) allocate (backend%dud_recv_e(SZ, 1, n_groups)) allocate (backend%d2u_send_s(SZ, 1, n_groups)) allocate (backend%d2u_send_e(SZ, 1, n_groups)) allocate (backend%d2u_recv_s(SZ, 1, n_groups)) allocate (backend%d2u_recv_e(SZ, 1, n_groups)) end function init subroutine alloc_omp_tdsops( & self, tdsops, n_tds, delta, operation, scheme, bc_start, bc_end, & stretch, stretch_correct, n_halo, from_to, sym, c_nu, nu0_nu, & filter_alpha & ) implicit none class(omp_backend_t) :: self class(tdsops_t), allocatable, intent(inout) :: tdsops integer, intent(in) :: n_tds real(dp), intent(in) :: delta character(*), intent(in) :: operation, scheme integer, intent(in) :: bc_start, bc_end real(dp), optional, intent(in) :: stretch(:), stretch_correct(:) integer, optional, intent(in) :: n_halo character(*), optional, intent(in) :: from_to logical, optional, intent(in) :: sym real(dp), optional, intent(in) :: c_nu, nu0_nu real(dp), optional, intent(in) :: filter_alpha allocate (tdsops_t :: tdsops) select type (tdsops) type is (tdsops_t) tdsops = tdsops_t(n_tds, delta, operation, scheme, bc_start, bc_end, & stretch, stretch_correct, n_halo, from_to, sym, & c_nu, nu0_nu, filter_alpha) end select end subroutine alloc_omp_tdsops subroutine sync_omp(self) implicit none class(omp_backend_t) :: self end subroutine sync_omp subroutine get_device_bw_info_omp(self, mem_clock_rt, mem_bus_width, & available) implicit none class(omp_backend_t) :: self integer, intent(out) :: mem_clock_rt integer, intent(out) :: mem_bus_width logical, intent(out) :: available mem_clock_rt = 0 mem_bus_width = 0 available = .false. end subroutine get_device_bw_info_omp subroutine transeq_x_omp(self, du, dv, dw, u, v, w, nu, dirps) implicit none class(omp_backend_t) :: self class(field_t), intent(inout) :: du, dv, dw class(field_t), intent(in) :: u, v, w real(dp), intent(in) :: nu type(dirps_t), intent(in) :: dirps call self%transeq_omp_dist(du, dv, dw, u, v, w, nu, dirps) end subroutine transeq_x_omp subroutine transeq_y_omp(self, du, dv, dw, u, v, w, nu, dirps) implicit none class(omp_backend_t) :: self class(field_t), intent(inout) :: du, dv, dw class(field_t), intent(in) :: u, v, w real(dp), intent(in) :: nu type(dirps_t), intent(in) :: dirps ! u, v, w is reordered so that we pass v, u, w call self%transeq_omp_dist(dv, du, dw, v, u, w, nu, dirps) end subroutine transeq_y_omp subroutine transeq_z_omp(self, du, dv, dw, u, v, w, nu, dirps) implicit none class(omp_backend_t) :: self class(field_t), intent(inout) :: du, dv, dw class(field_t), intent(in) :: u, v, w real(dp), intent(in) :: nu type(dirps_t), intent(in) :: dirps ! u, v, w is reordered so that we pass w, u, v call self%transeq_omp_dist(dw, du, dv, w, u, v, nu, dirps) end subroutine transeq_z_omp subroutine transeq_species_omp(self, dspec, uvw, spec, nu, dirps, sync) !! Compute the convection and diffusion for the given field !! in the given direction. !! Halo exchange for the given field is necessary !! When sync is true, halo exchange of momentum is necessary implicit none class(omp_backend_t) :: self class(field_t), intent(inout) :: dspec class(field_t), intent(in) :: uvw, spec real(dp), intent(in) :: nu type(dirps_t), intent(in) :: dirps logical, intent(in) :: sync integer :: n_groups n_groups = self%allocator%get_n_groups(dirps%dir) ! Halo exchange for momentum if needed if (sync) then call copy_into_buffers(self%u_send_s, self%u_send_e, uvw%data, & dirps%der1st%n_tds, n_groups) call sendrecv_fields(self%u_recv_s, self%u_recv_e, & self%u_send_s, self%u_send_e, & SZ*self%n_halo*n_groups, & self%mesh%par%nproc_dir(dirps%dir), & self%mesh%par%pprev(dirps%dir), & self%mesh%par%pnext(dirps%dir)) end if ! Halo exchange for the given field call copy_into_buffers(self%v_send_s, self%v_send_e, spec%data, & dirps%der1st%n_tds, n_groups) call sendrecv_fields(self%v_recv_s, self%v_recv_e, & self%v_send_s, self%v_send_e, & SZ*self%n_halo*n_groups, & self%mesh%par%nproc_dir(dirps%dir), & self%mesh%par%pprev(dirps%dir), & self%mesh%par%pnext(dirps%dir)) ! combine convection and diffusion call transeq_dist_component(self, dspec, spec, uvw, nu, & self%v_recv_s, self%v_recv_e, & self%u_recv_s, self%u_recv_e, & dirps%der1st, dirps%der1st_sym, & dirps%der2nd, dirps%dir) end subroutine transeq_species_omp subroutine transeq_omp_dist(self, du, dv, dw, u, v, w, nu, dirps) implicit none class(omp_backend_t) :: self class(field_t), intent(inout) :: du, dv, dw class(field_t), intent(in) :: u, v, w real(dp), intent(in) :: nu type(dirps_t), intent(in) :: dirps call transeq_halo_exchange(self, u, v, w, dirps%dir) call transeq_dist_component(self, du, u, u, nu, & self%u_recv_s, self%u_recv_e, & self%u_recv_s, self%u_recv_e, & dirps%der1st, dirps%der1st_sym, & dirps%der2nd, dirps%dir) call transeq_dist_component(self, dv, v, u, nu, & self%v_recv_s, self%v_recv_e, & self%u_recv_s, self%u_recv_e, & dirps%der1st_sym, dirps%der1st, & dirps%der2nd_sym, dirps%dir) call transeq_dist_component(self, dw, w, u, nu, & self%w_recv_s, self%w_recv_e, & self%u_recv_s, self%u_recv_e, & dirps%der1st_sym, dirps%der1st, & dirps%der2nd_sym, dirps%dir) end subroutine transeq_omp_dist subroutine transeq_halo_exchange(self, u, v, w, dir) class(omp_backend_t) :: self class(field_t), intent(in) :: u, v, w integer, intent(in) :: dir integer :: n, nproc_dir, pprev, pnext integer :: n_groups n_groups = self%allocator%get_n_groups(dir) n = self%mesh%get_n(u) nproc_dir = self%mesh%par%nproc_dir(dir) pprev = self%mesh%par%pprev(dir) pnext = self%mesh%par%pnext(dir) call copy_into_buffers(self%u_send_s, self%u_send_e, u%data, & n, n_groups) call copy_into_buffers(self%v_send_s, self%v_send_e, v%data, & n, n_groups) call copy_into_buffers(self%w_send_s, self%w_send_e, w%data, & n, n_groups) call sendrecv_fields(self%u_recv_s, self%u_recv_e, & self%u_send_s, self%u_send_e, & SZ*self%n_halo*n_groups, & nproc_dir, pprev, pnext) call sendrecv_fields(self%v_recv_s, self%v_recv_e, & self%v_send_s, self%v_send_e, & SZ*self%n_halo*n_groups, & nproc_dir, pprev, pnext) call sendrecv_fields(self%w_recv_s, self%w_recv_e, & self%w_send_s, self%w_send_e, & SZ*self%n_halo*n_groups, & nproc_dir, pprev, pnext) end subroutine transeq_halo_exchange subroutine transeq_dist_component(self, rhs_du, u, conv, nu, & u_recv_s, u_recv_e, & conv_recv_s, conv_recv_e, & tdsops_du, tdsops_dud, tdsops_d2u, dir) !! Computes RHS_x^u following: !! !! rhs_x^u = -0.5*(conv*du/dx + d(u*conv)/dx) + nu*d2u/dx2 class(omp_backend_t) :: self !> The result field, it is also used as temporary storage class(field_t), intent(inout) :: rhs_du class(field_t), intent(in) :: u, conv real(dp), intent(in) :: nu real(dp), dimension(:, :, :), intent(in) :: u_recv_s, u_recv_e, & conv_recv_s, conv_recv_e class(tdsops_t), intent(in) :: tdsops_du class(tdsops_t), intent(in) :: tdsops_dud class(tdsops_t), intent(in) :: tdsops_d2u integer, intent(in) :: dir class(field_t), pointer :: d2u, dud dud => self%allocator%get_block(dir) d2u => self%allocator%get_block(dir) call exec_dist_transeq_compact( & rhs_du%data, dud%data, d2u%data, & self%du_send_s, self%du_send_e, self%du_recv_s, self%du_recv_e, & self%dud_send_s, self%dud_send_e, self%dud_recv_s, self%dud_recv_e, & self%d2u_send_s, self%d2u_send_e, self%d2u_recv_s, self%d2u_recv_e, & u%data, u_recv_s, u_recv_e, & conv%data, conv_recv_s, conv_recv_e, & tdsops_du, tdsops_dud, tdsops_d2u, nu, & self%mesh%par%nproc_dir(dir), self%mesh%par%pprev(dir), & self%mesh%par%pnext(dir), self%allocator%get_n_groups(dir)) call rhs_du%set_data_loc(u%data_loc) call self%allocator%release_block(dud) call self%allocator%release_block(d2u) end subroutine transeq_dist_component subroutine tds_solve_omp(self, du, u, tdsops) implicit none class(omp_backend_t) :: self class(field_t), intent(inout) :: du class(field_t), intent(in) :: u class(tdsops_t), intent(in) :: tdsops ! Thomas is exact along an undecomposed line (see tdsops%prefer_thomas) if (tdsops%prefer_thomas) then call self%thom_solve(du, u, tdsops) return end if ! Check if direction matches for both in/out fields if (u%dir /= du%dir) then call abort_run('DIR mismatch between fields in tds_solve.') end if if (u%data_loc /= NULL_LOC) then call du%set_data_loc(move_data_loc(u%data_loc, u%dir, tdsops%move)) end if call tds_solve_dist(self, du, u, tdsops) end subroutine tds_solve_omp subroutine thom_solve_omp(self, du, u, tdsops) implicit none class(omp_backend_t) :: self class(field_t), intent(inout) :: du class(field_t), intent(in) :: u class(tdsops_t), intent(in) :: tdsops if (u%dir /= du%dir) then call abort_run('DIR mismatch between fields in thom_solve.') end if if (u%data_loc /= NULL_LOC) then call du%set_data_loc(move_data_loc(u%data_loc, u%dir, tdsops%move)) end if call exec_thom_tds_compact(du%data, u%data, tdsops, & self%allocator%get_n_groups(u%dir)) end subroutine thom_solve_omp subroutine tds_solve_dist(self, du, u, tdsops) implicit none class(omp_backend_t) :: self class(field_t), intent(inout) :: du class(field_t), intent(in) :: u class(tdsops_t), intent(in) :: tdsops integer :: n_groups, dir dir = u%dir n_groups = self%allocator%get_n_groups(dir) call copy_into_buffers(self%u_send_s, self%u_send_e, u%data, & tdsops%n_tds, n_groups) ! halo exchange call sendrecv_fields(self%u_recv_s, self%u_recv_e, & self%u_send_s, self%u_send_e, & SZ*self%n_halo*n_groups, & self%mesh%par%nproc_dir(dir), & self%mesh%par%pprev(dir), & self%mesh%par%pnext(dir)) call exec_dist_tds_compact( & du%data, u%data, self%u_recv_s, self%u_recv_e, & self%du_send_s, self%du_send_e, self%du_recv_s, self%du_recv_e, & tdsops, self%mesh%par%nproc_dir(dir), & self%mesh%par%pprev(dir), self%mesh%par%pnext(dir), & n_groups) end subroutine tds_solve_dist subroutine reorder_omp(self, u_, u, direction) implicit none class(omp_backend_t) :: self class(field_t), intent(inout) :: u_ class(field_t), intent(in) :: u integer, intent(in) :: direction integer, dimension(3) :: dims, cart_padded integer :: i, j, k integer :: out_i, out_j, out_k integer :: base_i, base_j, base_k, str_i, str_j, str_k integer :: dir_from, dir_to dims = self%allocator%get_padded_dims(u%dir) cart_padded = self%allocator%get_padded_dims(DIR_C) call get_dirs_from_rdr(dir_from, dir_to, direction) ! The reordering map is affine in the leading index only when that leading ! dimension is SZ (a directional field): then the index stays within one ! SZ block and never crosses a mod/division boundary, so evaluate the map ! at i=1 and i=2 once per (j, k) to recover the base and stride, then run a ! plain inner loop with no per element index call. For a DIR_C input the ! leading dimension is nx_padded, the single block assumption fails (e.g. ! RDR_C2Z), so fall back to the exact per element map there. if (dims(1) == SZ) then !$omp parallel do collapse(2) private(i, out_i, out_j, out_k, base_i, base_j, base_k, str_i, str_j, str_k) do k = 1, dims(3) do j = 1, dims(2) call get_index_reordering(base_i, base_j, base_k, 1, j, k, & dir_from, dir_to, SZ, cart_padded) call get_index_reordering(str_i, str_j, str_k, 2, j, k, & dir_from, dir_to, SZ, cart_padded) str_i = str_i - base_i; str_j = str_j - base_j; str_k = str_k - base_k do i = 1, dims(1) out_i = base_i + (i - 1)*str_i out_j = base_j + (i - 1)*str_j out_k = base_k + (i - 1)*str_k u_%data(out_i, out_j, out_k) = u%data(i, j, k) end do end do end do !$omp end parallel do else !$omp parallel do private(out_i, out_j, out_k) collapse(2) do k = 1, dims(3) do j = 1, dims(2) do i = 1, dims(1) call get_index_reordering(out_i, out_j, out_k, i, j, k, & dir_from, dir_to, SZ, cart_padded) u_%data(out_i, out_j, out_k) = u%data(i, j, k) end do end do end do !$omp end parallel do end if ! reorder keeps the data_loc the same call u_%set_data_loc(u%data_loc) end subroutine reorder_omp subroutine sum_yintox_omp(self, u, u_) implicit none class(omp_backend_t) :: self class(field_t), intent(inout) :: u class(field_t), intent(in) :: u_ call sum_intox_omp(self, u, u_, DIR_Y) end subroutine sum_yintox_omp subroutine sum_zintox_omp(self, u, u_) implicit none class(omp_backend_t) :: self class(field_t), intent(inout) :: u class(field_t), intent(in) :: u_ call sum_intox_omp(self, u, u_, DIR_Z) end subroutine sum_zintox_omp subroutine sum_intox_omp(self, u, u_, dir_to) class(omp_backend_t) :: self class(field_t), intent(inout) :: u class(field_t), intent(in) :: u_ integer, intent(in) :: dir_to integer :: dir_from integer, dimension(3) :: dims, cart_padded integer :: i, j, k ! Working indices integer :: ii, jj, kk ! Transpose indices integer :: base_i, base_j, base_k, str_i, str_j, str_k dir_from = DIR_X dims = self%allocator%get_padded_dims(u%dir) cart_padded = self%allocator%get_padded_dims(DIR_C) ! Same trick as reorder_omp, with the same SZ guard. if (dims(1) == SZ) then !$omp parallel do collapse(2) private(i, ii, jj, kk, base_i, base_j, base_k, str_i, str_j, str_k) do k = 1, dims(3) do j = 1, dims(2) call get_index_reordering(base_i, base_j, base_k, 1, j, k, & dir_from, dir_to, SZ, cart_padded) call get_index_reordering(str_i, str_j, str_k, 2, j, k, & dir_from, dir_to, SZ, cart_padded) str_i = str_i - base_i; str_j = str_j - base_j; str_k = str_k - base_k do i = 1, dims(1) ii = base_i + (i - 1)*str_i jj = base_j + (i - 1)*str_j kk = base_k + (i - 1)*str_k u%data(i, j, k) = u%data(i, j, k) + u_%data(ii, jj, kk) end do end do end do !$omp end parallel do else !$omp parallel do private(i, ii, jj, kk) collapse(2) do k = 1, dims(3) do j = 1, dims(2) do i = 1, dims(1) call get_index_reordering(ii, jj, kk, i, j, k, & dir_from, dir_to, SZ, cart_padded) u%data(i, j, k) = u%data(i, j, k) + u_%data(ii, jj, kk) end do end do end do !$omp end parallel do end if end subroutine sum_intox_omp subroutine veccopy_omp(self, dst, src) implicit none class(omp_backend_t) :: self class(field_t), intent(inout) :: dst class(field_t), intent(in) :: src integer :: i, j, k if (src%dir /= dst%dir) then call abort_run("Called vector copy with incompatible fields") end if if (dst%dir == DIR_C) & call abort_run('veccopy does not support DIR_C fields') !$omp parallel do do k = 1, size(dst%data, 3) do j = 1, size(dst%data, 2) !$omp simd do i = 1, SZ dst%data(i, j, k) = src%data(i, j, k) end do !$omp end simd end do end do !$omp end parallel do end subroutine veccopy_omp subroutine vecadd_omp(self, a, x, b, y) implicit none class(omp_backend_t) :: self real(dp), intent(in) :: a class(field_t), intent(in) :: x real(dp), intent(in) :: b class(field_t), intent(inout) :: y integer :: i, j, k if (x%dir /= y%dir) then call abort_run("Called vector add with incompatible fields") end if if (y%dir == DIR_C) call abort_run('vecadd does not support DIR_C fields') !$omp parallel do do k = 1, size(y%data, 3) do j = 1, size(y%data, 2) !$omp simd do i = 1, SZ y%data(i, j, k) = a*x%data(i, j, k) + b*y%data(i, j, k) end do !$omp end simd end do end do !$omp end parallel do end subroutine vecadd_omp subroutine vecmult_omp(self, y, x) !! [[m_base_backend(module):vecmult_base(subroutine)]] implicit none class(omp_backend_t) :: self class(field_t), intent(inout) :: y class(field_t), intent(in) :: x integer :: i, j, k if (x%dir /= y%dir) then call abort_run("Called vector multiply with incompatible fields") end if if (y%dir == DIR_C) call abort_run('vecmult does not support DIR_C fields') !$omp parallel do do k = 1, size(y%data, 3) do j = 1, size(y%data, 2) !$omp simd do i = 1, SZ y%data(i, j, k) = y%data(i, j, k)*x%data(i, j, k) end do !$omp end simd end do end do !$omp end parallel do end subroutine vecmult_omp subroutine compute_vorticity_omp( & self, field_out, dudx, dudy, dudz, dvdx, dvdy, dvdz, dwdx, dwdy, dwdz) implicit none class(omp_backend_t) :: self class(field_t), intent(inout) :: field_out class(field_t), intent(in) :: dudx, dudy, dudz class(field_t), intent(in) :: dvdx, dvdy, dvdz class(field_t), intent(in) :: dwdx, dwdy, dwdz field_out%data = sqrt((dwdy%data - dvdz%data)*(dwdy%data - dvdz%data) + & (dudz%data - dwdx%data)*(dudz%data - dwdx%data) + & (dvdx%data - dudy%data)*(dvdx%data - dudy%data)) end subroutine compute_vorticity_omp subroutine compute_qcriterion_omp( & self, field_out, dudx, dudy, dudz, dvdx, dvdy, dvdz, dwdx, dwdy, dwdz) implicit none class(omp_backend_t) :: self class(field_t), intent(inout) :: field_out class(field_t), intent(in) :: dudx, dudy, dudz class(field_t), intent(in) :: dvdx, dvdy, dvdz class(field_t), intent(in) :: dwdx, dwdy, dwdz field_out%data = 0.5_dp*(dudx%data + dvdy%data + dwdz%data) & *(dudx%data + dvdy%data + dwdz%data) - & 0.5_dp*(dudx%data*dudx%data + & dvdy%data*dvdy%data + & dwdz%data*dwdz%data) - & dudy%data*dvdx%data - & dudz%data*dwdx%data - & dvdz%data*dwdy%data end subroutine compute_qcriterion_omp subroutine compute_smagorinsky_nut_omp( & self, nut, mixing_length_sq, dudx, dudy, dudz, dvdx, dvdy, dvdz, & dwdx, dwdy, dwdz) implicit none class(omp_backend_t) :: self class(field_t), intent(inout) :: nut class(field_t), intent(in) :: mixing_length_sq class(field_t), intent(in) :: dudx, dudy, dudz class(field_t), intent(in) :: dvdx, dvdy, dvdz class(field_t), intent(in) :: dwdx, dwdy, dwdz real(dp) :: sij_sq integer :: i, j, k !$omp parallel do private(sij_sq) collapse(2) do k = 1, size(nut%data, 3) do j = 1, size(nut%data, 2) !$omp simd private(sij_sq) do i = 1, size(nut%data, 1) if (mixing_length_sq%data(i, j, k) > 0._dp) then sij_sq = dudx%data(i, j, k)**2 + & dvdy%data(i, j, k)**2 + & dwdz%data(i, j, k)**2 + & 0.5_dp*(dudy%data(i, j, k) + & dvdx%data(i, j, k))**2 + & 0.5_dp*(dudz%data(i, j, k) + & dwdx%data(i, j, k))**2 + & 0.5_dp*(dvdz%data(i, j, k) + & dwdy%data(i, j, k))**2 nut%data(i, j, k) = mixing_length_sq%data(i, j, k)* & sqrt(2._dp*sij_sq) else nut%data(i, j, k) = 0._dp end if end do !$omp end simd end do end do !$omp end parallel do end subroutine compute_smagorinsky_nut_omp subroutine compute_sgs_stress_omp( & self, stress, nut, gradient_a, gradient_b, scale_a, scale_b) implicit none class(omp_backend_t) :: self class(field_t), intent(inout) :: stress class(field_t), intent(in) :: nut, gradient_a, gradient_b real(dp), intent(in) :: scale_a, scale_b integer :: i, j, k !$omp parallel do collapse(2) do k = 1, size(stress%data, 3) do j = 1, size(stress%data, 2) !$omp simd do i = 1, size(stress%data, 1) if (nut%data(i, j, k) > 0._dp) then stress%data(i, j, k) = nut%data(i, j, k)*( & scale_a*gradient_a%data(i, j, k) + & scale_b*gradient_b%data(i, j, k)) else stress%data(i, j, k) = 0._dp end if end do !$omp end simd end do end do !$omp end parallel do end subroutine compute_sgs_stress_omp real(dp) function scalar_product_omp(self, x, y) result(s) !! [[m_base_backend(module):scalar_product_base(function)]] implicit none class(omp_backend_t) :: self class(field_t), intent(in) :: x, y class(field_t), pointer :: x_, y_ integer, dimension(3) :: dims integer :: i, j, k, ii integer :: nvec, remstart integer :: ierr if ((x%data_loc == NULL_LOC) .or. (y%data_loc == NULL_LOC)) then call abort_run("You must set the data_loc before calling scalar product") end if if (x%data_loc /= y%data_loc) then call abort_run("Called scalar product with incompatible fields") end if ! Reorient data into temporary DIR_C storage x_ => self%allocator%get_block(DIR_C, x%data_loc) call self%get_field_data(x_%data, x) y_ => self%allocator%get_block(DIR_C, y%data_loc) call self%get_field_data(y_%data, y) dims = self%mesh%get_dims(x_%data_loc) nvec = dims(1)/SZ remstart = nvec*SZ + 1 s = 0.0_dp !$omp parallel do reduction(+:s) private(i, ii) collapse(2) do k = 1, dims(3) do j = 1, dims(2) ! Execute inner vectorised loops do ii = 1, nvec !$omp simd reduction(+:s) do i = 1, SZ s = s + x_%data(i + (ii - 1)*SZ, j, k)* & y_%data(i + (ii - 1)*SZ, j, k) end do !$omp end simd end do ! Remainder loop do i = remstart, dims(1) s = s + x_%data(i, j, k)*y_%data(i, j, k) end do end do end do !$omp end parallel do ! Release temporary storage call self%allocator%release_block(x_) call self%allocator%release_block(y_) ! Reduce the result call MPI_Allreduce(MPI_IN_PLACE, s, 1, MPI_X3D2_DP, & MPI_SUM, MPI_COMM_WORLD, & ierr) end function scalar_product_omp real(dp) function vector_norm_squared_omp(self, a, b, c) & result(norm_squared) !! Global sum of a**2 + b**2 + c**2, with one MPI reduction. implicit none class(omp_backend_t) :: self class(field_t), intent(in) :: a, b, c real(dp) :: local_sum, pencil_sum integer :: dims(3), stacked integer :: i, j, k, k_i, k_j, ierr if (a%data_loc == NULL_LOC .or. b%data_loc == NULL_LOC .or. & c%data_loc == NULL_LOC) then call abort_run('You must set data_loc before computing a vector norm.') end if if (a%data_loc /= b%data_loc .or. a%data_loc /= c%data_loc) then call abort_run('Vector-norm fields must use the same data location.') end if if (a%dir /= DIR_X .or. b%dir /= DIR_X .or. c%dir /= DIR_X) then call abort_run('Vector-norm fields must use DIR_X layout.') end if dims = self%mesh%get_dims(a%data_loc) stacked = (dims(2) - 1)/SZ + 1 local_sum = 0._dp !$omp parallel do collapse(2) reduction(+:local_sum) private(k, pencil_sum) do k_j = 1, stacked do k_i = 1, dims(3) k = k_j + (k_i - 1)*stacked pencil_sum = 0._dp do j = 1, dims(1) !$omp simd reduction(+:pencil_sum) do i = 1, min(SZ, dims(2) - (k_j - 1)*SZ) pencil_sum = pencil_sum + a%data(i, j, k)**2 + & b%data(i, j, k)**2 + c%data(i, j, k)**2 end do !$omp end simd end do local_sum = local_sum + pencil_sum end do end do !$omp end parallel do call MPI_Allreduce(local_sum, norm_squared, 1, MPI_X3D2_DP, MPI_SUM, & MPI_COMM_WORLD, ierr) end function vector_norm_squared_omp subroutine field_labelled_sums_omp(self, sums, g, label, a, b, c) !! [[m_base_backend(module):field_labelled_sums_base(subroutine)]] implicit none class(omp_backend_t) :: self real(dp), intent(out) :: sums(:, :) class(field_t), intent(in) :: g, label, a, b, c real(dp) :: local_sums(3, size(sums, 2)), thread_sums(3, size(sums, 2)) integer :: dims(3), n_y_blocks, y_stride, n_label integer :: i, j, k, k_i, k_j, l call check_labelled_fields(g, label, a, b, c) dims = self%mesh%get_dims(g%data_loc) n_y_blocks = (dims(2) - 1)/SZ + 1 y_stride = y_group_stride(self) n_label = size(sums, 2) ! Each thread accumulates into its own copy, merged at the end, instead ! of reduction(+:local_sums): flang's device pass rejects reductions on ! dynamically-shaped arrays even outside target regions. local_sums = 0._dp !$omp parallel private(thread_sums) thread_sums = 0._dp !$omp do collapse(2) private(k, l) do k_i = 1, dims(3) do k_j = 1, n_y_blocks ! y-block is the fast-varying component of the group index k = k_j + (k_i - 1)*y_stride do j = 1, dims(1) do i = 1, min(SZ, dims(2) - (k_j - 1)*SZ) l = nint(label%data(i, j, k)) if (l < 1 .or. l > n_label) cycle thread_sums(1, l) = thread_sums(1, l) & + g%data(i, j, k)*a%data(i, j, k) thread_sums(2, l) = thread_sums(2, l) & + g%data(i, j, k)*b%data(i, j, k) thread_sums(3, l) = thread_sums(3, l) & + g%data(i, j, k)*c%data(i, j, k) end do end do end do end do !$omp end do !$omp critical local_sums = local_sums + thread_sums !$omp end critical !$omp end parallel sums = local_sums end subroutine field_labelled_sums_omp subroutine field_labelled_axpy_omp(self, coeffs, g, label, a, b, c) !! [[m_base_backend(module):field_labelled_axpy_base(subroutine)]] implicit none class(omp_backend_t) :: self real(dp), intent(in) :: coeffs(:, :) class(field_t), intent(in) :: g, label class(field_t), intent(inout) :: a, b, c integer :: dims(3), n_y_blocks, y_stride, n_label integer :: i, j, k, k_i, k_j, l call check_labelled_fields(g, label, a, b, c) dims = self%mesh%get_dims(g%data_loc) n_y_blocks = (dims(2) - 1)/SZ + 1 y_stride = y_group_stride(self) n_label = size(coeffs, 2) !$omp parallel do collapse(2) private(k, l) do k_i = 1, dims(3) do k_j = 1, n_y_blocks ! y-block is the fast-varying component of the group index k = k_j + (k_i - 1)*y_stride do j = 1, dims(1) do i = 1, min(SZ, dims(2) - (k_j - 1)*SZ) l = nint(label%data(i, j, k)) if (l < 1 .or. l > n_label) cycle a%data(i, j, k) = a%data(i, j, k) + coeffs(1, l)*g%data(i, j, k) b%data(i, j, k) = b%data(i, j, k) + coeffs(2, l)*g%data(i, j, k) c%data(i, j, k) = c%data(i, j, k) + coeffs(3, l)*g%data(i, j, k) end do end do end do end do !$omp end parallel do end subroutine field_labelled_axpy_omp subroutine field_sample_points_omp(self, vals, points, a, b, c) !! [[m_base_backend(module):field_sample_points_base(subroutine)]] implicit none class(omp_backend_t) :: self real(dp), intent(out) :: vals(:, :) real(dp), intent(in) :: points(:, :) class(field_t), intent(in) :: a, b, c integer :: dims(3), gdims(3), n_offset(3), y_stride integer :: p, o1, o2, o3, dir, g0(3), loc(3), i, k real(dp) :: frac(3), wgt(0:1, 3), w logical :: periodic(3), owned call check_point_fields(a, b, c) dims = self%mesh%get_dims(VERT) gdims = self%mesh%get_global_dims(VERT) n_offset = self%mesh%par%n_offset periodic = self%mesh%grid%periodic_BC y_stride = y_group_stride(self) vals = 0._dp do p = 1, size(points, 2) ! Zero-based global index of the lower corner and the weights do dir = 1, 3 frac(dir) = points(dir, p)/self%mesh%geo%d(dir) g0(dir) = floor(frac(dir)) frac(dir) = frac(dir) - g0(dir) wgt(0, dir) = 1._dp - frac(dir) wgt(1, dir) = frac(dir) end do do o3 = 0, 1 do o2 = 0, 1 do o1 = 0, 1 w = wgt(o1, 1)*wgt(o2, 2)*wgt(o3, 3) if (w == 0._dp) cycle call local_corner(g0 + [o1, o2, o3], gdims, n_offset, dims, & periodic, loc, owned) if (.not. owned) cycle ! DIR_X storage: x along j, y split into SZ-blocks i = mod(loc(2) - 1, SZ) + 1 k = (loc(2) - 1)/SZ + 1 + (loc(3) - 1)*y_stride vals(1, p) = vals(1, p) + w*a%data(i, loc(1), k) vals(2, p) = vals(2, p) + w*b%data(i, loc(1), k) vals(3, p) = vals(3, p) + w*c%data(i, loc(1), k) end do end do end do end do end subroutine field_sample_points_omp subroutine field_add_gaussian_sources_omp(self, a, b, c, points, coeffs, & eps, r_cut) !! [[m_base_backend(module):field_add_gaussian_sources_base(subroutine)]] implicit none class(omp_backend_t) :: self class(field_t), intent(inout) :: a, b, c real(dp), intent(in) :: points(:, :), coeffs(:, :) real(dp), intent(in) :: eps, r_cut integer :: lo(3, size(points, 2)), hi(3, size(points, 2)) integer :: dims(3), gdims(3), n_offset(3), y_stride integer :: s, gx, gy, gz, x, y, z, i, k real(dp) :: d(3), inv_eps2, r_cut2, dz2, dyz2, dist2, kern logical :: periodic(3) call check_point_fields(a, b, c) d = self%mesh%geo%d dims = self%mesh%get_dims(VERT) gdims = self%mesh%get_global_dims(VERT) n_offset = self%mesh%par%n_offset periodic = self%mesh%grid%periodic_BC y_stride = y_group_stride(self) inv_eps2 = 1._dp/eps**2 r_cut2 = r_cut**2 call source_index_boxes(self%mesh, points, r_cut, lo, hi) ! Sources in order, so every point accumulates them in the same order ! whatever the thread count. Distances use the unwrapped global index g; ! the point it lands on is found by wrapping g. do s = 1, size(points, 2) !$omp parallel do private(dz2, dyz2, dist2, kern, i, k, gx, gy, x, y, z) do gz = lo(3, s), hi(3, s) z = wrapped_local(gz, gdims(3), n_offset(3), dims(3), periodic(3)) if (z == 0) cycle dz2 = (points(3, s) - gz*d(3))**2 if (dz2 > r_cut2) cycle do gy = lo(2, s), hi(2, s) y = wrapped_local(gy, gdims(2), n_offset(2), dims(2), periodic(2)) if (y == 0) cycle dyz2 = dz2 + (points(2, s) - gy*d(2))**2 if (dyz2 > r_cut2) cycle i = mod(y - 1, SZ) + 1 k = (y - 1)/SZ + 1 + (z - 1)*y_stride do gx = lo(1, s), hi(1, s) x = wrapped_local(gx, gdims(1), n_offset(1), dims(1), & periodic(1)) if (x == 0) cycle dist2 = dyz2 + (points(1, s) - gx*d(1))**2 if (dist2 > r_cut2) cycle kern = exp(-dist2*inv_eps2) a%data(i, x, k) = a%data(i, x, k) + coeffs(1, s)*kern b%data(i, x, k) = b%data(i, x, k) + coeffs(2, s)*kern c%data(i, x, k) = c%data(i, x, k) + coeffs(3, s)*kern end do end do end do !$omp end parallel do end do end subroutine field_add_gaussian_sources_omp pure integer function wrapped_local(g, n_glob, offset, n_loc, periodic) !! Local one-based index of the zero-based global VERT index g, wrapped !! when periodic; 0 when it is off this rank or outside the domain. implicit none integer, intent(in) :: g, n_glob, offset, n_loc logical, intent(in) :: periodic integer :: gw gw = g if (periodic) gw = modulo(g, n_glob) wrapped_local = gw - offset + 1 if (gw < 0 .or. gw >= n_glob .or. wrapped_local < 1 & .or. wrapped_local > n_loc) wrapped_local = 0 end function wrapped_local pure subroutine local_corner(g, gdims, n_offset, dims, periodic, loc, & owned) !! Local one-based index of the zero-based global VERT index g, wrapped !! across periodic boundaries; owned is false off this rank or outside !! a non-periodic boundary. implicit none integer, intent(in) :: g(3), gdims(3), n_offset(3), dims(3) logical, intent(in) :: periodic(3) integer, intent(out) :: loc(3) logical, intent(out) :: owned integer :: dir, gw owned = .true. do dir = 1, 3 gw = g(dir) if (periodic(dir)) then gw = modulo(gw, gdims(dir)) else if (gw < 0 .or. gw >= gdims(dir)) then owned = .false. end if loc(dir) = gw - n_offset(dir) + 1 if (loc(dir) < 1 .or. loc(dir) > dims(dir)) owned = .false. end do end subroutine local_corner subroutine copy_into_buffers(u_send_s, u_send_e, u, n, n_groups) implicit none real(dp), dimension(:, :, :), intent(out) :: u_send_s, u_send_e real(dp), dimension(:, :, :), intent(in) :: u integer, intent(in) :: n integer, intent(in) :: n_groups integer :: i, j, k integer :: n_halo = 4 !$omp parallel do do k = 1, n_groups do j = 1, n_halo !$omp simd do i = 1, SZ u_send_s(i, j, k) = u(i, j, k) u_send_e(i, j, k) = u(i, n - n_halo + j, k) end do !$omp end simd end do end do !$omp end parallel do end subroutine copy_into_buffers subroutine field_max_mean_omp(self, max_val, mean_val, f, enforced_data_loc) !! [[m_base_backend(module):field_max_mean_base(subroutine)]] implicit none class(omp_backend_t) :: self real(dp), intent(out) :: max_val, mean_val class(field_t), intent(in) :: f integer, optional, intent(in) :: enforced_data_loc real(dp) :: val, max_p, sum_p, max_pncl, sum_pncl integer :: data_loc, dims(3), dims_padded(3), n, n_i, n_groups_pad, n_j integer :: i, j, k, k_i, k_j, ierr if (f%data_loc == NULL_LOC .and. (.not. present(enforced_data_loc))) then call abort_run('The input field to omp::field_max_mean does not have a & &valid f%data_loc. You may enforce a data_loc of your & &choice as last argument to carry on at your own risk!') end if if (present(enforced_data_loc)) then data_loc = enforced_data_loc else data_loc = f%data_loc end if dims = self%mesh%get_dims(data_loc) dims_padded = self%allocator%get_padded_dims(DIR_C) ! n_groups_pad is the stride between successive outer entries in the ! group index, and it comes from the padded extent of the grouped ! dimension, not the unpadded one (see get_index_dir in m_ordering). ! The two agree only when that dimension needs no padding, so using the ! unpadded count walks into other groups as soon as it does. ! Set for every supported direction; abort_run stops the others. n = 0; n_i = 0; n_j = 0; n_groups_pad = 0 if (f%dir == DIR_X) then n = dims(1); n_j = dims(2); n_i = dims(3) n_groups_pad = dims_padded(2)/SZ else if (f%dir == DIR_Y) then n = dims(2); n_j = dims(1); n_i = dims(3) n_groups_pad = dims_padded(1)/SZ else if (f%dir == DIR_Z) then n = dims(3); n_j = dims(1); n_i = dims(2) n_groups_pad = dims_padded(1)/SZ else call abort_run('field_max_mean does not support DIR_C fields!') end if sum_p = 0._dp max_p = 0._dp !$omp parallel do collapse(2) reduction(+:sum_p) reduction(max:max_p) & !$omp private(k, val, sum_pncl, max_pncl) do k_j = 1, (n_j - 1)/SZ + 1 ! loop over stacked groups do k_i = 1, n_i k = k_j + (k_i - 1)*n_groups_pad sum_pncl = 0._dp max_pncl = 0._dp do j = 1, n ! loop over only non-padded entries in the present group do i = 1, min(SZ, n_j - (k_j - 1)*SZ) val = abs(f%data(i, j, k)) sum_pncl = sum_pncl + val max_pncl = max(max_pncl, val) end do end do sum_p = sum_p + sum_pncl max_p = max(max_p, max_pncl) end do end do !$omp end parallel do ! rank-local values max_val = max_p mean_val = sum_p/product(self%mesh%get_global_dims(data_loc)) ! make sure all ranks have final values call MPI_Allreduce(MPI_IN_PLACE, max_val, 1, MPI_X3D2_DP, & MPI_MAX, MPI_COMM_WORLD, ierr) call MPI_Allreduce(MPI_IN_PLACE, mean_val, 1, MPI_X3D2_DP, & MPI_SUM, MPI_COMM_WORLD, ierr) end subroutine field_max_mean_omp ! ========================================================================= ! Add to the OMP backend module, next to field_max_mean_omp ! ========================================================================= subroutine slice_max_sum_omp(self, max_val, sum_val, f, i_slice, & enforced_data_loc) !! [[m_base_backend(module):slice_max_sum_base(subroutine)]] implicit none class(omp_backend_t) :: self real(dp), intent(out) :: max_val, sum_val class(field_t), intent(in) :: f integer, intent(in) :: i_slice integer, optional, intent(in) :: enforced_data_loc real(dp) :: val, max_p, sum_p integer :: data_loc, dims(3), dims_padded(3), n, n_i, n_groups_pad, n_j integer :: i, j, k, k_i, k_j if (f%data_loc == NULL_LOC .and. (.not. present(enforced_data_loc))) then call abort_run('The input field to omp::slice_max_sum does not have a & &valid f%data_loc. You may enforce a data_loc of your & &choice as last argument to carry on at your own risk!') end if if (present(enforced_data_loc)) then data_loc = enforced_data_loc else data_loc = f%data_loc end if dims = self%mesh%get_dims(data_loc) dims_padded = self%allocator%get_padded_dims(DIR_C) ! See the note on n_groups_pad in field_max_mean_omp above. ! Set for every supported direction; abort_run stops the others. n = 0; n_i = 0; n_j = 0; n_groups_pad = 0 if (f%dir == DIR_X) then n = dims(1); n_j = dims(2); n_i = dims(3) n_groups_pad = dims_padded(2)/SZ else if (f%dir == DIR_Y) then n = dims(2); n_j = dims(1); n_i = dims(3) n_groups_pad = dims_padded(1)/SZ else if (f%dir == DIR_Z) then n = dims(3); n_j = dims(1); n_i = dims(2) n_groups_pad = dims_padded(1)/SZ else call abort_run('slice_max_sum does not support DIR_C fields!') end if if (i_slice < 1 .or. i_slice > n) then call abort_run('slice_max_sum: i_slice out of range') end if j = i_slice sum_p = 0._dp max_p = -huge(1._dp) !$omp parallel do collapse(2) reduction(+:sum_p) reduction(max:max_p) & !$omp private(k, val) do k_j = 1, (n_j - 1)/SZ + 1 do k_i = 1, n_i k = k_j + (k_i - 1)*n_groups_pad do i = 1, min(SZ, n_j - (k_j - 1)*SZ) val = f%data(i, j, k) sum_p = sum_p + val max_p = max(max_p, val) end do end do end do !$omp end parallel do ! Rank-local values; caller is responsible for MPI_Allreduce. max_val = max_p sum_val = sum_p end subroutine slice_max_sum_omp subroutine x_slice_weighted_sum_omp(self, sum_val, f, i_slice, y_weights) !! [[m_base_backend(module):x_slice_weighted_sum_base(subroutine)]] implicit none class(omp_backend_t) :: self real(dp), intent(out) :: sum_val class(field_t), intent(in) :: f integer, intent(in) :: i_slice real(dp), intent(in) :: y_weights(:) real(dp) :: sum_p integer :: dims(3), y_stride, i, k, y_block, z if (f%dir /= DIR_X) & call abort_run('x_slice_weighted_sum: only supported for DIR_X fields.') if (f%data_loc == NULL_LOC) & call abort_run('x_slice_weighted_sum: requires a valid data_loc.') dims = self%mesh%get_dims(f%data_loc) y_stride = y_group_stride(self) if (i_slice < 1 .or. i_slice > dims(1)) & call abort_run('x_slice_weighted_sum: i_slice out of range') if (size(y_weights) /= dims(2)) & call abort_run('x_slice_weighted_sum: one weight per y point needed') sum_p = 0._dp ! y-block is the fast-varying component of the group index !$omp parallel do collapse(2) reduction(+:sum_p) private(k) do z = 1, dims(3) do y_block = 1, (dims(2) - 1)/SZ + 1 k = y_stride*(z - 1) + y_block do i = 1, min(SZ, dims(2) - (y_block - 1)*SZ) sum_p = sum_p + y_weights(i + (y_block - 1)*SZ)*f%data(i, i_slice, k) end do end do end do !$omp end parallel do ! Rank-local value; caller is responsible for MPI_Allreduce. sum_val = sum_p end subroutine x_slice_weighted_sum_omp subroutine field_scale_omp(self, f, a) implicit none class(omp_backend_t) :: self class(field_t), intent(in) :: f real(dp), intent(in) :: a f%data = a*f%data end subroutine field_scale_omp subroutine field_shift_omp(self, f, a) implicit none class(omp_backend_t) :: self class(field_t), intent(in) :: f real(dp), intent(in) :: a f%data = f%data + a end subroutine field_shift_omp subroutine field_set_face_omp(self, f, c_start, c_end, face, & bc_start, bc_end, flow_rate_diff) !! [[m_base_backend(module):field_set_face_base(subroutine)]] implicit none class(omp_backend_t) :: self class(field_t), intent(inout) :: f real(dp), intent(in) :: c_start, c_end integer, intent(in) :: face integer, optional, intent(in) :: bc_start integer, optional, intent(in) :: bc_end real(dp), optional, intent(in) :: flow_rate_diff integer :: dims(3), k, j, i_mod, k_start, k_end, bc_s, bc_e integer :: n_y_blocks, y_stride real(dp) :: fl_corr logical :: set_start, set_end if (f%dir /= DIR_X) then call abort_run('Setting a field face is only supported for DIR_X & &fields.') end if if (f%data_loc == NULL_LOC) then call abort_run('field_set_face require a valid data_loc.') end if fl_corr = 0._dp if (present(flow_rate_diff)) fl_corr = flow_rate_diff ! Only Dirichlet faces carry a prescribed value. Defaulting both to ! Dirichlet keeps the historical behaviour of writing both faces, while ! passing a non-Dirichlet BC leaves that face untouched -- needed where ! the two y-faces differ, as in the ABL (no-slip floor, free-slip lid). bc_s = BC_DIRICHLET bc_e = BC_DIRICHLET if (present(bc_start)) bc_s = bc_start if (present(bc_end)) bc_e = bc_end set_start = (bc_s == BC_DIRICHLET) set_end = (bc_e == BC_DIRICHLET) dims = self%mesh%get_dims(f%data_loc) select case (face) case (X_FACE) call abort_run('Setting X_FACE is not yet supported.') case (Y_FACE) ! DIR_X groups are ordered y-block fastest: b = (z-1)*y_stride + yb. ! The stride is the padded y-block count (see get_index_dir), which ! is larger than n_y_blocks for a cell field on a non-periodic y. i_mod = mod(dims(2) - 1, SZ) + 1 n_y_blocks = (dims(2) - 1)/SZ + 1 y_stride = y_group_stride(self) !$omp parallel do private(k_start, k_end) do k = 1, dims(3) k_start = 1 + (k - 1)*y_stride k_end = (k - 1)*y_stride + n_y_blocks do j = 1, dims(1) if (set_start) f%data(1, j, k_start) = c_start if (set_end) f%data(i_mod, j, k_end) = c_end end do end do !$omp end parallel do case (Z_FACE) call abort_run('Setting Z_FACE is not yet supported.') case default call abort_run('face is undefined.') end select end subroutine field_set_face_omp integer function y_group_stride(self) result(stride) !! Groups between successive z in the OMP DIR_X ordering, !! group = stride*(z - 1) + y_block. It is the padded y-block count !! (see get_index_dir in m_ordering), so it can exceed the number of !! y-blocks a field uses, e.g. a cell field on a non-periodic y. implicit none class(omp_backend_t), intent(in) :: self integer :: dims_padded(3) dims_padded = self%allocator%get_padded_dims(DIR_C) stride = dims_padded(2)/SZ end function y_group_stride subroutine field_plane_sums_omp(self, sums, f) !! [[m_base_backend(module):field_plane_sums(subroutine)]] implicit none class(omp_backend_t) :: self real(dp), intent(out) :: sums(:) class(field_t), intent(in) :: f integer :: dims(3), y_stride, y_block, i, j, k, z real(dp) :: plane_sum if (f%dir /= DIR_X) & call abort_run('field_plane_sums is only supported for DIR_X fields.') if (f%data_loc == NULL_LOC) & call abort_run('field_plane_sums requires a valid data_loc.') dims = self%mesh%get_dims(f%data_loc) if (size(sums) < dims(2)) & call abort_run('field_plane_sums: sums is smaller than the y extent.') y_stride = y_group_stride(self) ! OMP DIR_X ordering: group = y_stride*(z - 1) + y_block. Each plane is ! summed by one thread in a fixed order, so the result is reproducible. !$omp parallel do private(y_block, i, k, plane_sum) do j = 1, dims(2) y_block = (j - 1)/SZ + 1 i = j - (y_block - 1)*SZ plane_sum = 0._dp do z = 1, dims(3) do k = 1, dims(1) plane_sum = plane_sum + f%data(i, k, y_stride*(z - 1) + y_block) end do end do sums(j) = plane_sum end do !$omp end parallel do end subroutine field_plane_sums_omp subroutine field_set_y_plane_omp(self, f, c, plane) !! [[m_base_backend(module):field_set_y_plane(subroutine)]] implicit none class(omp_backend_t) :: self class(field_t), intent(inout) :: f real(dp), intent(in) :: c integer, intent(in) :: plane integer :: dims(3), k, j, y_stride, y_block, i_in_block, group if (f%dir /= DIR_X) & call abort_run('field_set_y_plane is only supported for DIR_X fields.') if (f%data_loc == NULL_LOC) & call abort_run('field_set_y_plane requires a valid data_loc.') dims = self%mesh%get_dims(f%data_loc) if (plane < 1 .or. plane > dims(2)) & call abort_run('field_set_y_plane: plane outside the domain.') ! OMP DIR_X ordering: group = y_stride*(z - 1) + y_block. y_stride = y_group_stride(self) y_block = (plane - 1)/SZ + 1 i_in_block = mod(plane - 1, SZ) + 1 !$omp parallel do private(group) do k = 1, dims(3) group = y_stride*(k - 1) + y_block do j = 1, dims(1) f%data(i_in_block, j, group) = c end do end do !$omp end parallel do end subroutine field_set_y_plane_omp subroutine field_set_abl_wall_stress_omp( & self, stress, u, w, sample_plane, sample_weight, stress_plane, & drag_coeff, component) !! [[m_base_backend(module):field_set_abl_wall_stress(subroutine)]] implicit none class(omp_backend_t) :: self class(field_t), intent(inout) :: stress class(field_t), intent(in) :: u, w integer, intent(in) :: sample_plane, stress_plane, component real(dp), intent(in) :: sample_weight, drag_coeff integer :: dims(3), y_stride, upper_plane integer :: sample_block, upper_block, stress_block integer :: sample_i, upper_i, stress_i, k, i integer :: sample_group, upper_group, stress_group real(dp) :: us, ws, speed if (stress%dir /= DIR_X .or. u%dir /= DIR_X .or. w%dir /= DIR_X) & call abort_run('ABL wall stress requires DIR_X fields.') if (stress%data_loc /= VERT .or. u%data_loc /= VERT .or. & w%data_loc /= VERT) & call abort_run('ABL wall stress requires vertex fields.') dims = self%mesh%get_dims(VERT) if (min(sample_plane, stress_plane) < 1 .or. & max(sample_plane, stress_plane) > dims(2)) & call abort_run('ABL wall stress plane is outside the domain.') if (component /= 1 .and. component /= 3) & call abort_run('Invalid ABL wall stress component.') upper_plane = abl_wall_upper_plane(sample_plane, sample_weight, dims(2)) y_stride = y_group_stride(self) sample_block = (sample_plane - 1)/SZ + 1 upper_block = (upper_plane - 1)/SZ + 1 stress_block = (stress_plane - 1)/SZ + 1 sample_i = mod(sample_plane - 1, SZ) + 1 upper_i = mod(upper_plane - 1, SZ) + 1 stress_i = mod(stress_plane - 1, SZ) + 1 !$omp parallel do private(sample_group, upper_group, stress_group, i, & !$omp us, ws, speed) do k = 1, dims(3) sample_group = y_stride*(k - 1) + sample_block upper_group = y_stride*(k - 1) + upper_block stress_group = y_stride*(k - 1) + stress_block do i = 1, dims(1) us = (1._dp - sample_weight)*u%data(sample_i, i, sample_group) & + sample_weight*u%data(upper_i, i, upper_group) ws = (1._dp - sample_weight)*w%data(sample_i, i, sample_group) & + sample_weight*w%data(upper_i, i, upper_group) speed = sqrt(us**2 + ws**2) if (component == 1) then stress%data(stress_i, i, stress_group) = drag_coeff*us*speed else stress%data(stress_i, i, stress_group) = drag_coeff*ws*speed end if end do end do !$omp end parallel do end subroutine field_set_abl_wall_stress_omp subroutine field_set_face_from_field_omp(self, f, f_start, c_end, face, & bc_start, bc_end, flow_rate_diff) implicit none class(omp_backend_t) :: self class(field_t), intent(inout) :: f class(field_t), intent(in) :: f_start real(dp), intent(in) :: c_end integer, intent(in) :: face integer, optional, intent(in) :: bc_start, bc_end real(dp), optional, intent(in) :: flow_rate_diff integer :: dims(3), k, i, j, z, i_max, n_mod, n_y_blocks, y_block, & y_stride, k_start, k_end real(dp) :: flow_rate_diff_val if (f%dir /= DIR_X) & call abort_run('field_set_face_from_field: only supported for DIR_X & &fields.') if (f%data_loc == NULL_LOC) & call abort_run('field_set_face_from_field: requires a valid data_loc.') flow_rate_diff_val = 0._dp if (present(flow_rate_diff)) flow_rate_diff_val = flow_rate_diff dims = self%mesh%get_dims(f%data_loc) n_mod = mod(dims(2) - 1, SZ) + 1 n_y_blocks = (dims(2) - 1)/SZ + 1 y_stride = y_group_stride(self) select case (face) case (X_FACE) !$omp parallel do collapse(2) private(k, i, i_max) do z = 1, dims(3) do y_block = 1, n_y_blocks ! y-block is the fast-varying component of the group index k = y_stride*(z - 1) + y_block if (y_block == n_y_blocks) then i_max = n_mod else i_max = SZ end if do i = 1, i_max ! left face: spatially-varying Dirichlet from f_start at i=1 f%data(i, 1, k) = f_start%data(i, 1, k) ! right face: convective outflow associate (fd => f%data(i, dims(1), k), & fd1 => f%data(i, dims(1) - 1, k)) fd = fd - c_end*(fd - fd1) + flow_rate_diff_val end associate end do end do end do !$omp end parallel do case (Y_FACE) !$omp parallel do private(k_start, k_end) do z = 1, dims(3) ! bottom wall (y = 1) and top wall (y = ny) in OMP DIR_X ordering k_start = 1 + (z - 1)*y_stride k_end = (z - 1)*y_stride + n_y_blocks do j = 1, dims(1) f%data(1, j, k_start) = f_start%data(1, j, k_start) f%data(n_mod, j, k_end) = f_start%data(n_mod, j, k_end) end do end do !$omp end parallel do case default call abort_run('field_set_face_from_field: only X_FACE and Y_FACE & &supported.') end select end subroutine field_set_face_from_field_omp subroutine set_field_face_data_omp(self, f, data_start, data_end, face) !! [[m_base_backend(module):set_field_face_data(subroutine)]] implicit none class(omp_backend_t) :: self class(field_t), intent(inout) :: f real(dp), dimension(:, :), intent(in) :: data_start real(dp), dimension(:, :), optional, intent(in) :: data_end integer, intent(in) :: face integer :: dims(3), z, j, i, y_block, i_max, k, k_start, k_end, & n_mod, n_y_blocks, y_stride if (f%dir /= DIR_X) & call abort_run('set_field_face_data: only supported for DIR_X fields.') if (f%data_loc == NULL_LOC) & call abort_run('set_field_face_data: requires a valid data_loc.') dims = self%mesh%get_dims(f%data_loc) n_mod = mod(dims(2) - 1, SZ) + 1 n_y_blocks = (dims(2) - 1)/SZ + 1 y_stride = y_group_stride(self) select case (face) case (X_FACE) if (size(data_start, 1) /= dims(2) .or. size(data_start, 2) /= dims(3)) & call abort_run('set_field_face_data: X_FACE planes must be (ny, nz).') if (present(data_end)) then if (size(data_end, 1) /= dims(2) .or. size(data_end, 2) /= dims(3)) & call abort_run('set_field_face_data: X_FACE planes must be & &(ny, nz).') end if !$omp parallel do collapse(2) private(k, i, i_max) do z = 1, dims(3) do y_block = 1, n_y_blocks ! OMP DIR_X ordering: k = y_stride*(z - 1) + y_block k = y_stride*(z - 1) + y_block i_max = SZ if (y_block == n_y_blocks) i_max = n_mod do i = 1, i_max f%data(i, 1, k) = data_start((y_block - 1)*SZ + i, z) if (present(data_end)) & f%data(i, dims(1), k) = data_end((y_block - 1)*SZ + i, z) end do end do end do !$omp end parallel do case (Y_FACE) if (size(data_start, 1) /= dims(1) .or. size(data_start, 2) /= dims(3)) & call abort_run('set_field_face_data: Y_FACE planes must be (nx, nz).') if (present(data_end)) then if (size(data_end, 1) /= dims(1) .or. size(data_end, 2) /= dims(3)) & call abort_run('set_field_face_data: Y_FACE planes must be & &(nx, nz).') end if !$omp parallel do collapse(2) private(k_start, k_end) do z = 1, dims(3) do j = 1, dims(1) k_start = 1 + (z - 1)*y_stride k_end = (z - 1)*y_stride + n_y_blocks f%data(1, j, k_start) = data_start(j, z) if (present(data_end)) f%data(n_mod, j, k_end) = data_end(j, z) end do end do !$omp end parallel do case default call abort_run('set_field_face_data: only X_FACE and Y_FACE supported.') end select end subroutine set_field_face_data_omp subroutine field_add_const_x_face_omp(self, f, c, at_end) !! [[m_base_backend(module):field_add_const_x_face_base(subroutine)]] implicit none class(omp_backend_t) :: self class(field_t), intent(inout) :: f real(dp), intent(in) :: c logical, intent(in) :: at_end integer :: dims(3), k, i, z, i_plane, i_max, n_mod, n_y_blocks, & y_block, y_stride if (f%dir /= DIR_X) & call abort_run('field_add_const_x_face: only supported for DIR_X & &fields.') if (f%data_loc == NULL_LOC) & call abort_run('field_add_const_x_face: requires a valid data_loc.') dims = self%mesh%get_dims(f%data_loc) n_mod = mod(dims(2) - 1, SZ) + 1 n_y_blocks = (dims(2) - 1)/SZ + 1 y_stride = y_group_stride(self) if (at_end) then i_plane = dims(1) ! global x = nx (outlet) else i_plane = 1 ! global x = 1 (inlet) end if !$omp parallel do collapse(2) private(k, i, i_max) do z = 1, dims(3) do y_block = 1, n_y_blocks ! y-block is the fast-varying component of the group index k = y_stride*(z - 1) + y_block if (y_block == n_y_blocks) then i_max = n_mod else i_max = SZ end if do i = 1, i_max f%data(i, i_plane, k) = f%data(i, i_plane, k) + c end do end do end do !$omp end parallel do end subroutine field_add_const_x_face_omp subroutine field_add_face_from_field_omp(self, f, g, face, bc_start, bc_end) !! [[m_base_backend(module):field_add_face_from_field(subroutine)]] implicit none class(omp_backend_t) :: self class(field_t), intent(inout) :: f class(field_t), intent(in) :: g integer, intent(in) :: face, bc_start, bc_end integer :: dims(3), k, i, j, z, i_max, n_mod, n_y_blocks, y_stride integer :: y_block, k_start, k_end logical :: set_start, set_end if (f%dir /= DIR_X .or. g%dir /= DIR_X) & call abort_run('field_add_face_from_field: only supported for DIR_X & &fields.') if (f%data_loc == NULL_LOC) & call abort_run('field_add_face_from_field: requires a valid data_loc.') set_start = (bc_start == BC_DIRICHLET) set_end = (bc_end == BC_DIRICHLET) dims = self%mesh%get_dims(f%data_loc) n_mod = mod(dims(2) - 1, SZ) + 1 n_y_blocks = (dims(2) - 1)/SZ + 1 y_stride = y_group_stride(self) select case (face) case (X_FACE) !$omp parallel do collapse(2) private(k, i, i_max) do z = 1, dims(3) do y_block = 1, n_y_blocks ! y-block is the fast-varying component of the group index k = y_stride*(z - 1) + y_block if (y_block == n_y_blocks) then i_max = n_mod else i_max = SZ end if do i = 1, i_max if (set_start) f%data(i, 1, k) = f%data(i, 1, k) + g%data(i, 1, k) if (set_end) f%data(i, dims(1), k) = f%data(i, dims(1), k) & + g%data(i, dims(1), k) end do end do end do !$omp end parallel do case (Y_FACE) !$omp parallel do private(k_start, k_end) do z = 1, dims(3) k_start = 1 + (z - 1)*y_stride k_end = (z - 1)*y_stride + n_y_blocks do j = 1, dims(1) if (set_start) f%data(1, j, k_start) = f%data(1, j, k_start) & + g%data(1, j, k_start) if (set_end) f%data(n_mod, j, k_end) = f%data(n_mod, j, k_end) & + g%data(n_mod, j, k_end) end do end do !$omp end parallel do case default call abort_run('field_add_face_from_field: only X_FACE and Y_FACE & &supported.') end select end subroutine field_add_face_from_field_omp real(dp) function field_volume_integral_omp(self, f) result(s) !! volume integral of a field implicit none class(omp_backend_t) :: self class(field_t), intent(in) :: f real(dp) :: sum_p, sum_pncl integer :: dims(3), stacked, i, j, k, k_i, k_j, ierr if (f%data_loc == NULL_LOC) then call abort_run('You must set the data_loc before calling volume & &integral.') end if if (f%dir /= DIR_X) then call abort_run('Volume integral can only be called on DIR_X fields.') end if dims = self%mesh%get_dims(f%data_loc) stacked = (dims(2) - 1)/SZ + 1 sum_p = 0._dp !$omp parallel do collapse(2) reduction(+:sum_p) private(k, sum_pncl) do k_j = 1, stacked ! loop over stacked groups do k_i = 1, dims(3) k = k_j + (k_i - 1)*stacked sum_pncl = 0._dp do j = 1, dims(1) ! loop over only non-padded entries in the present group do i = 1, min(SZ, dims(2) - (k_j - 1)*SZ) sum_pncl = sum_pncl + f%data(i, j, k) end do end do sum_p = sum_p + sum_pncl end do end do !$omp end parallel do ! rank-local values s = sum_p call MPI_Allreduce(MPI_IN_PLACE, s, 1, MPI_X3D2_DP, MPI_SUM, & MPI_COMM_WORLD, ierr) end function field_volume_integral_omp subroutine copy_data_to_f_omp(self, f, data) class(omp_backend_t), intent(inout) :: self class(field_t), intent(inout) :: f real(dp), dimension(:, :, :), intent(in) :: data f%data = data end subroutine copy_data_to_f_omp subroutine copy_f_to_data_omp(self, data, f) class(omp_backend_t), intent(inout) :: self real(dp), dimension(:, :, :), intent(out) :: data class(field_t), intent(in) :: f data = f%data end subroutine copy_f_to_data_omp subroutine init_omp_poisson_fft(self, mesh, xdirps, ydirps, zdirps, lowmem) #ifdef WITH_2DECOMPFFT use m_omp_poisson_fft, only: create_omp_poisson_fft #endif implicit none class(omp_backend_t) :: self type(mesh_t), target, intent(in) :: mesh type(dirps_t), intent(in) :: xdirps, ydirps, zdirps logical, optional, intent(in) :: lowmem #ifdef WITH_2DECOMPFFT call create_omp_poisson_fft(self%poisson_fft, mesh, xdirps, ydirps, & zdirps, lowmem) #else call abort_run('This build does not support FFT based Poisson solver on & &the OpenMP backend!') #endif end subroutine init_omp_poisson_fft end module m_omp_backend