tree_amg.f90

Go to the documentation of this file.

! Copyright (c) 2024-2026, The Neko Authors
! All rights reserved.
!
! Redistribution and use in source and binary forms, with or without
! modification, are permitted provided that the following conditions
! are met:
!
!   * Redistributions of source code must retain the above copyright
!     notice, this list of conditions and the following disclaimer.
!
!   * Redistributions in binary form must reproduce the above
!     copyright notice, this list of conditions and the following
!     disclaimer in the documentation and/or other materials provided
!     with the distribution.
!
!   * Neither the name of the authors nor the names of its
!     contributors may be used to endorse or promote products derived
!     from this software without specific prior written permission.
!
! THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
! "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
! LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
! FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
! COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
! INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
! BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
! LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
! CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
! LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
! ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
! POSSIBILITY OF SUCH DAMAGE.
!
module tree_amg
  use num_types, only : rp
  use utils, only : neko_error
  use math, only : rzero, col2
  use device_math , only : device_rzero, device_col2, device_masked_atomic_reduction_0, &
       device_masked_gather_copy_0, device_cfill
  use coefs, only : coef_t
  use mesh, only : mesh_t
  use space, only : space_t
  use ax_product, only: ax_t
  use bc_list, only: bc_list_t
  use gather_scatter, only : gs_t, gs_op_add
  use device, only: device_map, device_unmap, device_sync, &
       device_stream_wait_event, glb_cmd_queue, glb_cmd_event
  use neko_config, only: neko_bcknd_device
  use, intrinsic :: iso_c_binding
  implicit none
  private
 
  type, private :: tamg_node_t
     logical :: isleaf = .true. 
     integer :: gid = -1
     integer :: lvl = -1
     integer :: ndofs = 0
     integer, allocatable :: dofs(:)
     real(kind=rp) :: xyz(3) 
     real(kind=rp), allocatable :: interp_r(:) 
     real(kind=rp), allocatable :: interp_p(:) 
   contains
     procedure, pass(this) :: free => node_free
  end type tamg_node_t
 
  type, private :: tamg_lvl_t
     integer :: lvl = -1
     integer :: nnodes = 0
     type(tamg_node_t), allocatable :: nodes(:)
     integer :: fine_lvl_dofs = 0
     real(kind=rp), allocatable :: wrk_in(:) 
     type(c_ptr) :: wrk_in_d = c_null_ptr
     real(kind=rp), allocatable :: wrk_out(:) 
     type(c_ptr) :: wrk_out_d = c_null_ptr
     integer, allocatable :: map_finest2lvl(:)
     type(c_ptr) :: map_finest2lvl_d = c_null_ptr
     !--!
     integer, allocatable :: map_f2c(:)
     type(c_ptr) :: map_f2c_d = c_null_ptr
   contains
     procedure, pass(this) :: free => lvl_free
  end type tamg_lvl_t
 
  type, public :: tamg_hierarchy_t
     integer :: nlvls
     type(tamg_lvl_t), allocatable :: lvl(:)
 
     class(ax_t), pointer :: ax
     type(mesh_t), pointer :: msh
     type(space_t), pointer :: xh
     type(coef_t), pointer :: coef
     type(gs_t), pointer :: gs_h
     type(bc_list_t), pointer :: blst
 
   contains
     procedure, pass(this) :: init => tamg_init
     procedure, pass(this) :: free => tamg_free
     procedure, pass(this) :: matvec => tamg_matvec
     procedure, pass(this) :: matvec_impl => tamg_matvec_impl
     procedure, pass(this) :: interp_f2c => tamg_restriction_operator
     procedure, pass(this) :: interp_c2f => tamg_prolongation_operator
     procedure, pass(this) :: interp_f2c_d => tamg_device_restriction_operator
     procedure, pass(this) :: interp_c2f_d => tamg_device_prolongation_operator
     procedure, pass(this) :: device_matvec => tamg_device_matvec_flat_impl
  end type tamg_hierarchy_t
 
  public tamg_lvl_init, tamg_node_init
 
