test_poisson_fft_batch.F90

Go to the documentation of this file.

!! Copyright (C) 2026 A. Buccheri, H. Menke
!!
!! This program is free software; you can redistribute it and/or modify
!! it under the terms of the GNU General Public License as published by
!! the Free Software Foundation; either version 2, or (at your option)
!! any later version.
!!
!! This program is distributed in the hope that it will be useful,
!! but WITHOUT ANY WARRANTY; without even the implied warranty of
!! MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
!! GNU General Public License for more details.
!!
!! You should have received a copy of the GNU General Public License
!! along with this program; if not, write to the Free Software
!! Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
!! 02110-1301, USA.
 
#include "global.h"
 
module test_poisson_fft_batch_oct_m
  use, intrinsic :: iso_fortran_env
 
  use calc_mode_par_oct_m, only: calc_mode_par, p_strategy_states
  use debug_oct_m
  use electrons_oct_m,     only: electrons_t
  use fft_oct_m,           only: fft_batch_packed, fft_batch_unpacked
  use fourier_space_oct_m, only: fourier_space_op_t, fourier_space_op_end
  use global_oct_m
  use kpoints_oct_m,       only: kpoints_point_is_gamma
  use math_oct_m
  use messages_oct_m
  use mesh_oct_m,          only: mesh_t
  use mpi_oct_m,           only: mpi_world
  use namespace_oct_m,     only: namespace_t
  use poisson_oct_m
  use profiling_oct_m
  use states_elec_oct_m
  use xc_cam_oct_m,        only: cam_null
 
  implicit none
  private
  public :: test_poisson_fft_batch
 
contains
 
  subroutine mock_plane_wave_state(mesh, g_vector, state)
    class(mesh_t),   intent(in)  :: mesh
    real(real64),    intent(in)  :: g_vector(:)
    complex(real64), intent(out) :: state(:)
 
    complex(real64), parameter :: two_pi_i = m_two * m_pi * m_zi
    integer :: ip
 
    assert(size(g_vector) == size(mesh%chi, 1))
    assert(size(state) == mesh%np)
 
    !$omp parallel do default(shared)
    do ip = 1, mesh%np
      state(ip) = exp(two_pi_i * dot_product(g_vector, mesh%chi(:, ip)))
    enddo
    !$omp end parallel do
 
  end subroutine mock_plane_wave_state
 
 
  subroutine test_poisson_fft_batch(namespace)
    type(namespace_t), intent(in) :: namespace
 
    integer,      parameter :: batch_size = 4
    integer,      parameter :: count_partial = 3
#ifndef NDEBUG
    real(real64), parameter :: tolerance = 1.0e-10_real64
