diff --git a/.gitignore b/.gitignore index 9b1ee24bd..628d4c306 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,7 @@ *~ a.out *.mod +*.smod build/ build.* *.o diff --git a/CMakeLists.txt b/CMakeLists.txt index 044a71aec..8fe562b4b 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -2,6 +2,11 @@ cmake_minimum_required(VERSION 3.10) project(x3d2 LANGUAGES Fortran) enable_testing() +# +# Set the Poisson solver choice +# +set(POISSON_SOLVER FFT CACHE STRING + "Select the Poisson solver: FFT or ITER") set(WITH_2DECOMPFFT ON CACHE BOOL "Enable Poisson based FFT solver on the OpenMP backend.") diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 06785145d..597585b84 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -6,6 +6,7 @@ set(SRC mesh.f90 ordering.f90 poisson_fft.f90 + poisson_cg_types.f90 solver.f90 tdsops.f90 time_integrator.f90 @@ -41,6 +42,17 @@ set(CUDASRC backend/cuda/sendrecv.f90 backend/cuda/tdsops.f90 ) +if (${POISSON_SOLVER} STREQUAL "ITER") + set(CGSRC + petsc/poisson_cg.f90 + poisson_cg.f90 + ) +else() + set(CGSRC + dummy/poisson_cg.f90 + poisson_cg.f90 + ) +endif() set(BACKENDSRC backend/omp/backend.f90 ) @@ -66,6 +78,7 @@ if(WITH_2DECOMPFFT) else() list(APPEND SRC ${DUMMYDECOMPSRC}) endif() +list(APPEND SRC ${CGSRC}) add_library(x3d2 STATIC ${SRC}) target_include_directories(x3d2 INTERFACE ${CMAKE_CURRENT_BINARY_DIR}) @@ -120,3 +133,16 @@ target_link_libraries(xcompact PRIVATE OpenMP::OpenMP_Fortran) find_package(MPI REQUIRED) target_link_libraries(x3d2 PRIVATE MPI::MPI_Fortran) target_link_libraries(xcompact PRIVATE MPI::MPI_Fortran) + +if (${POISSON_SOLVER} STREQUAL "FFT") + message(STATUS "Using the FFT Poisson solver") +elseif (${POISSON_SOLVER} STREQUAL "ITER") + message(STATUS "Using the iterative Poisson solver") + + # Find and link PETSc + find_package(PkgConfig REQUIRED) + pkg_check_modules(PETSC REQUIRED IMPORTED_TARGET petsc) + target_link_libraries(x3d2 PUBLIC PkgConfig::PETSC) +else() + message(FATAL_ERROR "Unrecognised Poisson solver: ${POISSON_SOLVER}") +endif() diff --git a/src/dummy/poisson_cg.f90 b/src/dummy/poisson_cg.f90 new file mode 100644 index 000000000..fcf4f65ce --- /dev/null +++ b/src/dummy/poisson_cg.f90 @@ -0,0 +1,92 @@ +module m_poisson_cg_backend + + use m_base_poisson_cg, only: poisson_solver_t, poisson_precon_t, laplace_operator_t + + use m_allocator, only: field_t + use m_base_backend, only: base_backend_t + use m_mesh, only: mesh_t + use m_vector_calculus, only: vector_calculus_t + + implicit none + + private + public :: init_solver + public :: init_precon + + type, extends(poisson_solver_t), public :: poisson_solver_impl + private + type(laplace_operator_t) :: lapl ! The high-order Laplacian operator + class(poisson_precon_t), allocatable :: precon ! The (low-order) preconditioner + contains + private + ! Solves the Poisson problem + procedure, public :: solve => solve + end type poisson_solver_impl + + type, extends(poisson_precon_t), public :: poisson_precon_impl + !! Wrapper definition of the Poisson preconditioner. + !! User code should use this class which will instantiate the + !! backend-specific code (determined at compile time). + private + contains + private + ! Applies the preconditioner to compute b=Px, mostly for testing purposes. + procedure, public :: apply => apply_precon + end type poisson_precon_impl + + interface poisson_precon_t + !! Public constructor for the Poisson preconditioner object. + module procedure init_precon + end interface poisson_precon_t + +contains + + module subroutine init_solver(solver, backend, mesh) + class(poisson_solver_t), allocatable, intent(out) :: solver + class(base_backend_t), target, intent(in) :: backend + type(mesh_t), intent(in) :: mesh + + allocate (poisson_solver_impl :: solver) + + select type (solver) + type is (poisson_solver_impl) + solver%precon = init_precon(backend) + solver%lapl = laplace_operator_t(backend, mesh) + class default + error stop "Dummy CG solver init failed" + end select + + error stop "This dummy module does not implement CG, recompile the code with PETSc" + end subroutine init_solver + + subroutine solve(self, p, f, backend) + class(poisson_solver_impl) :: self + class(field_t), intent(inout) :: p + class(field_t), intent(in) :: f + class(base_backend_t), intent(in) :: backend + + ! Call the backend-specific subroutine + error stop "Dummy CG backend called" + end subroutine solve + + function init_precon(backend) result(precon) + !! Public constructor for the Poisson preconditioner object. + class(base_backend_t), target, intent(in) :: backend + class(poisson_precon_t), allocatable :: precon + + allocate (poisson_precon_impl :: precon) + + end function init_precon + + subroutine apply_precon(self, p, b, backend) + ! Applies the preconditioner to compute b=Px, mostly for testing purposes. + class(poisson_precon_impl) :: self + class(field_t), intent(in) :: p ! Pressure solution + class(field_t), intent(inout) :: b ! The evaluated matrix-vector product + class(base_backend_t), intent(in) :: backend + + ! Call the backend-specific subroutine + error stop "Dummy CG backend called" + end subroutine apply_precon + +end module m_poisson_cg_backend diff --git a/src/mesh.f90 b/src/mesh.f90 index b0b3a6dd5..b79ade78f 100644 --- a/src/mesh.f90 +++ b/src/mesh.f90 @@ -189,7 +189,7 @@ subroutine decomposition_generic(grid, par) class(grid_t), intent(inout) :: grid class(par_t), intent(inout) :: par integer, allocatable, dimension(:, :, :) :: global_ranks - integer :: i, nproc_x, nproc_y, nproc_z + integer :: nproc_x, nproc_y, nproc_z if (par%is_root()) then print *, "Domain decomposition by x3d2 (generic)" @@ -207,8 +207,7 @@ subroutine decomposition_generic(grid, par) allocate (global_ranks(nproc_x, nproc_y, nproc_z)) ! set the corresponding global rank for each sub-domain - global_ranks = reshape([(i, i=0, par%nproc - 1)], & - shape=[nproc_x, nproc_y, nproc_z]) + global_ranks = par%compute_global_rank_layout() call par%compute_rank_pos_from_global(global_ranks) call grid%copy_vert2cell_dims(par) diff --git a/src/mesh_content.f90 b/src/mesh_content.f90 index c1851db6f..bb406067b 100644 --- a/src/mesh_content.f90 +++ b/src/mesh_content.f90 @@ -56,6 +56,7 @@ module m_mesh_content integer, dimension(3) :: pprev ! rank ID of the next rank in each direction contains procedure :: is_root ! returns if the current rank is the root rank + procedure :: compute_global_rank_layout ! determines the layout of ranks in a 3-D grid procedure :: compute_rank_pos_from_global ! fills in pnext, pprev and nrank_dir from global ranks map end type @@ -70,6 +71,19 @@ pure function is_root(self) result(is_root_rank) end function + pure function compute_global_rank_layout(self) result(global_rank_layout) + class(par_t), intent(in) :: self + integer, dimension(self%nproc_dir(1), self%nproc_dir(2), self%nproc_dir(3)) :: global_rank_layout + + integer :: i + + global_rank_layout = reshape([(i, i=0, self%nproc - 1)], & + shape=[self%nproc_dir(1), & + self%nproc_dir(2), & + self%nproc_dir(3)]) + + end function compute_global_rank_layout + pure subroutine compute_rank_pos_from_global(self, global_ranks) !! From the global rank maps, fills in the rank position as well !! as the previous and next rank in the `par` structure diff --git a/src/petsc/poisson_cg.f90 b/src/petsc/poisson_cg.f90 new file mode 100644 index 000000000..b2f0feb6b --- /dev/null +++ b/src/petsc/poisson_cg.f90 @@ -0,0 +1,695 @@ +!!! PETSc-based implementation of the iterative Poisson solver + +module m_cg_types + !! Types module providing the context type required by the PETSc matrix-free + !! operator. + + use petsc + + use m_common, only: CELL + use m_base_backend, only: base_backend_t + use m_allocator, only: field_t + use m_base_poisson_cg, only: laplace_operator_t + + implicit none + + private + + type, public :: mat_ctx_t + class(base_backend_t), pointer :: backend + type(laplace_operator_t) :: lapl + class(field_t), pointer :: xfield + class(field_t), pointer :: ffield + end type mat_ctx_t + + interface mat_ctx_t + module procedure init_ctx + end interface mat_ctx_t + +contains + + function init_ctx(backend, lapl, dir) result(ctx) + + class(base_backend_t), target, intent(in) :: backend + class(laplace_operator_t), intent(in) :: lapl + integer, intent(in) :: dir + type(mat_ctx_t) :: ctx + + ctx%xfield => backend%allocator%get_block(dir, CELL) + ctx%ffield => backend%allocator%get_block(dir, CELL) + ctx%backend => backend + ctx%lapl = lapl + + end function init_ctx + +end module m_cg_types + +module m_poisson_cg_backend + !! Module implementing a Poisson solver based on the (preconditioned) + !! Conjugate Gradient method using PETSc. + + use m_base_poisson_cg, only: poisson_solver_t, poisson_precon_t, laplace_operator_t + + use petsc + + use m_cg_types + + use m_common, only: dp, DIR_Z, DIR_C, CELL + use m_allocator, only: field_t + use m_base_backend, only: base_backend_t + use m_mesh, only: mesh_t + + implicit none + + private + public :: init_solver + public :: init_precon + + type, extends(poisson_solver_t), public :: poisson_solver_impl + !! Conjugate Gradient based Poisson solver using PETSc as a backend. + !! Supports any decomposition that is also supported by the underlying + !! finite difference schemes. + type(mat_ctx_t) :: ctx + type(tKSP) :: ksp ! The solver + type(tMat) :: Amat ! The operator matrix + type(tVec) :: fvec ! The RHS vector + type(tVec) :: pvec ! The solution vector + type(laplace_operator_t) :: lapl + class(poisson_precon_t), allocatable :: precon + contains + procedure, public :: solve => solve_petsc + procedure :: create_operator + procedure :: create_vectors + procedure :: create_solver + end type poisson_solver_impl + + type, extends(poisson_precon_t), public :: poisson_precon_impl + !! The PETSc implementation of the Poisson preconditioner, implements a 2nd + !! order finite difference approximation of the Laplacian. + type(tMat) :: Pmat ! The preconditioner matrix + type(tDM), pointer :: da + contains + procedure :: init => init_precon_petsc + procedure :: apply => petsc_apply_precon + end type poisson_precon_impl + + interface MatCreateShell + !! Defines the interface to the external (PETSc) function to create a + !! matrix-free operator. + subroutine MatCreateShell(comm, nrow_l, ncol_l, nrow_g, ncol_g, ctx, M, & + ierr) + use petsc + use m_cg_types + integer :: comm + integer :: nrow_l ! Local number of rows + integer :: ncol_l ! Local number of columns + integer :: nrow_g ! Global number of rows + integer :: ncol_g ! Global number of columns + type(mat_ctx_t) :: ctx ! The shell matrix context + type(tMat) :: M ! The matrix object + integer :: ierr + end subroutine MatCreateShell + end interface MatCreateShell + + interface MatShellSetContext + !! Defines the interface to the external (PETSc) function to store + !! application-dependent information required by the matrix-free operator. + subroutine MatShellSetContext(M, ctx, ierr) + use petsc + use m_cg_types + type(tMat) :: M ! The matrix object + type(mat_ctx_t) :: ctx ! The shell matrix context + integer :: ierr + end subroutine MatShellSetContext + end interface MatShellSetContext + + interface MatShellGetContext + !! Defines the interface to the external (PETSc) function to retrieve + !! application-dependent information from the matrix-free operator. + subroutine MatShellGetContext(M, ctx, ierr) + use petsc + use m_cg_types + type(tMat) :: M ! The matrix object + type(mat_ctx_t) :: ctx ! The shell matrix context + integer :: ierr + end subroutine MatShellGetContext + end interface MatShellGetContext + + interface MatShellSetOperation + !! Defines the interface to the external (PETSc) function to set the + !! matrix-free operator procedure that evaluates `f = Mx`. + subroutine MatShellSetOperation(M, OP, fn, ierr) + use petsc + type(tMat) :: M + integer :: OP + interface + subroutine fn(M, x, f, ierr) + use petsc + type(tMat) :: M ! The operator + type(tVec) :: x ! The input vector + type(tVec) :: f ! The output vector + integer :: ierr ! The error code + end subroutine fn + end interface + integer :: ierr + end subroutine MatShellSetOperation + end interface MatShellSetOperation + + type(mat_ctx_t), save :: ctx_global ! XXX: This sucks! + type(tDM), pointer, save :: da_ptr_global ! XXX: This sucks! + type(tDM), target, save :: da_tgt + +contains + + function init_precon(backend) result(precon) + ! Constructs the PETSc preconditioner implementation + class(base_backend_t), intent(in) :: backend + class(poisson_precon_t), allocatable :: precon + + allocate (poisson_precon_impl :: precon) + select type (precon) + type is (poisson_precon_impl) + call precon%init(backend) + class default + error stop "IMPOSSIBLE" + end select + + end function init_precon + + subroutine init_precon_petsc(self, backend) + ! Initialise the PETSc implementation of the preconditioner object +#include "petsc/finclude/petscmat.h" + + class(poisson_precon_impl), intent(out) :: self + class(base_backend_t), intent(in) :: backend + + type(tMatNullSpace) :: nsp + integer :: ierr + + integer, parameter :: nnb = 26 ! Number of neighbours (27-point stencil has 26 neighbours) + + integer, dimension(3) :: dims + integer :: i, j, k + real(dp) :: dx, dy, dz + + MatStencil :: row(4, 1) + MatStencil :: col(4, 27) + real(dp), dimension(nnb + 1) :: coeffs + + logical :: initialised + + integer, dimension(3), parameter :: stencil1d = [1, -2, 1] + real(dp), dimension(3, 3, 3) :: stencil3d + ! integer, dimension(3, 3, 3) :: stencil3d_x, stencil3d_y, stencil3d_z + + integer, dimension(:), allocatable :: procx, procy, procz + integer, dimension(:), allocatable :: nxglobal, nyglobal, nzglobal + integer, dimension(:), allocatable :: lx, ly, lz + integer, dimension(:, :, :), allocatable :: procgrid + integer, parameter :: dof = 1 ! Variables per point in the linear system (P) + integer, parameter :: stencil_width = 1 + integer :: ctr + integer :: ii, jj, kk + integer :: ifirst, jfirst, kfirst + integer :: ilast, jlast, klast + + ! Ensure PETSc is initialised + call PetscInitialized(initialised, ierr) + if (.not. initialised) then + if (backend%mesh%par%nrank == 0) then + print *, "Initialising PETSc" + end if + call PetscInitialize(PETSC_NULL_CHARACTER, ierr) + end if + if (backend%mesh%par%nrank == 0) then + print *, "PETSc Initialised" + end if + + self%da => da_tgt + + ! Create an explicit preconditioner matrix + associate (mesh => backend%mesh) + + allocate(procgrid(mesh%par%nproc_dir(1), mesh%par%nproc_dir(2), mesh%par%nproc_dir(3))) + procgrid = mesh%par%compute_global_rank_layout() + + procx = procgrid(:, mesh%par%nrank_dir(2) + 1, mesh%par%nrank_dir(3) + 1) + procy = procgrid(mesh%par%nrank_dir(1) + 1, :, mesh%par%nrank_dir(3) + 1) + procz = procgrid(mesh%par%nrank_dir(1) + 1, mesh%par%nrank_dir(2) + 1, :) + deallocate(procgrid) + + dims = mesh%get_dims(CELL) + allocate(nxglobal(mesh%par%nproc)) + allocate(nyglobal(mesh%par%nproc)) + allocate(nzglobal(mesh%par%nproc)) + nxglobal = 0; nyglobal = 0; nzglobal = 0 + nxglobal(mesh%par%nrank + 1) = dims(1) + nyglobal(mesh%par%nrank + 1) = dims(2) + nzglobal(mesh%par%nrank + 1) = dims(3) + + call MPI_Allreduce(MPI_IN_PLACE, nxglobal, mesh%par%nproc, & + MPI_INTEGER, MPI_SUM, MPI_COMM_WORLD, ierr) + call MPI_Allreduce(MPI_IN_PLACE, nyglobal, mesh%par%nproc, & + MPI_INTEGER, MPI_SUM, MPI_COMM_WORLD, ierr) + call MPI_Allreduce(MPI_IN_PLACE, nzglobal, mesh%par%nproc, & + MPI_INTEGER, MPI_SUM, MPI_COMM_WORLD, ierr) + + lx = nxglobal(procx + 1) + ly = nyglobal(procy + 1) + lz = nzglobal(procz + 1) + + deallocate(nxglobal) + deallocate(nyglobal) + deallocate(nzglobal) + + dims = mesh%get_global_dims(CELL) + call DMDACreate3d(PETSC_COMM_WORLD, & + DM_BOUNDARY_PERIODIC, DM_BOUNDARY_PERIODIC, DM_BOUNDARY_PERIODIC, & + DMDA_STENCIL_BOX, & + dims(1), dims(2), dims(3), & + mesh%par%nproc_dir(1), mesh%par%nproc_dir(2), mesh%par%nproc_dir(3), & + dof, & + stencil_width, & + lx, ly, lz, & + self%da, & + ierr) + call DMSetFromOptions(self%da, ierr) + call DMSetUp(self%da, ierr) + end associate + + call DMCreateMatrix(self%da, self%Pmat, ierr) + call MatSetFromOptions(self%Pmat, ierr) + call MatSetUp(self%Pmat, ierr) + + ! Set up stencils + do k = 1, 3 + do j = 1, 3 + stencil3d(:, j, k) = stencil1d + end do + end do + stencil3d(:, 2, 1) = 2 * stencil3d(:, 2, 1) + stencil3d(:, 1, 2) = 2 * stencil3d(:, 1, 2) + stencil3d(:, 2, 2) = 4 * stencil3d(:, 2, 2) + stencil3d(:, 3, 2) = 2 * stencil3d(:, 3, 2) + stencil3d(:, 2, 3) = 2 * stencil3d(:, 2, 3) + + ! stencil3d_x = stencil3d + ! stencil3d_y = reshape(stencil3d, shape=[3, 3, 3], order=[2, 1, 3]) + ! stencil3d_z = reshape(stencil3d, shape=[3, 3, 3], order=[3, 2, 1]) + + !! Stencil according to Marcin Krotkiewski & Marcin Dabrowski, "Efficient 3D + !! stencil computations using CUDA", Parallel Computing (2019) + + ! ! Central + ! stencil3d(2, 2, 2) = 8.0_dp / 3.0_dp + + ! ! Face centre + ! stencil3d([1, 3], 2, 2) = 0 + ! stencil3d(2, [1, 3], 2) = 0 + ! stencil3d(2, 2, [1, 3]) = 0 + + ! ! Edge centre + ! stencil3d([1, 3], [1, 3], 2) = -1.0_dp / 6.0_dp + ! stencil3d([1, 3], 2, [1, 3]) = -1.0_dp / 6.0_dp + ! stencil3d(2, [1, 3], [1, 3]) = -1.0_dp / 6.0_dp + + ! ! Corner + ! stencil3d([1, 3], [1, 3], [1, 3]) = -1.0_dp / 12.0_dp + + ! !! Stencil according to PETSc bench_kspsolve + + ! ! Central + ! stencil3d(2, 2, 2) = 44.0_dp / 13.0_dp + + ! ! Face centre + ! stencil3d([1, 3], 2, 2) = -3.0_dp / 13.0_dp + ! stencil3d(2, [1, 3], 2) = -3.0_dp / 13.0_dp + ! stencil3d(2, 2, [1, 3]) = -3.0_dp / 13.0_dp + + ! ! Edge centre + ! stencil3d([1, 3], [1, 3], 2) = -3.0_dp / 26.0_dp + ! stencil3d([1, 3], 2, [1, 3]) = -3.0_dp / 26.0_dp + ! stencil3d(2, [1, 3], [1, 3]) = -3.0_dp / 26.0_dp + + ! ! Corner + ! stencil3d([1, 3], [1, 3], [1, 3]) = -1.0_dp / 13.0_dp + + !! 7-point star stencil + stencil3d = 0 + stencil3d(1, 2, 2) = 1 + stencil3d(3, 2, 2) = 1 + stencil3d(2, 1, 2) = 1 + stencil3d(2, 3, 2) = 1 + stencil3d(2, 2, 1) = 1 + stencil3d(2, 2, 3) = 1 + stencil3d(2, 2, 2) = -6 + + ! !! 27-point stencil + + ! stencil3d = 1.0_dp + ! stencil3d(2, 2, 2) = -26.0_dp + ! stencil3d = stencil3d / 9.0_dp + + ! Set the Poisson coefficients + associate (mesh => backend%mesh) + dims = mesh%get_dims(CELL) + dx = mesh%geo%d(1); dy = mesh%geo%d(2); dz = mesh%geo%d(3) + print *, dx, dy, dz + + call DMDAGetCorners(self%da, ifirst, jfirst, kfirst, ilast, jlast, klast, ierr) + ilast = ifirst + (ilast - 1) + jlast = jfirst + (jlast - 1) + klast = kfirst + (klast - 1) + print *, ifirst, ilast + print *, jfirst, jlast + print *, kfirst, klast + do k = kfirst, klast + do j = jfirst, jlast + do i = ifirst, ilast + + row(MatStencil_i, 1) = i + row(MatStencil_j, 1) = j + row(MatStencil_k, 1) = k + ctr = 1 + do kk = -1, 1 + do jj = -1, 1 + do ii = -1, 1 + col(MatStencil_i, ctr) = i + ii + col(MatStencil_j, ctr) = j + jj + col(MatStencil_k, ctr) = k + kk + ctr = ctr + 1 + end do + end do + end do + ! coeffs = reshape(stencil3d_x / dx**2 + stencil3d_y / dy**2 + stencil3d_z / dz**2, shape=[27]) / 16.0_dp + coeffs = reshape(stencil3d / dx**2, shape=[27]) + + ! Push to matrix + call MatSetValuesStencil(self%Pmat, 1, row, nnb + 1, col, & + coeffs, INSERT_VALUES, ierr) + end do + end do + end do + end associate + ! print *, coeffs + + call MatAssemblyBegin(self%Pmat, MAT_FINAL_ASSEMBLY, ierr) + call MatAssemblyEnd(self%Pmat, MAT_FINAL_ASSEMBLY, ierr) + + call MatNullSpaceCreate(PETSC_COMM_WORLD, PETSC_TRUE, 0, PETSC_NULL_VEC, nsp, ierr) + call MatSetnullSpace(self%Pmat, nsp, ierr) + call MatNullSpaceDestroy(nsp, ierr) + + ! call MatView(self%Pmat, PETSC_VIEWER_STDOUT_SELF, ierr) + + end subroutine init_precon_petsc + + module subroutine petsc_apply_precon(self, p, b, backend) + class(poisson_precon_impl) :: self + class(field_t), intent(in) :: p + class(field_t), intent(inout) :: b + class(base_backend_t), intent(in) :: backend + + type(tVec) :: pVec, bVec + integer :: ierr + + integer :: n + + n = product(backend%mesh%get_dims(CELL)) + call create_vec(pVec, self%da) + call create_vec(bVec, self%da) + + call copy_field_to_vec(pVec, p, backend, self%da) + call MatMult(self%PMat, pVec, bVec, ierr) + call copy_vec_to_field(b, bVec, backend, self%da) + + call VecDestroy(pVec, ierr) + call VecDestroy(bVec, ierr) + + end subroutine petsc_apply_precon + + module subroutine solve_petsc(self, p, f, backend) + class(poisson_solver_impl) :: self + class(field_t), intent(inout) :: p ! Pressure solution + class(field_t), intent(in) :: f ! Poisson RHS + class(base_backend_t), intent(in) :: backend + + integer :: ierr + + associate (precon => self%precon) + select type(precon) + type is (poisson_precon_impl) + da_ptr_global => precon%da + class default + error stop + end select + end associate + ctx_global = mat_ctx_t(backend, self%lapl, DIR_Z) + call copy_field_to_vec(self%fvec, f, backend, da_ptr_global) + call KSPSolve(self%ksp, self%fvec, self%pvec, ierr) + call copy_vec_to_field(p, self%pvec, backend, da_ptr_global) + + end subroutine solve_petsc + + module subroutine init_solver(solver, backend, mesh) + !! Public constructor for the poisson_cg_t type. + class(poisson_solver_t), allocatable, intent(out) :: solver + class(base_backend_t), target, intent(in) :: backend + type(mesh_t), intent(in) :: mesh + + allocate (poisson_solver_impl :: solver) + + select type (solver) + type is (poisson_solver_impl) + solver%precon = init_precon(backend) + solver%lapl = laplace_operator_t(backend, mesh) + call init_petsc_cg(solver, backend) + class default + ! This should be impossible + error stop "Failure in allocating PETSc Poisson solver -- this indicates a serious problem" + end select + end subroutine init_solver + + subroutine init_petsc_cg(self, backend) + !! Private constructor for the poisson_cg_t type. + type(poisson_solver_impl), intent(inout) :: self + class(base_backend_t), target, intent(in) :: backend + + integer :: n ! Local problem size + + integer :: ierr + + logical :: initialised + + call PetscInitialized(initialised, ierr) + if (.not. initialised) then + if (backend%mesh%par%nrank == 0) then + print *, "Initialising PETSc" + end if + call PetscInitialize(PETSC_NULL_CHARACTER, ierr) + end if + if (backend%mesh%par%nrank == 0) then + print *, "PETSc Initialised" + end if + + ! Determine local problem size + n = product(backend%mesh%get_dims(CELL)) + + ! Initialise preconditioner and operator matrices + ! XXX: Add option to use preconditioner as operator (would imply low-order + ! solution)? + call self%create_operator(n) + + ! Initialise RHS and solution vectors + call self%create_vectors(n) + + ! Create the linear system + call self%create_solver() + + end subroutine init_petsc_cg + + subroutine create_operator(self, n) + ! Set the PETSc MATVEC to use the x3d2 high-order Laplacian operator + class(poisson_solver_impl) :: self + integer, intent(in) :: n ! The local problem size + + type(tMatNullSpace) :: nsp + integer :: ierr + + call MatCreateShell(PETSC_COMM_WORLD, n, n, PETSC_DETERMINE, & + PETSC_DETERMINE, self%ctx, self%Amat, ierr) + call MatShellSetContext(self%Amat, self%ctx, ierr) ! Is this necessary? + call MatShellSetOperation(self%Amat, MATOP_MULT, poissmult_petsc, ierr) + call MatSetUp(self%Amat, ierr) + + call MatAssemblyBegin(self%Amat, MAT_FINAL_ASSEMBLY, ierr) + call MatAssemblyEnd(self%Amat, MAT_FINAL_ASSEMBLY, ierr) + + call MatNullSpaceCreate(PETSC_COMM_WORLD, PETSC_TRUE, 0, PETSC_NULL_VEC, nsp, ierr) + call MatSetnullSpace(self%Amat, nsp, ierr) + call MatNullSpaceDestroy(nsp, ierr) + + end subroutine create_operator + + subroutine create_vectors(self, n) + ! Allocates the pressure and forcing vectors. + + class(poisson_solver_impl) :: self ! The Poisson solver + integer, intent(in) :: n ! The local vector size + + associate(precon => self%precon) + select type(precon) + type is (poisson_precon_impl) + call create_vec(self%fvec, precon%da) + call create_vec(self%pvec, precon%da) + class default + error stop + end select + end associate + + end subroutine create_vectors + + subroutine create_vec(v, dm) + ! Utility subroutine to allocate a PETSc vector. + + type(tVec), intent(out) :: v ! The vector + ! integer, intent(in) :: n ! The local vector size + type(tDM), intent(in) :: dm + + integer :: ierr + + ! call VecCreate(PETSC_COMM_WORLD, v, ierr) + ! call VecSetSizes(v, n, PETSC_DETERMINE, ierr) + ! call VecSetFromOptions(v, ierr) + + call DMCreateGlobalVector(dm, v, ierr) + + call VecAssemblyBegin(v, ierr) + call VecAssemblyEnd(v, ierr) + + end subroutine create_vec + + subroutine create_solver(self) + ! Sets up the PETSc linear solver. + + class(poisson_solver_impl) :: self + + integer :: ierr + + associate (precon => self%precon) + select type (precon) + type is (poisson_precon_impl) + call KSPCreate(PETSC_COMM_WORLD, self%ksp, ierr) + call KSPSetOperators(self%ksp, self%Amat, precon%Pmat, ierr) + ! call KSPSetOperators(self%ksp, precon%Pmat, precon%Pmat, ierr) + call KSPSetFromOptions(self%ksp, ierr) + call KSPSetInitialGuessNonzero(self%ksp, PETSC_TRUE, ierr) + class default + error stop "Poisson preconditioner type is wrong" + end select + end associate + + end subroutine create_solver + + subroutine poissmult_petsc(M, x, f, ierr) + !! Computes the action of the Poisson operator, i.e. `f = Mx` where `M` is + !! the discrete Laplacian. + type(tMat) :: M ! The operator + type(tVec) :: x ! The input vector + type(tVec) :: f ! The output vector + integer :: ierr ! The error code + + type(mat_ctx_t) :: ctx + + ierr = 0; + ! XXX: Fixme + ! call MatShellGetContext(M, ctx, ierr) + associate (matrix => M); end associate ! Silence unused argument + ! print *, ctx%foo + ctx = ctx_global + + call copy_vec_to_field(ctx%xfield, x, ctx%backend, da_ptr_global) + call ctx%lapl%apply(ctx%ffield, ctx%xfield) + call copy_field_to_vec(f, ctx%ffield, ctx%backend, da_ptr_global) + + end subroutine poissmult_petsc + + subroutine copy_vec_to_field(f, v, backend, da) + !! Copies the contents of a PETSc vector into an x3d2 field object + ! XXX: This can be avoided if a field can wrap the vector memory + class(field_t), intent(inout) :: f ! The destination field + type(tVec) :: v ! The source vector + class(base_backend_t), intent(in) :: backend + type(tDM), intent(in) :: da + + ! real(dp), dimension(:), pointer :: vdata + real(dp), dimension(:, :, :), pointer :: vdata3d + integer :: ierr + + integer, dimension(3) :: dims + integer :: nx, ny, nz + + dims = backend%mesh%get_dims(CELL) + nx = dims(1) + ny = dims(2) + nz = dims(3) + + ! Local copy + ! call VecGetArrayReadF90(v, vdata, ierr) + ! if (nx*ny*nz /= size(vdata)) then + ! print *, "Vector and field sizes are incompatible (padding?)" + ! stop 1 + ! end if + ! vdata3d(1:nx, 1:ny, 1:nz) => vdata(:) ! Get a 3D representation of the vector + ! call backend%set_field_data(f, vdata3d, DIR_C) + ! call VecRestoreArrayReadF90(v, vdata, ierr) + ! nullify (vdata3d) + call DMDAVecGetArrayReadF90(da, v, vdata3d, ierr) + call backend%set_field_data(f, vdata3d, DIR_C) + call DMDAVecRestoreArrayReadF90(da, v, vdata3d, ierr) + + ! Halo exchange + + end subroutine copy_vec_to_field + + subroutine copy_field_to_vec(v, f, backend, da) + !! Copies the contents of an x3d2 field object into a PETSc vector + ! XXX: This can be avoided if a field can wrap the vector memory + type(tVec) :: v ! The destination vector. + class(field_t), intent(in) :: f ! The source field. + class(base_backend_t), intent(in) :: backend + type(tDM), intent(in) :: da + + ! real(dp), dimension(:), pointer :: vdata + real(dp), dimension(:, :, :), pointer :: vdata3d + integer :: ierr + + integer, dimension(3) :: dims + integer :: nx, ny, nz + + dims = backend%mesh%get_dims(CELL) + nx = dims(1) + ny = dims(2) + nz = dims(3) + + ! Local copy + ! call VecGetArrayF90(v, vdata, ierr) + ! if (nx*ny*nz /= size(vdata)) then + ! print *, "Vector and field sizes are incompatible (padding?)" + ! stop 1 + ! end if + ! vdata3d(1:nx, 1:ny, 1:nz) => vdata(:) ! Get a 3D representation of the vector + ! call backend%get_field_data(vdata3d, f, DIR_C) + ! call VecRestoreArrayF90(v, vdata, ierr) + ! nullify (vdata3d) + call DMDAVecGetArrayF90(da, v, vdata3d, ierr) + call backend%get_field_data(vdata3d, f, DIR_C) + call DMDAVecRestoreArrayF90(da, v, vdata3d, ierr) + + ! Halo exchange + call VecAssemblyBegin(v, ierr) + call VecAssemblyEnd(v, ierr) + + end subroutine copy_field_to_vec + +end module m_poisson_cg_backend diff --git a/src/poisson_cg.f90 b/src/poisson_cg.f90 new file mode 100644 index 000000000..596c01aea --- /dev/null +++ b/src/poisson_cg.f90 @@ -0,0 +1,71 @@ +module m_poisson_cg + !! Module defining a Poisson solver based on the (preconditioned) Conjugate + !! Gradient method. + + use m_base_poisson_cg, only: poisson_solver_t + use m_poisson_cg_backend, only: init_solver + + use m_allocator, only: allocator_t, field_t + use m_base_backend, only: base_backend_t + use m_mesh, only: mesh_t + + implicit none + + private + + type, public :: poisson_cg_t + !! Conjugate Gradient based Poisson solver. + !! Supports any decomposition that is also supported by the underlying + !! finite difference schemes. + + ! Prevent default access to components of type. + private + class(poisson_solver_t), allocatable :: solver + + contains + private + ! Solves the Poisson problem + procedure, public :: solve + procedure, public :: init => init_internal + end type poisson_cg_t + + interface poisson_cg_t + !! Public constructor for the poisson_cg_t type. + module procedure init_cg + end interface poisson_cg_t + +contains + + subroutine solve(self, p, f, backend) + ! Solves the Poisson problem + class(poisson_cg_t) :: self + class(field_t), intent(inout) :: p ! Pressure solution + class(field_t), intent(in) :: f ! Poisson RHS + class(base_backend_t), intent(in) :: backend + + call self%solver%solve(p, f, backend) + end subroutine solve + + function init_cg(backend, mesh) result(solver) + !! Initialises the conjugate gradient (CG) solver. + !! XXX: The solver implementation is responsible for initialising the + !! preconditioner, i.e. it should at some point during its initialisation + !! do the equivalent of: `call self%precon = poisson_precon_t(backend)`. + class(base_backend_t), target, intent(in) :: backend + type(mesh_t), intent(in) :: mesh + type(poisson_cg_t) :: solver + + call solver%init(backend, mesh) + + end function init_cg + + subroutine init_internal(self, backend, mesh) + class(poisson_cg_t), intent(out) :: self + class(base_backend_t), target, intent(in) :: backend + type(mesh_t), intent(in) :: mesh + + call init_solver(self%solver, backend, mesh) + + end subroutine init_internal + +end module m_poisson_cg diff --git a/src/poisson_cg_types.f90 b/src/poisson_cg_types.f90 new file mode 100644 index 000000000..204a8d99a --- /dev/null +++ b/src/poisson_cg_types.f90 @@ -0,0 +1,201 @@ +module m_base_poisson_cg + + use m_allocator, only: field_t + use m_base_backend, only: base_backend_t + use m_common, only: dp, & + RDR_X2Z, RDR_Y2Z, RDR_C2Z, & + RDR_Z2X, RDR_Z2Y, RDR_Z2C, & + DIR_X, DIR_Y, DIR_Z, DIR_C, & + CELL, VERT, & + BC_PERIODIC + use m_mesh, only: mesh_t + use m_tdsops, only: tdsops_t, dirps_t + use m_vector_calculus, only: vector_calculus_t + + implicit none + + private + + type, public :: laplace_operator_t + !! Operator that computes the Laplacian of a field. + private + type(dirps_t) :: xdirps, ydirps, zdirps + type(vector_calculus_t) :: vector_calculus + contains + procedure :: apply => poissmult + end type laplace_operator_t + + interface laplace_operator_t + !! Public constructor for the laplace_operator_t type. + procedure init_lapl + end interface laplace_operator_t + + type, abstract, public :: poisson_precon_t + !! Wrapper definition of the Poisson preconditioner. + !! User code should use this class which will instantiate the + !! backend-specific code (determined at compile time). + private + contains + private + ! Applies the preconditioner to compute b=Px, mostly for testing purposes. + procedure(apply_precon), public, deferred :: apply + end type poisson_precon_t + + abstract interface + subroutine apply_precon(self, p, b, backend) + ! Applies the preconditioner to compute b=Px, mostly for testing purposes. + import poisson_precon_t + import field_t + import base_backend_t + implicit none + class(poisson_precon_t) :: self + class(field_t), intent(in) :: p ! Pressure solution + class(field_t), intent(inout) :: b ! The evaluated matrix-vector product + class(base_backend_t), intent(in) :: backend + end subroutine apply_precon + end interface + + type, abstract, public :: poisson_solver_t + !! Base definition of the backend-specific implementation of the iterative + !! Poisson solver. + contains + private + ! Solves the Poisson problem + procedure(solve), public, deferred :: solve + end type poisson_solver_t + + abstract interface + subroutine solve(self, p, f, backend) + import poisson_solver_t + import field_t + import base_backend_t + implicit none + class(poisson_solver_t) :: self + class(field_t), intent(inout) :: p + class(field_t), intent(in) :: f + class(base_backend_t), intent(in) :: backend + end subroutine solve + end interface + +contains + + function init_lapl(backend, mesh) result(lapl) + !! Public constructor for the laplace_operator_t type. + class(base_backend_t), target, intent(in) :: backend + type(mesh_t), intent(in) :: mesh + type(laplace_operator_t) :: lapl + + ! TODO: read these from config + character(len=*), parameter :: stagder_scheme = "compact6" + character(len=*), parameter :: interpl_scheme = "classic" + + lapl%vector_calculus = vector_calculus_t(backend) + + lapl%xdirps%dir = DIR_X; lapl%ydirps%dir = DIR_Y; lapl%zdirps%dir = DIR_Z + + call allocate_lapl_tdsops(lapl%xdirps, backend, mesh, stagder_scheme, interpl_scheme) + call allocate_lapl_tdsops(lapl%ydirps, backend, mesh, stagder_scheme, interpl_scheme) + call allocate_lapl_tdsops(lapl%zdirps, backend, mesh, stagder_scheme, interpl_scheme) + + end function init_lapl + + subroutine allocate_lapl_tdsops(dirps, backend, mesh, stagder_scheme, interpl_scheme) + type(dirps_t), intent(inout) :: dirps + class(base_backend_t), intent(in) :: backend + type(mesh_t), intent(in) :: mesh + character(len=*), intent(in) :: stagder_scheme, interpl_scheme + + integer :: dir, bc_start, bc_end, n_vert, n_cell + real(dp) :: d + + dir = dirps%dir + bc_start = mesh%grid%BCs(dir, 1) + bc_end = mesh%grid%BCs(dir, 2) + d = mesh%geo%d(dir) + + n_vert = mesh%get_n(dir, VERT) + n_cell = mesh%get_n(dir, CELL) + + call backend%alloc_tdsops(dirps%interpl_v2p, n_cell, d, 'interpolate', & + interpl_scheme, bc_start, bc_end, from_to='v2p') + call backend%alloc_tdsops(dirps%interpl_p2v, n_vert, d, 'interpolate', & + interpl_scheme, bc_start, bc_end, from_to='p2v') + call backend%alloc_tdsops(dirps%stagder_v2p, n_cell, d, 'stag-deriv', & + stagder_scheme, bc_start, bc_end, from_to='v2p') + call backend%alloc_tdsops(dirps%stagder_p2v, n_vert, d, 'stag-deriv', & + stagder_scheme, bc_start, bc_end, from_to='p2v') + + end subroutine allocate_lapl_tdsops + + subroutine poissmult(self, f, p) + !! Computes the action of the Laplace operator, i.e. `f = Ax` where `A` is + !! the discrete Laplacian. + class(laplace_operator_t) :: self + class(field_t), intent(inout) :: f ! The output field + class(field_t), intent(in) :: p ! The input field + + class(field_t), pointer :: f_z, p_z + + integer :: reorder_op, reorder_op2z + + if (p%dir == DIR_Z .and. f%dir == DIR_Z) then + call poissmult_dirz(self, f, p) + else + if (p%dir /= f%dir) then + error stop "Currently orientations of P and F must match" + end if + if (f%dir == DIR_X) then + reorder_op2z = RDR_X2Z + reorder_op = RDR_Z2X + else if (f%dir == DIR_Y) then + reorder_op2z = RDR_Y2Z + reorder_op = RDR_Z2Y + else if (f%dir == DIR_C) then + reorder_op2z = RDR_C2Z + reorder_op = RDR_Z2C + else + error stop "Unsupported Poisson orientation" + end if + + f_z => self%vector_calculus%backend%allocator%get_block(DIR_Z, CELL) + p_z => self%vector_calculus%backend%allocator%get_block(DIR_Z, CELL) + + call self%vector_calculus%backend%reorder(p_z, p, reorder_op2z) + + call poissmult_dirz(self, f_z, p_z) + + call self%vector_calculus%backend%reorder(f, f_z, reorder_op) + + call self%vector_calculus%backend%allocator%release_block(f_z) + call self%vector_calculus%backend%allocator%release_block(p_z) + end if + + end subroutine poissmult + + subroutine poissmult_dirz(lapl, f, p) + class(laplace_operator_t), intent(in) :: lapl + class(field_t), intent(inout) :: f ! The output field + class(field_t), intent(in) :: p ! The input field + + if (p%dir /= DIR_Z .or. f%dir /= DIR_Z) then + error stop "Currently orientations of P and F must be in Z" + end if + if (p%data_loc /= CELL .or. f%data_loc /= CELL) then + error stop "The pressure Poisson equation must be evaluated at cell centres" + end if + + ! call lapl%vector_calculus%divgrad(f, p, & + ! lapl%xdirps%stagder_p2v, lapl%xdirps%interpl_p2v, & + ! lapl%ydirps%stagder_p2v, lapl%ydirps%interpl_p2v, & + ! lapl%zdirps%stagder_p2v, lapl%zdirps%interpl_p2v, & + ! lapl%xdirps%stagder_v2p, lapl%xdirps%interpl_v2p, & + ! lapl%ydirps%stagder_v2p, lapl%ydirps%interpl_v2p, & + ! lapl%zdirps%stagder_v2p, lapl%zdirps%interpl_v2p) + call lapl%vector_calculus%divgrad_stag(f, p, & + lapl%xdirps%stagder_p2v, lapl%ydirps%stagder_p2v, lapl%zdirps%stagder_p2v, & + lapl%xdirps%stagder_v2p, lapl%ydirps%stagder_v2p, lapl%zdirps%stagder_v2p) + + end subroutine poissmult_dirz + + +end module m_base_poisson_cg diff --git a/src/solver.f90 b/src/solver.f90 index 1c6510554..8c5e7fab4 100644 --- a/src/solver.f90 +++ b/src/solver.f90 @@ -15,6 +15,8 @@ module m_solver use m_time_integrator, only: time_intg_t use m_vector_calculus, only: vector_calculus_t + use m_poisson_cg, only: poisson_cg_t + implicit none type :: solver_t @@ -54,6 +56,7 @@ module m_solver class(base_backend_t), pointer :: backend type(mesh_t), pointer :: mesh type(time_intg_t) :: time_integrator + type(poisson_cg_t) :: poisson_cg type(allocator_t), pointer :: host_allocator type(dirps_t), pointer :: xdirps, ydirps, zdirps type(vector_calculus_t) :: vector_calculus @@ -148,8 +151,11 @@ function init(backend, mesh, host_allocator) result(solver) solver%ydirps, solver%zdirps) solver%poisson => poisson_fft case ('CG') - if (solver%mesh%par%is_root()) & - print *, 'Poisson solver: CG, not yet implemented' + if (solver%mesh%par%is_root()) then + print *, 'Poisson solver: CG' + end if + + solver%poisson_cg = poisson_cg_t(solver%backend, mesh) solver%poisson => poisson_cg case default error stop 'poisson_solver_type is not valid. Use "FFT" or "CG".' @@ -382,6 +388,8 @@ subroutine poisson_cg(self, pressure, div_u) class(field_t), intent(inout) :: pressure class(field_t), intent(in) :: div_u + call self%poisson_cg%solve(pressure, div_u, self%backend) + end subroutine poisson_cg subroutine pressure_correction(self, u, v, w) diff --git a/src/vector_calculus.f90 b/src/vector_calculus.f90 index 1b70438f0..2008dcaa3 100644 --- a/src/vector_calculus.f90 +++ b/src/vector_calculus.f90 @@ -4,7 +4,8 @@ module m_vector_calculus use m_allocator, only: allocator_t use m_base_backend, only: base_backend_t use m_common, only: dp, DIR_X, DIR_Y, DIR_Z, & - RDR_X2Y, RDR_X2Z, RDR_Y2X, RDR_Y2Z, RDR_Z2X, RDR_Z2Y + RDR_X2Y, RDR_X2Z, RDR_Y2X, RDR_Y2Z, RDR_Z2X, RDR_Z2Y, & + CELL, X_FACE, Y_FACE, Z_FACE, VERT use m_field, only: field_t use m_tdsops, only: tdsops_t @@ -18,6 +19,8 @@ module m_vector_calculus procedure :: divergence_v2c procedure :: gradient_c2v procedure :: laplacian + procedure :: divgrad + procedure :: divgrad_stag end type vector_calculus_t interface vector_calculus_t @@ -392,4 +395,111 @@ subroutine laplacian(self, lapl_u, u, x_der2nd, y_der2nd, z_der2nd) end subroutine laplacian + subroutine divgrad(self, d2pdx2, p, & + x_stagder_c2v, x_interpl_c2v, & + y_stagder_c2v, y_interpl_c2v, & + z_stagder_c2v, z_interpl_c2v, & + x_stagder_v2c, x_interpl_v2c, & + y_stagder_v2c, y_interpl_v2c, & + z_stagder_v2c, z_interpl_v2c) + !! Computes the Laplacian of a scalar field 'p' using the div(grad) + !! formulation. The expected use case is evaluating the pressure Laplacian + !! consistently with div(u). + !! + !! Evaluated at the cell centres (data_loc=CELL). + !! + !! Input and output fields are in DIR_Z layout. + class(vector_calculus_t) :: self + class(field_t), intent(inout) :: d2pdx2 + class(field_t), intent(in) :: p + class(tdsops_t), intent(in) :: x_stagder_c2v, x_interpl_c2v, & + y_stagder_c2v, y_interpl_c2v, & + z_stagder_c2v, z_interpl_c2v + class(tdsops_t), intent(in) :: x_stagder_v2c, x_interpl_v2c, & + y_stagder_v2c, y_interpl_v2c, & + z_stagder_v2c, z_interpl_v2c + + class(field_t), pointer :: dpdx, dpdy, dpdz + + if (d2pdx2%data_loc /= CELL .or. p%data_loc /= CELL) then + error stop "p and d2pdx2 must both be cell-centred for div(grad)" + end if + + if (d2pdx2%dir /= DIR_Z .or. p%dir /= DIR_Z) then + error stop "p and d2pdx2 must both be Z-oriented for div(grad)" + end if + + dpdx => self%backend%allocator%get_block(DIR_X, VERT) + dpdy => self%backend%allocator%get_block(DIR_X, VERT) + dpdz => self%backend%allocator%get_block(DIR_X, VERT) + + call self%gradient_c2v(dpdx, dpdy, dpdz, p, & + x_stagder_c2v, x_interpl_c2v, & + y_stagder_c2v, y_interpl_c2v, & + z_stagder_c2v, z_interpl_c2v) + call self%divergence_v2c(d2pdx2, dpdx, dpdy, dpdz, & + x_stagder_v2c, x_interpl_v2c, & + y_stagder_v2c, y_interpl_v2c, & + z_stagder_v2c, z_interpl_v2c) + + call self%backend%allocator%release_block(dpdx) + call self%backend%allocator%release_block(dpdy) + call self%backend%allocator%release_block(dpdz) + + end subroutine divgrad + + subroutine divgrad_stag(self, d2pdx2_z, p, & + x_stagder_c2v, y_stagder_c2v, z_stagder_c2v, & + x_stagder_v2c, y_stagder_v2c, z_stagder_v2c) + + class(vector_calculus_t) :: self + class(field_t), intent(inout) :: d2pdx2_z + class(field_t), intent(in) :: p + class(tdsops_t), intent(in) :: x_stagder_c2v, y_stagder_c2v, z_stagder_c2v + class(tdsops_t), intent(in) :: x_stagder_v2c, y_stagder_v2c, z_stagder_v2c + + class(field_t), pointer :: p_x, p_y + class(field_t), pointer :: dpdx, dpdy, dpdz + class(field_t), pointer :: d2pdx2, d2pdy2, d2pdz2 + + ! Compute staggered (face) gradients + dpdx => self%backend%allocator%get_block(DIR_X, X_FACE) + dpdy => self%backend%allocator%get_block(DIR_Y, Y_FACE) + dpdz => self%backend%allocator%get_block(DIR_Z, Z_FACE) + p_x => self%backend%allocator%get_block(DIR_X, CELL) + p_y => self%backend%allocator%get_block(DIR_Y, CELL) + + call self%backend%reorder(p_x, p, RDR_Z2X) + call self%backend%reorder(p_y, p, RDR_Z2Y) + call self%backend%tds_solve(dpdx, p_x, x_stagder_c2v) + call self%backend%tds_solve(dpdy, p_y, y_stagder_c2v) + call self%backend%tds_solve(dpdz, p, z_stagder_c2v) + + call self%backend%allocator%release_block(p_x) + call self%backend%allocator%release_block(p_y) + + ! Compute staggered (cell) derder + d2pdx2 => self%backend%allocator%get_block(DIR_X, CELL) + d2pdy2 => self%backend%allocator%get_block(DIR_Y, CELL) + d2pdz2 => self%backend%allocator%get_block(DIR_Z, CELL) + call self%backend%tds_solve(d2pdx2, dpdx, x_stagder_v2c) + call self%backend%tds_solve(d2pdy2, dpdy, y_stagder_v2c) + call self%backend%tds_solve(d2pdz2, dpdz, z_stagder_v2c) + + call self%backend%allocator%release_block(dpdx) + call self%backend%allocator%release_block(dpdy) + call self%backend%allocator%release_block(dpdz) + + ! Accumulate Laplacian + call self%backend%sum_yintox(d2pdx2, d2pdy2) + call self%backend%sum_zintox(d2pdx2, d2pdz2) + + call self%backend%reorder(d2pdx2_z, d2pdx2, RDR_X2Z) + + call self%backend%allocator%release_block(d2pdx2) + call self%backend%allocator%release_block(d2pdy2) + call self%backend%allocator%release_block(d2pdz2) + + end subroutine divgrad_stag + end module m_vector_calculus diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 23539ac67..f9a72c084 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -25,6 +25,9 @@ function(define_test testfile np backend) endif() find_package(OpenMP REQUIRED) target_link_libraries(${test_name} PRIVATE OpenMP::OpenMP_Fortran) + if (${POISSON_SOLVER} STREQUAL "ITER") + target_compile_definitions(${test_name} PUBLIC HAVE_CG) + endif() endif() target_link_libraries(${test_name} PRIVATE x3d2) target_link_libraries(${test_name} PRIVATE x3d2_backends) @@ -82,3 +85,9 @@ if(${CMAKE_Fortran_COMPILER_ID} STREQUAL "PGI" OR define_test(cuda/test_cuda_transeq.f90 ${CMAKE_CTEST_NPROCS} cuda) endif() +define_test(test_poisson_cg_eval.f90 1 omp) +define_test(test_poisson_cg_eval.f90 ${CMAKE_CTEST_NPROCS} omp) +if (${POISSON_SOLVER} STREQUAL "ITER") + define_test(test_poisson_cg_solve.f90 1 omp) + define_test(test_poisson_cg_solve.f90 ${CMAKE_CTEST_NPROCS} omp) +endif() diff --git a/tests/test_mesh.f90 b/tests/test_mesh.f90 index 6d6e4030a..029bdd50f 100644 --- a/tests/test_mesh.f90 +++ b/tests/test_mesh.f90 @@ -107,8 +107,9 @@ subroutine run_test_mesh(use_2decomp, allpass) if (.not. (n_cell == n_vert_z(mesh%par%nrank + 1) & .and. n_vert == n_vert_z(mesh%par%nrank + 1))) then allpass = .false. - print *, mesh%par%nrank, "error in get_n, n_cell=", & - n_cell, "n_vert=", n_vert + print *, mesh%par%nrank, & + "error in get_n, n_cell=", n_cell, & + "n_vert=", n_vert end if end if diff --git a/tests/test_poisson_cg_eval.f90 b/tests/test_poisson_cg_eval.f90 new file mode 100644 index 000000000..87e87fc78 --- /dev/null +++ b/tests/test_poisson_cg_eval.f90 @@ -0,0 +1,369 @@ +!!! test_poisson_cg_eval.f90 +!! +!! SPDX-License-Identifier: BSD-3-Clause + +program test_poisson_cg_eval + !! Tests evaluating the Poisson/Laplace operator Lapl(p) on the pressure grid, + !! used by the iterative Poisson solver. + + use MPI + + use m_allocator, only: field_t + use m_base_backend, only: base_backend_t + use m_mesh, only: mesh_t +#ifdef CUDA +#else + use m_omp_backend +#endif + use m_common, only: DIR_X, DIR_Y, DIR_Z, CELL, POISSON_SOLVER_CG + use m_base_poisson_cg, only: laplace_operator_t, poisson_precon_t + use m_poisson_cg_backend, only: init_precon, poisson_precon_impl + + implicit none + + class(allocator_t), allocatable :: allocator + class(base_backend_t), allocatable :: backend + class(field_t), pointer :: pressure + class(field_t), pointer :: f + + integer :: irank, nproc + integer :: ierr + + logical :: test_pass + + real(dp), parameter :: Lx = 1.0_dp + real(dp), parameter :: Ly = 1.0_dp + real(dp), parameter :: Lz = 1.0_dp + integer, parameter :: nref = 4 + + real(dp), dimension(nref) :: e + + type(laplace_operator_t) :: lapl +#ifdef HAVE_CG + class(poisson_precon_t), allocatable :: precon +#endif + + call MPI_Init(ierr) + call MPI_Comm_rank(MPI_COMM_WORLD, irank, ierr) + call MPI_Comm_size(MPI_COMM_WORLD, nproc, ierr) + if (irank == 0) then + print *, "Testing the pressure Laplacian operator" + print *, "Parallel run with", nproc, "ranks" + end if + + test_pass = .true. + + call test_driver(lapl) +#ifdef HAVE_CG + call test_driver(precon) +#endif + + call MPI_Allreduce(MPI_IN_PLACE, test_pass, 1, & + MPI_LOGICAL, MPI_LAND, MPI_COMM_WORLD, & + ierr) + call MPI_Finalize(ierr) + + if (irank == 0) then + if (.not. test_pass) then + error stop "Test failed" + end if + end if + +contains + + subroutine test_driver(linear_operator) + + class(*), intent(in) :: linear_operator + + integer :: nx + integer :: ny + integer :: nz + integer :: i + + integer :: order + character(len=:), allocatable :: opname + + select type (linear_operator) + type is (laplace_operator_t) + order = 6 + opname = "High-Order Laplacian" + type is (poisson_precon_impl) + order = 2 + opname = "Low-Order preconditioner" + class default + error stop "Unknown linear operator" + end select + + nx = 16; ny = 16; nz = 16 + ! nx = nx * nproc + ny = ny*nproc + ! nz = nz * nproc + do i = 1, nref + if (irank == 0) then + print *, "---------------------------------" + print *, "Testing refinement level ", i - 1 + print *, "Using the ", opname + end if + e(i) = run_test([nx, ny, nz], linear_operator) + nx = 2*nx; ny = 2*ny; nz = 2*nz + end do + + do i = 2, nref + if (e(i) > 2.2*(e(i - 1)/(2**order))) then + if (irank == 0) then + print *, "Error convergence ", i, " failed ", & + e(i), e(i - 1)/(2**order) + end if + test_pass = .false. + else + if (irank == 0) then + print *, "Error convergence ", i, "satisfied ", & + e(i), e(i - 1)/(2**order) + end if + end if + end do + + end subroutine test_driver + + real(dp) function run_test(n, linear_operator) + + integer, dimension(3), intent(in) :: n + class(*), intent(in) :: linear_operator + + type(mesh_t) :: mesh + real(dp) :: rms_err + + call initialise_test(mesh, n) + + call test_constant_field(f, linear_operator, pressure, mesh) + rms_err = test_variable_field(f, linear_operator, pressure, mesh) + + ! Finalise test + call backend%allocator%release_block(pressure) + + run_test = rms_err + + end function run_test + + subroutine initialise_test(mesh, n) + + type(mesh_t), intent(out) :: mesh + integer, dimension(3), intent(in) :: n + + call init_globs(mesh, n, nproc, [Lx, Ly, Lz]) +#ifdef CUDA + if (irank == 0) then + error stop "CUDA iterative solver not currently supported" + end if +#else + if (allocated(allocator)) then + deallocate (allocator) + end if + allocate (allocator) + allocator = allocator_t(mesh, SZ) + + if (allocated(backend)) then + deallocate (backend) + end if + allocate (omp_backend_t :: backend) + backend = omp_backend_t(mesh, allocator) +#endif + lapl = laplace_operator_t(backend, mesh) +#ifdef HAVE_CG + precon = init_precon(backend) +#endif + + ! Main solver calls Poisson in the DIR_Z orientation + pressure => backend%allocator%get_block(DIR_Z, CELL) + f => backend%allocator%get_block(DIR_Z, CELL) + + call MPI_Barrier(MPI_COMM_WORLD, ierr) + if (irank == 0) then + print *, "Initialisation complete" + end if + + end subroutine initialise_test + + subroutine init_globs(mesh, n, nproc, L) + !! Initialisation for the globs object + + type(mesh_t), intent(out) :: mesh + integer, dimension(3), intent(in) :: n ! The grid sizes + integer, intent(in) :: nproc ! The number of processors + real(dp), dimension(3), intent(in) :: L ! The domain dimensions + + mesh = mesh_t(n, [1, nproc, 1], L, & + ["periodic", "periodic"], & + ["periodic", "periodic"], & + ["periodic", "periodic"]) + + end subroutine init_globs + + subroutine test_constant_field(f, linear_operator, pressure, mesh) + + class(field_t), intent(inout) :: f + class(*), intent(in) :: linear_operator + class(field_t), intent(in) :: pressure + type(mesh_t), intent(in) :: mesh + + real(dp), dimension(:, :, :), allocatable :: expect + + real(dp) :: rms_err + + logical :: check_pass + + if (irank == 0) then + print *, "Testing constant field" + end if + + ! Set pressure field to some constant + pressure%data = 0 !42 + allocate (expect, mold=f%data) + expect = 0 ! Correct answer + f%data = 17 ! Initialise with wrong answer + + select type (linear_operator) + type is (laplace_operator_t) + call linear_operator%apply(f, pressure) + type is (poisson_precon_impl) + call linear_operator%apply(pressure, f, backend) + class default + error stop "Unsupported linear operator type" + end select + + ! Check Laplacian evaluation (expect zero) + rms_err = check_soln(check_pass, mesh, f, expect) + + if (.not. check_pass) then + test_pass = .false. + + if (irank == 0) then + print *, "- FAILED RMS(err) = ", rms_err + end if + else + if (irank == 0) then + print *, "- PASS" + end if + end if + + end subroutine test_constant_field + + real(dp) function test_variable_field(f, linear_operator, pressure, mesh) + + use m_common, only: pi + use m_ordering, only: get_index_dir + + class(field_t), intent(inout) :: f + class(*), intent(in) :: linear_operator + class(field_t), intent(in) :: pressure + type(mesh_t), intent(in) :: mesh + + real(dp), dimension(:, :, :), allocatable :: expect + + integer, dimension(3) :: n + real(dp) :: x, y, z + real(dp), dimension(3) :: coords + integer :: i, j, k + integer :: ii, jj, kk + + real(dp) :: dx, dy, dz + real(dp) :: Lx, Ly, Lz + + logical :: check_pass + + if (irank == 0) then + print *, "Testing variable field" + end if + + dx = mesh%geo%d(1); dy = mesh%geo%d(2); dz = mesh%geo%d(3) + Lx = mesh%geo%L(1); Ly = mesh%geo%L(2); Lz = mesh%geo%L(3) + n = mesh%get_dims(CELL) + + ! Set pressure field to some variable + allocate (expect, mold=f%data) + do k = 1, n(3) + do j = 1, n(2) + do i = 1, n(1) + coords = mesh%get_coordinates(i, j, k) + x = coords(1); y = coords(2); z = coords(3) + + ! Need to get Cartesian -> memory layout mapping + call get_index_dir(ii, jj, kk, i, j, k, & + DIR_Z, & + SZ, n(1), n(2), n(3)) + + pressure%data(ii, jj, kk) = cos(2*pi*(x/Lx)) + & + cos(2*pi*(y/Ly)) + & + cos(2*pi*(z/Lz)) + expect(ii, jj, kk) = -((2*pi/Lx)**2*cos(2*pi*(x/Lx)) + & + (2*pi/Ly)**2*cos(2*pi*(y/Ly)) + & + (2*pi/Lz)**2*cos(2*pi*(z/Lz))) + ! pressure%data(ii, jj, kk) = cos(2 * pi * (z / Lz)) + ! expect(ii, jj, kk) = -((2 * pi / Lz)**2 * cos(2 * pi * (z / Lz))) + end do + end do + end do + f%data = 0 ! Initialise with wrong answer + + select type (linear_operator) + type is (laplace_operator_t) + call linear_operator%apply(f, pressure) + type is (poisson_precon_impl) + call linear_operator%apply(pressure, f, backend) + class default + error stop "Unsupported linear operator type" + end select + + ! Check Laplacian evaluation + ! XXX: Note had to relax the tolerance, otherwise obtains RMS(err)~=8e-7 + test_variable_field = check_soln(check_pass, mesh, f, expect, & + opttol=1.0e-6_dp) + + end function test_variable_field + + real(dp) function check_soln(check_pass, mesh, soln, expect, opttol) + + logical, intent(out) :: check_pass + type(mesh_t), intent(in) :: mesh + class(field_t), intent(in) :: soln + real(dp), dimension(:, :, :), intent(in) :: expect + real(dp), intent(in), optional :: opttol + + real(dp) :: rms + real(dp) :: tol + + integer, dimension(3) :: ng + integer :: n + + if (present(opttol)) then + tol = opttol + else + tol = 1.0e-8_dp + end if + + ng = mesh%get_global_dims(CELL) + n = product(ng) + + rms = sum((soln%data - expect)**2) + call MPI_Allreduce(MPI_IN_PLACE, rms, 1, & + MPI_DOUBLE_PRECISION, MPI_SUM, MPI_COMM_WORLD, & + ierr) + rms = sqrt(rms/n) + + if (rms /= rms) then ! NAN check + print *, "- SEVERE ERROR: RMS=NAN" + test_pass = .false. + check_pass = .false. + else + if (rms > tol) then + check_pass = .false. + else + check_pass = .true. + end if + end if + + check_soln = rms + + end function check_soln + +end program test_poisson_cg_eval diff --git a/tests/test_poisson_cg_solve.f90 b/tests/test_poisson_cg_solve.f90 new file mode 100644 index 000000000..8711ac2a1 --- /dev/null +++ b/tests/test_poisson_cg_solve.f90 @@ -0,0 +1,253 @@ +!!! test_poisson_cg_solve.f90 +!! +!! SPDX-License-Identifier: BSD-3-Clause + +program test_poisson_cg_solve + !! Tests solving the Poisson problem using CG. + + use mpi + use petsc + + use m_common, only: dp, DIR_Z, CELL, pi + use m_mesh, only: mesh_t + use m_base_backend, only: base_backend_t +#ifdef CUDA +#else + use m_omp_backend +#endif + use m_poisson_cg, only: poisson_cg_t + use m_ordering, only: get_index_dir + + implicit none + + integer :: irank, nproc + + integer, parameter :: nref = 3 ! Number of refinements to perform + real(dp), parameter :: Lx = 1.0_dp + real(dp), parameter :: Ly = 1.0_dp + real(dp), parameter :: Lz = 1.0_dp + + logical :: test_pass + + call setup() + call run() + call teardown() + +contains + + subroutine setup() + !! Perform test program setup + + integer :: ierr + + call MPI_Init(ierr) + call PETScInitialize(PETSC_NULL_CHARACTER, ierr) + call MPI_Comm_rank(MPI_COMM_WORLD, irank, ierr) + call MPI_Comm_size(MPI_COMM_WORLD, nproc, ierr) + + test_pass = .true. + + end subroutine setup + + subroutine run() + + integer :: nx, ny, nz + integer :: i + + type(mesh_t), allocatable :: mesh + class(allocator_t), allocatable :: allocator + class(base_backend_t), allocatable :: backend + class(poisson_cg_t), allocatable :: poisson_cg + class(field_t), pointer :: p, f + + real(dp), dimension(nref) :: rms + + ! Test across multiple refinement levels + nx = 16; ny = 16; nz = 16 + do i = 1, nref + mesh = mesh_t(dims_global=[nx, ny, nz], & + nproc_dir=[1, 1, nproc], & + L_global=[Lx, Ly, Lz], & + BC_x=["periodic", "periodic"], & + BC_y=["periodic", "periodic"], & + BC_z=["periodic", "periodic"], & + use_2decomp=.false.) +#ifdef CUDA + error stop "CUDA iterative solver not currently supported" +#else + allocator = allocator_t(mesh, SZ) + if (allocated(backend)) then + deallocate (backend) + end if + allocate (omp_backend_t :: backend) + backend = omp_backend_t(mesh, allocator) +#endif + poisson_cg = poisson_cg_t(backend, mesh) + + ! Main solver calls Poisson in the DIR_Z orientation + p => backend%allocator%get_block(DIR_Z, CELL) + f => backend%allocator%get_block(DIR_Z, CELL) + + call set_rhs(f, backend%mesh) + + call poisson_cg%solve(p, f, backend) + + rms(i) = calc_err(p, backend%mesh) + + call backend%allocator%release_block(p) + call backend%allocator%release_block(f) + + nx = 2*nx + ny = 2*ny + nz = 2*nz + end do + + call test_convergence(rms) + + end subroutine run + + subroutine set_rhs(f, mesh) + class(field_t), intent(inout) :: f + type(mesh_t), intent(in) :: mesh + + integer, dimension(3) :: n + real(dp), dimension(3) :: d, L + + integer :: i, j, k + integer :: ii, jj, kk + real(dp) :: x, y, z + + n = mesh%get_dims(CELL) + d = mesh%geo%d + L = mesh%geo%L + + do k = 1, n(3) + do j = 1, n(2) + do i = 1, n(1) + x = (mesh%par%n_offset(1) + (i - 1))*d(1) + y = (mesh%par%n_offset(2) + (j - 1))*d(2) + z = (mesh%par%n_offset(3) + (k - 1))*d(3) + + ! Need to get Cartesian -> memory layout mapping + call get_index_dir(ii, jj, kk, i, j, k, & + DIR_Z, & + SZ, n(1), n(2), n(3)) + + f%data(ii, jj, kk) = -((2*pi/L(1))**2*cos(2*pi*(x/L(1))) + & + (2*pi/L(2))**2*cos(2*pi*(y/L(2))) + & + (2*pi/L(3))**2*cos(2*pi*(z/L(3)))) + end do + end do + end do + + end subroutine set_rhs + + real(dp) function calc_err(p, mesh) + class(field_t), intent(in) :: p + type(mesh_t), intent(in) :: mesh + + real(dp) :: err_rms + real(dp) :: p_an + + integer, dimension(3) :: n + real(dp), dimension(3) :: d, L + + integer :: i, j, k + integer :: ii, jj, kk + real(dp) :: x, y, z + + integer :: ierr + + err_rms = 0.0_dp + + n = mesh%get_dims(CELL) + d = mesh%geo%d + L = mesh%geo%L + + do k = 1, n(3) + do j = 1, n(2) + do i = 1, n(1) + x = (mesh%par%n_offset(1) + (i - 1))*d(1) + y = (mesh%par%n_offset(2) + (j - 1))*d(2) + z = (mesh%par%n_offset(3) + (k - 1))*d(3) + + ! Need to get Cartesian -> memory layout mapping + call get_index_dir(ii, jj, kk, i, j, k, & + DIR_Z, & + SZ, n(1), n(2), n(3)) + + p_an = cos(2*pi*(x/L(1))) + & + cos(2*pi*(y/L(2))) + & + cos(2*pi*(z/L(3))) + + err_rms = err_rms + (p%data(ii, jj, kk) - p_an)**2 + end do + end do + end do + + call MPI_Allreduce(MPI_IN_PLACE, err_rms, 1, & + MPI_DOUBLE_PRECISION, MPI_SUM, MPI_COMM_WORLD, & + ierr) + n = mesh%get_global_dims(CELL) + calc_err = sqrt(err_rms/product(n)) + + end function calc_err + + subroutine test_convergence(rms_err) + real(dp), dimension(:), intent(in) :: rms_err + + integer :: i + real(dp) :: r + real(dp) :: tol + + do i = 1, size(rms_err) + if (irank == 0) then + print *, "RMS = ", rms_err(i) + end if + end do + do i = 2, size(rms_err) + r = rms_err(i)/rms_err(i - 1) + if (irank == 0) then + print *, "Convergence ratio ", i, ": ", r + end if + + tol = 1.1_dp/(2**6) + if (r > tol) then + if (irank == 0) then + print *, "- Error convergence failed, tolerance = ", tol + end if + + test_pass = .false. + end if + end do + + end subroutine test_convergence + + subroutine teardown() + !! Perform test program cleanup + use petsc + + integer :: ierr + + ! Reduce test status + call MPI_Allreduce(MPI_IN_PLACE, test_pass, 1, & + MPI_LOGICAL, MPI_LAND, MPI_COMM_WORLD, & + ierr) + + ! Cleanup parallelism + call PetscFinalize(ierr) + call MPI_Finalize(ierr) + + ! Report pass/fail + if (irank == 0) then + if (test_pass) then + print *, "PASS" + else + print *, "FAIL" + stop 1 + end if + end if + + end subroutine teardown + +end program test_poisson_cg_solve diff --git a/tests/test_setget_field.f90 b/tests/test_setget_field.f90 index 89ebd2365..c954fccb7 100644 --- a/tests/test_setget_field.f90 +++ b/tests/test_setget_field.f90 @@ -1,6 +1,6 @@ program test_setget_field - use mpi + use mpi use m_allocator, only: allocator_t, field_t use m_base_backend, only: base_backend_t diff --git a/tests/test_sum_yintox.f90 b/tests/test_sum_yintox.f90 new file mode 100644 index 000000000..7a31086e0 --- /dev/null +++ b/tests/test_sum_yintox.f90 @@ -0,0 +1,103 @@ +program test_sum_yintox + !! Tests the implementation of summing a Y-oriented field into an X-oriented + !! one. + + use m_common, only: DIR_X, DIR_Y + + use m_allocator + use m_base_backend +#ifdef CUDA +#else + use m_omp_backend + use m_omp_common, only: SZ +#endif + + implicit none + + integer, parameter :: nx = 17, ny = 32, nz = 59 + real(dp), parameter :: lx = 1.618, ly = 3.141529, lz = 1.729 + + class(base_backend_t), pointer :: backend + class(allocator_t), pointer :: allocator +#ifdef CUDA +#else + type(omp_backend_t), target :: omp_backend + type(allocator_t), target :: omp_allocator +#endif + + type(mesh_t) :: mesh + + integer :: nrank, nproc + integer :: ierr + + logical :: test_pass = .true. + + call MPI_Init(ierr) + call MPI_Comm_rank(MPI_COMM_WORLD, nrank, ierr) + call MPI_Comm_size(MPI_COMM_WORLD, nproc, ierr) + + mesh = mesh_t([nx, ny, nz], & + [1, 1, nproc], & + [lx, ly, lz]) + +#ifdef CUDA +#else + omp_allocator = allocator_t(mesh, SZ) + allocator => omp_allocator + omp_backend = omp_backend_t(mesh, allocator) + backend => omp_backend +#endif + + call runtest() + + if (nrank == 0) then + if (.not. test_pass) then + error stop "Test failed" + end if + end if + + call MPI_Finalize(ierr) + +contains + + subroutine runtest() + + use m_ordering, only: get_index_dir + + class(field_t), pointer :: a, b + integer :: ctr + integer :: i, j, k + integer :: ii, jj, kk + + integer, dimension(3) :: dims + + a => backend%allocator%get_block(DIR_X) + b => backend%allocator%get_block(DIR_Y) + + dims = mesh%get_padded_dims(DIR_C) + + ! Initialise fields so that b = -a + ctr = 0 + do k = 1, dims(3) + do j = 1, dims(2) + do i = 1, dims(1) + call get_index_dir(ii, jj, kk, i, j, k, DIR_X, SZ, & + dims(1), dims(2), dims(3)) + a%data(ii, jj, kk) = ctr + call get_index_dir(ii, jj, kk, i, j, k, DIR_Y, SZ, & + dims(1), dims(2), dims(3)) + b%data(ii, jj, kk) = -ctr + ctr = ctr + 1 + end do + end do + end do + + call backend%sum_yintox(a, b) + + if ((minval(a%data) /= 0) .or. (maxval(a%data) /= 0)) then + test_pass = .false. + end if + + end subroutine runtest + +end program test_sum_yintox diff --git a/tests/test_sum_zintox.f90 b/tests/test_sum_zintox.f90 new file mode 100644 index 000000000..0325ff3d6 --- /dev/null +++ b/tests/test_sum_zintox.f90 @@ -0,0 +1,103 @@ +program test_sum_zintox + !! Tests the implementation of summing a Z-oriented field into an X-oriented + !! one. + + use m_common, only: DIR_X, DIR_Z + + use m_allocator + use m_base_backend +#ifdef CUDA +#else + use m_omp_backend + use m_omp_common, only: SZ +#endif + + implicit none + + integer, parameter :: nx = 17, ny = 32, nz = 59 + real(dp), parameter :: lx = 1.618, ly = 3.141529, lz = 1.729 + + class(base_backend_t), pointer :: backend + class(allocator_t), pointer :: allocator +#ifdef CUDA +#else + type(omp_backend_t), target :: omp_backend + type(allocator_t), target :: omp_allocator +#endif + + type(mesh_t) :: mesh + + integer :: nrank, nproc + integer :: ierr + + logical :: test_pass = .true. + + call MPI_Init(ierr) + call MPI_Comm_rank(MPI_COMM_WORLD, nrank, ierr) + call MPI_Comm_size(MPI_COMM_WORLD, nproc, ierr) + + mesh = mesh_t([nx, ny, nz], & + [1, 1, nproc], & + [lx, ly, lz]) + +#ifdef CUDA +#else + omp_allocator = allocator_t(mesh, SZ) + allocator => omp_allocator + omp_backend = omp_backend_t(mesh, allocator) + backend => omp_backend +#endif + + call runtest() + + if (nrank == 0) then + if (.not. test_pass) then + error stop "Test failed" + end if + end if + + call MPI_Finalize(ierr) + +contains + + subroutine runtest() + + use m_ordering, only: get_index_dir + + class(field_t), pointer :: a, b + integer :: ctr + integer :: i, j, k + integer :: ii, jj, kk + + integer, dimension(3) :: dims + + a => backend%allocator%get_block(DIR_X) + b => backend%allocator%get_block(DIR_Z) + + dims = mesh%get_padded_dims(DIR_C) + + ! Initialise fields so that b = -a + ctr = 0 + do k = 1, dims(3) + do j = 1, dims(2) + do i = 1, dims(1) + call get_index_dir(ii, jj, kk, i, j, k, DIR_X, SZ, & + dims(1), dims(2), dims(3)) + a%data(ii, jj, kk) = ctr + call get_index_dir(ii, jj, kk, i, j, k, DIR_Z, SZ, & + dims(1), dims(2), dims(3)) + b%data(ii, jj, kk) = -ctr + ctr = ctr + 1 + end do + end do + end do + + call backend%sum_zintox(a, b) + + if ((minval(a%data) /= 0) .or. (maxval(a%data) /= 0)) then + test_pass = .false. + end if + + end subroutine runtest + +end program test_sum_zintox