contains
 
  subroutine tamg_init(this, ax, Xh, coef, msh, gs_h, nlvls, blst)
    class(tamg_hierarchy_t), target, intent(inout) :: this
    class(ax_t), target, intent(in) :: ax
    type(space_t),target, intent(in) :: Xh
    type(coef_t), target, intent(in) :: coef
    type(mesh_t), target, intent(in) :: msh
    type(gs_t), target, intent(in) :: gs_h
    integer, intent(in) :: nlvls
    type(bc_list_t), target, intent(in) :: blst
    integer :: i, n
 
    this%ax => ax
    this%msh => msh
    this%Xh => xh
    this%coef => coef
    this%gs_h => gs_h
    this%blst => blst
 
    if (nlvls .lt. 2) then
       call neko_error("Need to request at least two multigrid levels.")
    end if
 
    this%nlvls = nlvls
    allocate( this%lvl(this%nlvls) )
 
    do i = 1, nlvls
       allocate( this%lvl(i)%map_finest2lvl( 0:coef%dof%size() ))
       if (neko_bcknd_device .eq. 1) then
          call device_map(this%lvl(i)%map_finest2lvl, this%lvl(i)%map_finest2lvl_d, coef%dof%size()+1)
       end if
    end do
 
  end subroutine tamg_init
 
  subroutine tamg_free(this)
    class(tamg_hierarchy_t), intent(inout) :: this
    integer :: i
    if (allocated(this%lvl)) then
       ! using size() instead of this%nlvls since
       ! this%nlvls may be less than the allocated number of levels due to early
       ! termination of aggregation
       do i = 1, size(this%lvl)
          call this%lvl(i)%free()
       end do
       deallocate(this%lvl)
    end if
    nullify(this%ax)
    nullify(this%msh)
    nullify(this%Xh)
    nullify(this%coef)
    nullify(this%gs_h)
    nullify(this%blst)
  end subroutine tamg_free
 
  subroutine tamg_lvl_init(tamg_lvl, lvl, nnodes, ndofs)
    type(tamg_lvl_t), intent(inout) :: tamg_lvl
    integer, intent(in) :: lvl
    integer, intent(in) :: nnodes
    integer, intent(in) :: ndofs
 
    tamg_lvl%lvl = lvl
    tamg_lvl%nnodes = nnodes
    allocate( tamg_lvl%nodes(tamg_lvl%nnodes) )
    allocate( tamg_lvl%map_f2c(0:ndofs) )
    if (neko_bcknd_device .eq. 1) then
       call device_map(tamg_lvl%map_f2c, tamg_lvl%map_f2c_d, ndofs+1)
    end if
 
    tamg_lvl%fine_lvl_dofs = ndofs
    allocate( tamg_lvl%wrk_in( ndofs ) )
    allocate( tamg_lvl%wrk_out( ndofs ) )
    if (neko_bcknd_device .eq. 1) then
       call device_map(tamg_lvl%wrk_in, tamg_lvl%wrk_in_d, ndofs)
       call device_cfill(tamg_lvl%wrk_in_d, 0.0_rp, ndofs)
       call device_map(tamg_lvl%wrk_out, tamg_lvl%wrk_out_d, ndofs)
       call device_cfill(tamg_lvl%wrk_out_d, 0.0_rp, ndofs)
       ! Order the async fills against host writes; on unified memory
       ! the device pointers may alias the work arrays
       call device_sync()
    end if
  end subroutine tamg_lvl_init
 
  subroutine lvl_free(this)
    class(tamg_lvl_t), intent(inout) :: this
    integer :: i
 
    if (allocated(this%nodes)) then
       do i = 1, this%nnodes
          call this%nodes(i)%free()
       end do
       deallocate(this%nodes)
    end if
    if (allocated(this%wrk_in)) then
       if (neko_bcknd_device .eq. 1 .and. c_associated(this%wrk_in_d)) then
          call device_unmap(this%wrk_in, this%wrk_in_d)
       end if
       deallocate(this%wrk_in)
    end if
    if (allocated(this%wrk_out)) then
       if (neko_bcknd_device .eq. 1 .and. c_associated(this%wrk_out_d)) then
          call device_unmap(this%wrk_out, this%wrk_out_d)
       end if
       deallocate(this%wrk_out)
    end if
    if (allocated(this%map_f2c)) then
       if (neko_bcknd_device .eq. 1 .and. c_associated(this%map_f2c_d)) then
          call device_unmap(this%map_f2c, this%map_f2c_d)
       end if
       deallocate(this%map_f2c)
    end if
    if (allocated(this%map_finest2lvl)) then
       if (neko_bcknd_device .eq. 1 .and. c_associated(this%map_finest2lvl_d)) then
          call device_unmap(this%map_finest2lvl, this%map_finest2lvl_d)
       end if
       deallocate(this%map_finest2lvl)
    end if
    this%nnodes = 0
    this%lvl = -1
    this%fine_lvl_dofs = 0
  end subroutine lvl_free
 
  subroutine tamg_node_init(node, gid, ndofs)
    type(tamg_node_t), intent(inout) :: node
    integer, intent(in) :: gid
    integer, intent(in) :: ndofs
 
    node%gid = gid
    node%ndofs = ndofs
    allocate( node%dofs( node%ndofs) )
    node%dofs = -1
    allocate( node%interp_r( node%ndofs) )
    allocate( node%interp_p( node%ndofs) )
    node%interp_r = 1.0_rp
    node%interp_p = 1.0_rp
  end subroutine tamg_node_init
 
  subroutine node_free(this)
    class(tamg_node_t), intent(inout) :: this
 
    if (allocated(this%dofs)) then
       deallocate(this%dofs)
    end if
    if (allocated(this%interp_r)) then
       deallocate(this%interp_r)
    end if
    if (allocated(this%interp_p)) then
       deallocate(this%interp_p)
    end if
  end subroutine node_free
 
  recursive subroutine tamg_matvec(this, vec_out, vec_in, lvl_out)
    class(tamg_hierarchy_t), intent(inout) :: this
    real(kind=rp), intent(inout) :: vec_out(:)
    real(kind=rp), intent(inout) :: vec_in(:)
    integer, intent(in) :: lvl_out
    integer :: i, n, e
    !call this%matvec_impl(vec_out, vec_in, this%nlvls, lvl_out)
    call tamg_matvec_flat_impl(this, vec_out, vec_in, this%nlvls, lvl_out)
  end subroutine tamg_matvec
 
  recursive subroutine tamg_matvec_impl(this, vec_out, vec_in, lvl, lvl_out)
    class(tamg_hierarchy_t), intent(inout) :: this
    real(kind=rp), intent(inout) :: vec_out(:)
    real(kind=rp), intent(inout) :: vec_in(:)
    integer, intent(in) :: lvl
    integer, intent(in) :: lvl_out
    integer :: i, n, e
 
    if (lvl .eq. 0) then
       n = size(vec_in)
       call this%gs_h%op(vec_in, n, gs_op_add)
       call col2( vec_in, this%coef%mult, n)
 
       call this%ax%compute(vec_out, vec_in, this%coef, this%msh, this%Xh)
       call this%gs_h%op(vec_out, n, gs_op_add)
       call this%blst%apply(vec_out, n)
 
       if (lvl_out .ne. 0) then
          call col2(vec_out, this%coef%mult, n)
       end if
    else
       if (lvl_out .ge. lvl) then
          associate( wrk_in => this%lvl(lvl)%wrk_in, wrk_out => this%lvl(lvl)%wrk_out)
            n = this%lvl(lvl)%fine_lvl_dofs
            call rzero(wrk_in, n)
            call rzero(vec_out, this%lvl(lvl)%nnodes)
            do n = 1, this%lvl(lvl)%nnodes
               associate(node => this%lvl(lvl)%nodes(n))
                 do i = 1, node%ndofs
                    wrk_in( node%dofs(i) ) = wrk_in( node%dofs(i) ) + vec_in( node%gid ) * node%interp_p( i )
                 end do
               end associate
            end do
 
            call this%matvec_impl(wrk_out, wrk_in, lvl-1, lvl_out)
 
            do n = 1, this%lvl(lvl)%nnodes
               associate(node => this%lvl(lvl)%nodes(n))
                 do i = 1, node%ndofs
                    vec_out( node%gid ) = vec_out(node%gid ) + wrk_out( node%dofs(i) ) * node%interp_r( i )
                 end do
               end associate
            end do
          end associate
       else if (lvl_out .lt. lvl) then
          call this%matvec_impl(vec_out, vec_in, lvl-1, lvl_out)
       else
          call neko_error("TAMG: matvec level numbering problem.")
       end if
    end if
  end subroutine tamg_matvec_impl
 
 
  recursive subroutine tamg_matvec_flat_impl(this, vec_out, vec_in, lvl_blah, lvl_out)
    class(tamg_hierarchy_t), intent(inout) :: this
    real(kind=rp), intent(inout) :: vec_out(:)
    real(kind=rp), intent(inout) :: vec_in(:)
    integer, intent(in) :: lvl_blah
    integer, intent(in) :: lvl_out
    integer :: i, n, cdof, lvl
 
    lvl = lvl_out
    n = this%lvl(1)%fine_lvl_dofs
    if (lvl .eq. 0) then
       call this%ax%compute(vec_out, vec_in, this%coef, this%msh, this%Xh)
       call this%gs_h%op(vec_out, n, gs_op_add)
       call this%blst%apply(vec_out, n)
    else
       associate( wrk_in => this%lvl(1)%wrk_in, wrk_out => this%lvl(1)%wrk_out)
         do i = 1, n
            cdof = this%lvl(lvl)%map_finest2lvl(i)
            wrk_in(i) = vec_in( cdof )
         end do
 
         call this%gs_h%op(wrk_in, n, gs_op_add)
         call col2( wrk_in, this%coef%mult, n)
         call this%blst%apply(wrk_in, n)
 
         call this%ax%compute(wrk_out, wrk_in, this%coef, this%msh, this%Xh)
         call this%gs_h%op(wrk_out, n, gs_op_add)
         call this%blst%apply(wrk_out, n)
 
         call col2(wrk_out, this%coef%mult, n)
 
         call rzero(vec_out, this%lvl(lvl)%nnodes)
         !$omp parallel do private(cdof)
         do i = 1, n
            cdof = this%lvl(lvl)%map_finest2lvl(i)
            !$omp atomic
            vec_out(cdof) = vec_out(cdof) + wrk_out( i )
         end do
         !$omp end parallel do
       end associate
    end if
  end subroutine tamg_matvec_flat_impl
 
 
 
  subroutine tamg_restriction_operator(this, vec_out, vec_in, lvl)
    class(tamg_hierarchy_t), intent(inout) :: this
    real(kind=rp), intent(inout) :: vec_out(:)
    real(kind=rp), intent(inout) :: vec_in(:)
    integer, intent(in) :: lvl
    integer :: i, n, node_start, node_end, node_id
 
    vec_out = 0d0
    if (lvl-1 .eq. 0) then
       call col2(vec_in, this%coef%mult, this%lvl(lvl)%fine_lvl_dofs)
    end if
    do n = 1, this%lvl(lvl)%nnodes
       associate(node => this%lvl(lvl)%nodes(n))
         do i = 1, node%ndofs
            vec_out( node%gid ) = vec_out( node%gid ) + vec_in( node%dofs(i) ) * node%interp_r( i )
         end do
       end associate
    end do
  end subroutine tamg_restriction_operator
 
  subroutine tamg_prolongation_operator(this, vec_out, vec_in, lvl)
    class(tamg_hierarchy_t), intent(inout) :: this
    real(kind=rp), intent(inout) :: vec_out(:)
    real(kind=rp), intent(inout) :: vec_in(:)
    integer, intent(in) :: lvl
    integer :: i, n, node_start, node_end, node_id
 
    vec_out = 0d0
    do n = 1, this%lvl(lvl)%nnodes
       associate(node => this%lvl(lvl)%nodes(n))
         do i = 1, node%ndofs
            vec_out( node%dofs(i) ) = vec_out( node%dofs(i) ) + vec_in( node%gid ) * node%interp_p( i )
         end do
       end associate
    end do
    if (lvl-1 .eq. 0) then
       call this%gs_h%op(vec_out, this%lvl(lvl)%fine_lvl_dofs, gs_op_add)
       call col2(vec_out, this%coef%mult, this%lvl(lvl)%fine_lvl_dofs)
       call this%blst%apply(vec_out, this%lvl(lvl)%fine_lvl_dofs)
    end if
  end subroutine tamg_prolongation_operator
 
 
  subroutine tamg_device_matvec_flat_impl(this, vec_out, vec_in, vec_out_d, vec_in_d, lvl_out)
    class(tamg_hierarchy_t), intent(inout) :: this
    real(kind=rp), intent(inout) :: vec_out(:)
    real(kind=rp), intent(inout) :: vec_in(:)
    type(c_ptr) :: vec_out_d
    type(c_ptr) :: vec_in_d
    integer, intent(in) :: lvl_out
    integer :: i, n, cdof, lvl
 
    lvl = lvl_out
    n = this%lvl(1)%fine_lvl_dofs
    if (lvl .eq. 0) then
       call this%ax%compute(vec_out, vec_in, this%coef, this%msh, this%Xh)
       call this%gs_h%op(vec_out, n, gs_op_add, glb_cmd_event)
       call device_stream_wait_event(glb_cmd_queue, glb_cmd_event, 0)
       call this%blst%apply(vec_out, n)
    else
 
       associate( wrk_in_d => this%lvl(1)%wrk_in_d, wrk_out_d => this%lvl(1)%wrk_out_d)
         call device_masked_gather_copy_0(wrk_in_d, vec_in_d, this%lvl(lvl)%map_finest2lvl_d, this%lvl(lvl)%nnodes, n)
         call this%gs_h%op(this%lvl(1)%wrk_in, n, gs_op_add, glb_cmd_event)
         call device_stream_wait_event(glb_cmd_queue, glb_cmd_event, 0)
         call device_col2( wrk_in_d, this%coef%mult_d, n)
         call this%blst%apply(this%lvl(1)%wrk_in, n)
 
         call this%ax%compute(this%lvl(1)%wrk_out, this%lvl(1)%wrk_in, this%coef, this%msh, this%Xh)
         call this%gs_h%op(this%lvl(1)%wrk_out, n, gs_op_add, glb_cmd_event)
         call device_stream_wait_event(glb_cmd_queue, glb_cmd_event, 0)
         call this%blst%apply(this%lvl(1)%wrk_out, n)
 
         call device_col2( wrk_out_d, this%coef%mult_d, n)
 
         call device_rzero(vec_out_d, this%lvl(lvl)%nnodes)
         call device_masked_atomic_reduction_0(vec_out_d, wrk_out_d, this%lvl(lvl)%map_finest2lvl_d, this%lvl(lvl)%nnodes, n)
       end associate
 
    end if
  end subroutine tamg_device_matvec_flat_impl
 
  subroutine tamg_device_restriction_operator(this, vec_out_d, vec_in_d, lvl)
    class(tamg_hierarchy_t), intent(inout) :: this
    type(c_ptr) :: vec_out_d
    type(c_ptr) :: vec_in_d
    integer, intent(in) :: lvl
    integer :: i, n, m
    n = this%lvl(lvl)%nnodes
    m = this%lvl(lvl)%fine_lvl_dofs
    if (lvl-1 .eq. 0) then
       call device_col2(vec_in_d, this%coef%mult_d, m)
    end if
    call device_rzero(vec_out_d, n)
    call device_masked_atomic_reduction_0(vec_out_d, vec_in_d, this%lvl(lvl)%map_f2c_d, n, m)
  end subroutine tamg_device_restriction_operator
 
  subroutine tamg_device_prolongation_operator(this, vec_out_d, vec_in_d, lvl, vec_out)
    class(tamg_hierarchy_t), intent(inout) :: this
    real(kind=rp), intent(inout) :: vec_out(:)
    type(c_ptr) :: vec_out_d
    type(c_ptr) :: vec_in_d
    integer, intent(in) :: lvl
    integer :: i, n, m
    n = this%lvl(lvl)%nnodes
    m = this%lvl(lvl)%fine_lvl_dofs
    call device_masked_gather_copy_0(vec_out_d, vec_in_d, this%lvl(lvl)%map_f2c_d, n, m)
    if (lvl-1 .eq. 0) then
       call this%gs_h%op(vec_out, m, gs_op_add, glb_cmd_event)
       call device_stream_wait_event(glb_cmd_queue, glb_cmd_event, 0)
       call device_col2( vec_out_d, this%coef%mult_d, m)
       call this%blst%apply( vec_out, m)
    end if
  end subroutine tamg_device_prolongation_operator
 
end module tree_amg