#endif
    real(real64), parameter :: reinit_q_value = 1e9_real64 
 
    type(electrons_t),   pointer :: sys
    type(poisson_t)              :: psolver_single, psolver_batch, psolver_batch_last
    type(poisson_t)              :: psolver_real_batch, psolver_real_single
    type(fourier_space_op_t)     :: coulb
    complex(real64), allocatable :: rho_batch(:, :), pot(:), pot_batch(:, :), pot_partial(:, :), pot_batch_last(:, :)
    complex(real64), allocatable :: pot_real_batch(:, :), pot_real(:)
    real(real64),    allocatable :: diff(:), max_abs_diff(:), mean_diff(:), max_abs_diff_partial(:), max_abs_diff_last(:)
    real(real64),    allocatable :: max_abs_diff_real(:)
    integer                      :: ipair
    real(real64)                 :: qred(3)
 
    real(real64), parameter :: dG(3, batch_size) = reshape([ &
      1,  0,  0, &
      0,  1,  0, &
      1,  1,  0, &
      -1,  2,  1  &
      ], [3, batch_size])
 
    push_sub(test_poisson_fft_batch)
 
    call messages_write('Info: Testing batched FFT Poisson solve for periodic k-point pair products')
    call messages_info(namespace=namespace)
 
    ! Non-Gamma q-point in reduced coordinates.
    qred = [0.25_real64, 0.125_real64, 0.0_real64]
 
    ! Initialise the system; only the grid (sys%gr) is used below, as the codensities are mocked
    call calc_mode_par%set_parallelization(p_strategy_states, default=.false.)
    sys => electrons_t(namespace, mpi_world, calc_mode_id=1)
 
    ! Mocked densities
    safe_allocate(rho_batch(1:sys%gr%np, 1:batch_size))
    do ipair = 1, batch_size
      call mock_plane_wave_state(sys%gr, dg(:, ipair), rho_batch(:, ipair))
    enddo
 
    ! -----------------------
    ! Batch Poisson solve
    ! -----------------------
    call poisson_init(psolver_batch, namespace, sys%space, sys%gr%der, sys%mc, sys%gr%stencil, &
      sys%st%qtot, label='Pair-product Poisson batch size 4', solver=poisson_fft, verbose=.false., &
      force_serial=.true., force_cmplx=.true., fft_batch_size=batch_size)
 
    ! Build the kernel
    ! Initialisation done internally, triggered by change in q
    safe_allocate(coulb%qq(1:sys%space%dim))
    coulb%qq = reinit_q_value
    call poisson_build_kernel(psolver_batch, namespace, sys%space, coulb, qred, cam_null)
 
    safe_allocate(pot_batch(1:sys%gr%np, 1:batch_size))
    call zpoisson_solve_batch(psolver_batch, pot_batch, rho_batch, view=fft_batch_unpacked, kernel=coulb)
 
    ! Partial batch: solve only count_partial < batch_size columns with the same (capacity batch_size) solver.
    ! This exercises the active-count path: the unused capacity lanes must not affect the result.
    safe_allocate(pot_partial(1:sys%gr%np, 1:count_partial))
    call zpoisson_solve_batch(psolver_batch, pot_partial, rho_batch(:, 1:count_partial), &
      view=fft_batch_unpacked, kernel=coulb)
    call poisson_end(psolver_batch)
 
    ! -----------------------
    ! Batch-last (batch_axis = 4) solve: exercises the cube batch-last layout (CPU FFTW only).
    ! The same per-G Coulomb kernel applies (it is independent of the batch layout).
    ! -----------------------
    call poisson_init(psolver_batch_last, namespace, sys%space, sys%gr%der, sys%mc, sys%gr%stencil, &
      sys%st%qtot, label='Pair-product Poisson batch size 4 (batch-last)', solver=poisson_fft, verbose=.false., &
      force_serial=.true., force_cmplx=.true., fft_batch_size=batch_size, fft_batch_axis=4)
    safe_allocate(pot_batch_last(1:sys%gr%np, 1:batch_size))
    call zpoisson_solve_batch(psolver_batch_last, pot_batch_last, rho_batch, view=fft_batch_unpacked, kernel=coulb)
    call poisson_end(psolver_batch_last)
 
    ! -----------------------
    ! Complex data on a REAL cube (force_cmplx = .false.): the batched solver cannot perform a
    ! complex-to-complex FFT on a real cube, so it must fall back to transforming the real and
    ! imaginary parts separately with the real FFT, mirroring the scalar zpoisson_solve fallback.
    ! This is the layout the time-dependent exact-exchange / OEP-KLI path hits: the cube is sized
    ! for the real ground state, while the TD codensities are complex. The real/imag split requires
    ! a Gamma (q = 0) kernel, so no external q-shifted kernel is used here; the batched solve is
    ! checked against the per-column scalar solve on the same real cube.
    ! -----------------------
    call poisson_init(psolver_real_batch, namespace, sys%space, sys%gr%der, sys%mc, sys%gr%stencil, &
      sys%st%qtot, label='Pair-product Poisson real cube batch size 4', solver=poisson_fft, &
      verbose=.false., force_serial=.true., force_cmplx=.false., fft_batch_size=batch_size)
    safe_allocate(pot_real_batch(1:sys%gr%np, 1:batch_size))
    call zpoisson_solve_batch(psolver_real_batch, pot_real_batch, rho_batch, view=fft_batch_unpacked)
    call poisson_end(psolver_real_batch)
 
    write(message(1), '(a)') 'Test Poisson FFT batching: Performed batch solve'
    call messages_info(1, namespace=namespace)
 
    ! -----------------------
    ! Single Poisson solves
    ! -----------------------
    call poisson_init(psolver_single, namespace, sys%space, sys%gr%der, sys%mc, sys%gr%stencil, &
      sys%st%qtot, label='Pair-product Poisson batch size 1', solver=poisson_fft, verbose=.false., &
      force_serial=.true., force_cmplx=.true., fft_batch_size=1)
 
    ! Single-solve reference on a real cube, to validate the complex-on-real-cube batched solve above.
    call poisson_init(psolver_real_single, namespace, sys%space, sys%gr%der, sys%mc, sys%gr%stencil, &
      sys%st%qtot, label='Pair-product Poisson real cube batch size 1', solver=poisson_fft, &
      verbose=.false., force_serial=.true., force_cmplx=.false., fft_batch_size=1)
 
    safe_allocate(pot(1:sys%gr%np))
    safe_allocate(pot_real(1:sys%gr%np))
    safe_allocate(diff(1:sys%gr%np))
    safe_allocate(mean_diff(1:batch_size))
    safe_allocate(max_abs_diff(1:batch_size))
    safe_allocate(max_abs_diff_partial(1:count_partial))
    safe_allocate(max_abs_diff_last(1:batch_size))
    safe_allocate(max_abs_diff_real(1:batch_size))
 
    do ipair = 1, batch_size
      ! Reuse the same kernel
      call zpoisson_solve(psolver_single, namespace, pot, rho_batch(:, ipair), all_nodes=.false., kernel=coulb)
      diff = abs(pot_batch(:, ipair) - pot(:))
      mean_diff(ipair) = sum(diff)/real(size(diff), real64)
      max_abs_diff(ipair) = maxval(diff)
      ! The batch-last solve must also match the single solve, lane by lane
      max_abs_diff_last(ipair) = maxval(abs(pot_batch_last(:, ipair) - pot(:)))
      ! The first count_partial columns must also match the partial (count < capacity) batch solve
      if (ipair <= count_partial) then
        max_abs_diff_partial(ipair) = maxval(abs(pot_partial(:, ipair) - pot(:)))
      end if
      ! The complex-on-real-cube batched solve must match the per-column scalar solve on the real
      ! cube (both go through the real/imag split). No external kernel here (Gamma point).
      call zpoisson_solve(psolver_real_single, namespace, pot_real, rho_batch(:, ipair), all_nodes=.false.)
      max_abs_diff_real(ipair) = maxval(abs(pot_real_batch(:, ipair) - pot_real(:)))
    end do
 
    call poisson_end(psolver_single)
    call poisson_end(psolver_real_single)
    call fourier_space_op_end(coulb)
 
    safe_deallocate_a(rho_batch)
    safe_deallocate_a(pot_batch)
    safe_deallocate_a(pot_batch_last)
    safe_deallocate_a(pot_real_batch)
    safe_deallocate_a(pot_partial)
    safe_deallocate_a(pot)
    safe_deallocate_a(pot_real)
    safe_deallocate_a(diff)
 
    write(message(1), '(a)') 'Test Poisson FFT batching: Performed a set of single solves'
    call messages_info(1, namespace=namespace)
 
    ! -----------------------
    ! Error reporting
    ! -----------------------
    do ipair = 1, batch_size
      write(message(1), '(a,i0,a,es18.10,a,es18.10)') 'Poisson FFT pair ', ipair, &
        ': mean abs diff = ', mean_diff(ipair), ', max abs diff = ', max_abs_diff(ipair)
      call messages_info(1, namespace=namespace)
    end do
    write(message(1), '(a,es18.10)') 'Poisson FFT batch overall mean abs diff = ', sum(mean_diff)/real(batch_size, real64)
    write(message(2), '(a,es18.10)') 'Poisson FFT batch overall max abs diff = ', maxval(max_abs_diff)
    call messages_info(2, namespace=namespace)
 
    assert(all(max_abs_diff < tolerance))
 
    write(message(1), '(a,es18.10)') 'Poisson FFT batch-last (batch_axis=4) overall max abs diff = ', &
      maxval(max_abs_diff_last)
    call messages_info(1, namespace=namespace)
 
    assert(all(max_abs_diff_last < tolerance))
 
    write(message(1), '(a,i0,a,i0,a,es18.10)') 'Poisson FFT partial batch (count ', count_partial, &
      ' of capacity ', batch_size, ') overall max abs diff = ', maxval(max_abs_diff_partial)
    call messages_info(1, namespace=namespace)
 
    assert(all(max_abs_diff_partial < tolerance))
 
    write(message(1), '(a,es18.10)') 'Poisson FFT complex-on-real-cube batch overall max abs diff = ', &
      maxval(max_abs_diff_real)
    call messages_info(1, namespace=namespace)
 
    assert(all(max_abs_diff_real < tolerance))
 
    call messages_write('Poisson FFT batch application test passed')
    call messages_info(namespace=namespace)
 
    safe_deallocate_a(mean_diff)
    safe_deallocate_a(max_abs_diff)
    safe_deallocate_a(max_abs_diff_partial)
    safe_deallocate_a(max_abs_diff_last)
    safe_deallocate_a(max_abs_diff_real)
 
    call states_elec_deallocate_wfns(sys%st)
    safe_deallocate_p(sys)
 
    pop_sub(test_poisson_fft_batch)
 
  end subroutine test_poisson_fft_batch
 
end module test_poisson_fft_batch_oct_m
 
!! Local Variables:
!! mode: f90
!! coding: utf-8
!! End: