schwarz.f90

Go to the documentation of this file.

! Copyright (c) 2008-2020, UCHICAGO ARGONNE, LLC.
!
! The UChicago Argonne, LLC as Operator of Argonne National
! Laboratory holds copyright in the Software. The copyright holder
! reserves all rights except those expressly granted to licensees,
! and U.S. Government license rights.
!
! Redistribution and use in source and binary forms, with or without
! modification, are permitted provided that the following conditions
! are met:
!
! 1. Redistributions of source code must retain the above copyright
! notice, this list of conditions and the disclaimer below.
!
! 2. Redistributions in binary form must reproduce the above copyright
! notice, this list of conditions and the disclaimer (as noted below)
! in the documentation and/or other materials provided with the
! distribution.
!
! 3. Neither the name of ANL 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
! UCHICAGO ARGONNE, LLC, THE U.S. DEPARTMENT OF
! ENERGY 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.
!
! Additional BSD Notice
! ---------------------
! 1. This notice is required to be provided under our contract with
! the U.S. Department of Energy (DOE). This work was produced at
! Argonne National Laboratory under Contract
! No. DE-AC02-06CH11357 with the DOE.
!
! 2. Neither the United States Government nor UCHICAGO ARGONNE,
! LLC nor any of their employees, makes any warranty,
! express or implied, or assumes any liability or responsibility for the
! accuracy, completeness, or usefulness of any information, apparatus,
! product, or process disclosed, or represents that its use would not
! infringe privately-owned rights.
!
! 3. Also, reference herein to any specific commercial products, process,
! or services by trade name, trademark, manufacturer or otherwise does
! not necessarily constitute or imply its endorsement, recommendation,
! or favoring by the United States Government or UCHICAGO ARGONNE LLC.
! The views and opinions of authors expressed
! herein do not necessarily state or reflect those of the United States
! Government or UCHICAGO ARGONNE, LLC, and shall
! not be used for advertising or product endorsement purposes.
!
module schwarz
  use num_types, only : rp, i8
  use math, only : rzero, rone
  use mesh, only : mesh_t
  use space, only : space_t, gll
  use dofmap, only : dofmap_t
  use gather_scatter, only : gs_t, gs_op_add
  use device_schwarz, only : device_schwarz_extrude, device_schwarz_toext3d, &
       device_schwarz_toreg3d
  use device_math, only : device_rzero, device_col2
  use fdm, only : fdm_t
  use device, only : device_map, device_alloc, device_memcpy, &
       device_event_create, host_to_device, device_to_host, &
       device_get_ptr, glb_cmd_queue, aux_cmd_queue, &
       device_event_record, device_event_sync, device_stream_wait_event, &
       device_event_destroy, device_free
  use neko_config, only : neko_bcknd_device
  use bc_list, only : bc_list_t
  use, intrinsic :: iso_c_binding, only : c_sizeof, c_ptr, c_null_ptr, &
       c_associated
  !$ use omp_lib
  implicit none
  private
 
  type, public :: schwarz_t
     real(kind=rp), allocatable :: work1(:)
     real(kind=rp), allocatable :: work2(:)
     real(kind=rp), allocatable :: wt(:,:,:,:,:)
     type(c_ptr) :: work1_d = c_null_ptr
     type(c_ptr) :: work2_d = c_null_ptr
     type(c_ptr) :: wt_d = c_null_ptr
     type(space_t) :: xh_schwarz
     type(gs_t) :: gs_schwarz
     type(dofmap_t) :: dm_schwarz
     type(fdm_t) :: fdm
     type(space_t), pointer :: xh => null()
     type(bc_list_t), pointer :: bclst => null()
     type(dofmap_t), pointer :: dof => null()
     type(gs_t), pointer :: gs_h => null()
     type(mesh_t), pointer :: msh => null()
     type(c_ptr) :: event = c_null_ptr
     logical :: local_gs = .false.
     type(gs_t), allocatable :: gs_h_local
   contains
     procedure, pass(this) :: init => schwarz_init
     procedure, pass(this) :: free => schwarz_free
     procedure, pass(this) :: compute => schwarz_compute
  end type schwarz_t
 
contains
 
  subroutine schwarz_init(this, Xh, dof, gs_h, bclst, msh)
    class(schwarz_t), target, intent(inout) :: this
    type(space_t), target, intent(inout) :: Xh
    type(dofmap_t), target, intent(in) :: dof
    type(gs_t), target, intent(inout) :: gs_h
    type(mesh_t), target, intent(inout) :: msh
    type(bc_list_t), target, intent(inout) :: bclst
    integer :: nthrds
 
    call this%free()
 
    call this%Xh_schwarz%init(gll, xh%lx+2, xh%lx+2, xh%lx+2)
    call this%dm_schwarz%init(msh, this%Xh_schwarz)
    call this%gs_schwarz%init(this%dm_schwarz)
 
    allocate(this%work1(this%dm_schwarz%size()))
    allocate(this%work2(this%dm_schwarz%size()))
    allocate(this%wt(xh%lx, xh%lx, 4, msh%gdim, msh%nelv))
 
    call this%fdm%init(xh, dof, gs_h)
 
 
    this%msh => msh
    this%Xh => xh
    this%bclst => bclst
    this%dof => dof
 
    nthrds = 1
    !$omp parallel
    !$ nthrds = omp_get_num_threads()
    !$omp end parallel
 
    ! If we are running multithreaded, we need a local gs object,
    ! otherwise we can reuse the external one
    if (nthrds .gt. 1) then
       allocate(this%gs_h_local)
       call this%gs_h_local%init(this%dof)
       this%gs_h => this%gs_h_local
       this%local_gs = .true.
    else
       this%gs_h => gs_h
       this%local_gs = .false.
    end if
 
    if (neko_bcknd_device .eq. 1) then
       call device_map(this%work1, this%work1_d, this%dm_schwarz%size())
       call device_map(this%work2, this%work2_d, this%dm_schwarz%size())
    end if
 
    call schwarz_setup_wt(this)
    if (neko_bcknd_device .eq. 1) then
       call device_alloc(this%wt_d, &
            int(this%dof%size(), i8) * int(c_sizeof(this%work1(1)), i8))
       call rone(this%work1, this%dof%size())
       call schwarz_wt3d(this%work1, this%wt, xh%lx, msh%nelv)
       call device_memcpy(this%work1, this%wt_d, this%dof%size(), &
            host_to_device, sync = .false.)
       call device_event_create(this%event, 2)
    end if
  end subroutine schwarz_init
 
  subroutine schwarz_free(this)
    class(schwarz_t), intent(inout) :: this
 
    if (allocated(this%work1)) deallocate(this%work1)
    if (allocated(this%work2)) deallocate(this%work2)
    if (allocated(this%wt)) deallocate(this%wt)
 
    if (c_associated(this%work1_d)) then
       call device_free(this%work1_d)
    end if
 
    if (c_associated(this%work2_d)) then
       call device_free(this%work2_d)
    end if
 
    if (c_associated(this%wt_d)) then
       call device_free(this%wt_d)
    end if
 
    call this%Xh_schwarz%free()
    call this%gs_schwarz%free()
    !why cant I do this?
    !call dofmap_free(this%dm_schwarz)
    call this%dm_schwarz%free()
    call this%fdm%free()
 
    nullify(this%Xh)
    nullify(this%bclst)
    nullify(this%dof)
    if (allocated(this%gs_h_local)) then
       call this%gs_h_local%free()
       deallocate(this%gs_h_local)
    end if
    nullify(this%gs_h)
    nullify(this%msh)
 
    this%local_gs = .false.
    if (c_associated(this%event)) then
       call device_event_destroy(this%event)
    end if
  end subroutine schwarz_free
  subroutine schwarz_setup_wt(this)
    class(schwarz_t), intent(inout) :: this
    integer :: enx, eny, enz, n, ie, k, ns
    real(kind=rp), parameter :: zero = 0.0_rp
    real(kind=rp), parameter :: one = 1.0_rp
    associate(work1 => this%work1, work2 => this%work2, msh => this%msh, &
         xh => this%Xh, xh_schwarz => this%Xh_schwarz)
 
      n = this%dof%size()
 
      enx = xh_schwarz%lx
      eny = xh_schwarz%ly
      enz = xh_schwarz%lz
      if (.not. msh%gdim .eq. 3) enz = 1
      ns = enx * eny * enz * msh%nelv
 
      call rone(work2, ns)
      call rzero(work1, ns)
 
      !   Sum overlap region (border excluded)
      !   Cred to PFF for this, very clever
      call schwarz_extrude(work1, 0, zero, work2, 0, one, enx, eny, enz, &
           msh%nelv)
      if (neko_bcknd_device .eq. 1) then
         call device_memcpy(work2, this%work2_d, ns, &
              host_to_device, sync = .false.)
         call this%gs_schwarz%op(work2, ns, gs_op_add)
         call device_memcpy(work2, this%work2_d, ns, &
              device_to_host, sync = .true.)
      else
         call this%gs_schwarz%op(work2, ns, gs_op_add)
      end if
      call schwarz_extrude(work2, 0, one, work1, 0, -one, enx, eny, enz, &
           msh%nelv)
      call schwarz_extrude_single(work2, 2, one, 0, one, enx, eny, enz, &
           msh%nelv)
 
      ! if(.not.if3d) then ! Go back to regular size array
      !    call hsmg_schwarz_toreg2d(mg_work,mg_work(i),mg_nh(l))
      ! else
      call schwarz_toreg3d(work1, work2, xh%lx, msh%nelv)
      ! endif
 
      if (neko_bcknd_device .eq. 1) then
         call device_memcpy(work1, this%work1_d, n, &
              host_to_device, sync = .false.)
         call this%gs_h%op(work1, n, gs_op_add)
         call device_memcpy(work1, this%work1_d, n, &
              device_to_host, sync = .true.)
      else
         call this%gs_h%op(work1, n, gs_op_add)
      end if
 
      k = 1
      do ie = 1, msh%nelv
         if (msh%gdim .eq. 2) then
            call schwarz_setup_schwarz_wt2d_2(this%wt, ie, xh%lx, &
                 work1(k), msh%nelv)
         end if
         if (this%msh%gdim .eq. 3) then
            call schwarz_setup_schwarz_wt3d_2(this%wt, ie, xh%lx, &
                 work1(k), msh%nelv)
            k = k + xh%lxyz
         end if
      end do
    end associate
  end subroutine schwarz_setup_wt
 
  subroutine schwarz_setup_schwarz_wt2d_2(wt, ie, n, work, nelv)
    integer, intent(in) :: n, nelv
    real(kind=rp), intent(inout) :: wt(n, 4, 2, nelv)
    real(kind=rp), intent(inout) :: work(n, n)
    integer :: ie, i, j
    do j = 1, n
       wt(j, 1, 1, ie) = 1.0_rp / work(1, j)
       wt(j, 2, 1, ie) = 1.0_rp / work(2, j)
       wt(j, 3, 1, ie) = 1.0_rp / work(n - 1, j)
       wt(j, 4, 1, ie) = 1.0_rp / work(n, j)
    end do
    do i = 1, n
       wt(i, 1, 2, ie) = 1.0_rp / work(i, 1)
       wt(i, 2, 2, ie) = 1.0_rp / work(i, 2)
       wt(i, 3, 2, ie) = 1.0_rp / work(i, n - 1)
       wt(i, 4, 2, ie) = 1.0_rp / work(i, n)
    end do
 
    return
  end subroutine schwarz_setup_schwarz_wt2d_2
 
  subroutine schwarz_setup_schwarz_wt3d_2(wt, ie, n, work, nelv)
    integer, intent(in) :: n, nelv, ie
    real(kind=rp), intent(inout) :: wt(n, n, 4, 3, nelv)
    real(kind=rp), intent(inout) :: work(n, n, n)
    integer :: i, j, k
 
    do k = 1, n
       do j = 1, n
          wt(j, k, 1, 1, ie) = 1.0_rp / work(1, j, k)
          wt(j, k, 2, 1, ie) = 1.0_rp / work(2, j, k)
          wt(j, k, 3, 1, ie) = 1.0_rp / work(n - 1, j, k)
          wt(j, k, 4, 1, ie) = 1.0_rp / work(n, j, k)
       end do
    end do
 
    do k = 1, n
       do i = 1, n
          wt(i, k, 1, 2, ie) = 1.0_rp / work(i, 1, k)
          wt(i, k, 2, 2, ie) = 1.0_rp / work(i, 2, k)
          wt(i, k, 3, 2, ie) = 1.0_rp / work(i, n - 1, k)
          wt(i, k, 4, 2, ie) = 1.0_rp / work(i, n, k)
       end do
    end do
 
    do j = 1, n
       do i = 1, n
          wt(i, j, 1, 3, ie) = 1.0_rp / work(i, j, 1)
          wt(i, j, 2, 3, ie) = 1.0_rp / work(i, j, 2)
          wt(i, j, 3, 3, ie) = 1.0_rp / work(i, j, n - 1)
          wt(i, j, 4, 3, ie) = 1.0_rp / work(i, j, n)
       end do
    end do
 
  end subroutine schwarz_setup_schwarz_wt3d_2
 
  subroutine schwarz_toreg3d(b, a, n, nelv)
    integer, intent(in) :: n, nelv
    real(kind=rp), intent(inout) :: a(0:n+1, 0:n+1, 0:n+1, nelv)
    real(kind=rp), intent(inout) :: b(n, n, n, nelv)
    integer :: i, j, k, ie
    !$omp parallel do private(ie, i, j, k)
    do ie = 1, nelv
       do k = 1, n
          do j = 1, n
             do i = 1, n
                b(i, j, k, ie) = a(i, j, k, ie)
             end do
          end do
       end do
    end do
    !$omp end parallel do
  end subroutine schwarz_toreg3d
 
  subroutine schwarz_toext3d(a, b, n, nelv)
    integer, intent(in) :: n, nelv
    real(kind=rp), intent(inout) :: a(0:n+1, 0:n+1, 0:n+1, nelv), &
         b(n, n, n, nelv)
    integer :: i, j, k, ie
 
    call rzero(a, (n+2)*(n+2)*(n+2)*nelv)
    !$omp parallel do private(ie, i, j, k)
    do ie = 1, nelv
       do k = 1, n
          do j = 1, n
             do i = 1, n
                a(i, j, k, ie) = b(i, j, k, ie)
             end do
          end do
       end do
    end do
    !$omp end parallel do
  end subroutine schwarz_toext3d
 
  subroutine schwarz_extrude(arr1, l1, f1, arr2, l2, f2, nx, ny, nz, nelv)
    integer, intent(in) :: l1, l2, nx, ny, nz, nelv
    real(kind=rp), intent(inout) :: arr1(nx, ny, nz, nelv)
    real(kind=rp), intent(in) :: arr2(nx, ny, nz, nelv)
    real(kind=rp), intent(in) :: f1, f2
    integer :: i, j, k, ie, i0, i1
    i0 = 2
    i1 = nx - 1
 
    if (nz .eq. 1) then
       !$omp parallel do private(ie, i, j)
       do ie = 1, nelv
          do j = i0, i1
             arr1(l1 + 1, j, 1, ie) = f1 * arr1(l1 + 1, j, 1, ie) &
                  + f2 * arr2(l2 + 1, j, 1, ie)
             arr1(nx - l1, j, 1, ie) = f1 * arr1(nx - l1, j, 1, ie) &
                  + f2 * arr2(nx - l2, j, 1, ie)
          end do
          do i = i0, i1
             arr1(i, l1 + 1, 1, ie) = f1 * arr1(i, l1 + 1, 1, ie) &
                  + f2 * arr2(i, l2 + 1, 1, ie)
             arr1(i, ny - l1, 1, ie) = f1 * arr1(i, ny - l1, 1, ie) &
                  + f2 * arr2(i, nx - l2, 1, ie)
          end do
       end do
       !$omp end parallel do
    else
       !$omp parallel do private(ie, i, j, k)
       do ie = 1, nelv
          do k = i0, i1
             do j = i0, i1
                arr1(l1 + 1, j, k, ie) = f1 * arr1(l1 + 1, j, k, ie) &
                     + f2 * arr2(l2 + 1, j, k, ie)
                arr1(nx - l1, j, k, ie) = f1 * arr1(nx - l1, j, k, ie) &
                     + f2 * arr2(nx - l2, j, k, ie)
             end do
          end do
          do k = i0, i1
             do i = i0, i1
                arr1(i, l1 + 1, k, ie) = f1 * arr1(i, l1 + 1, k, ie) &
                     + f2 * arr2(i, l2 + 1, k, ie)
                arr1(i, nx - l1, k, ie) = f1 * arr1(i, nx - l1, k, ie) &
                     + f2 * arr2(i, nx - l2, k, ie)
             end do
          end do
          do j = i0, i1
             do i = i0, i1
                arr1(i, j, l1 + 1, ie) = f1 * arr1(i, j, l1 + 1, ie) &
                     + f2 * arr2(i, j, l2 + 1, ie)
                arr1(i, j, nx - l1, ie) = f1 * arr1(i, j, nx - l1, ie) &
                     + f2 * arr2(i, j, nx - l2, ie)
             end do
          end do
       end do
       !$omp end parallel do
    end if
  end subroutine schwarz_extrude
 
  subroutine schwarz_extrude_single(arr, l1, f1, l2, f2, nx, ny, nz, nelv)
    integer, intent(in) :: l1, l2, nx, ny, nz, nelv
    real(kind=rp), intent(inout) :: arr(nx, ny, nz, nelv)
    real(kind=rp), intent(in) :: f1, f2
    integer :: i, j, k, ie, i0, i1
    i0 = 2
    i1 = nx - 1
 
    if (nz .eq. 1) then
       !$omp parallel do private(ie, i, j)
       do ie = 1, nelv
          do j = i0, i1
             arr(l1 + 1, j, 1, ie) = f1 * arr(l1 + 1, j, 1, ie) &
                  + f2 * arr(l2 + 1, j, 1, ie)
             arr(nx - l1, j, 1, ie) = f1 * arr(nx - l1, j, 1, ie) &
                  + f2 * arr(nx - l2, j, 1, ie)
          end do
          do i = i0, i1
             arr(i, l1 + 1, 1, ie) = f1 * arr(i, l1 + 1, 1, ie) &
                  + f2 * arr(i, l2 + 1, 1, ie)
             arr(i, ny - l1, 1, ie) = f1 * arr(i, ny - l1, 1, ie) &
                  + f2 * arr(i, nx - l2, 1, ie)
          end do
       end do
       !$omp end parallel do
    else
       !$omp parallel do private(ie, i, j, k)
       do ie = 1, nelv
          do k = i0, i1
             do j = i0, i1
                arr(l1 + 1, j, k, ie) = f1 * arr(l1 + 1, j, k, ie) &
                     + f2 * arr(l2 + 1, j, k, ie)
                arr(nx - l1, j, k, ie) = f1 * arr(nx - l1, j, k, ie) &
                     + f2 * arr(nx - l2, j, k, ie)
             end do
          end do
          do k = i0, i1
             do i = i0, i1
                arr(i, l1 + 1, k, ie) = f1 * arr(i, l1 + 1, k, ie) &
                     + f2 * arr(i, l2 + 1, k, ie)
                arr(i, nx - l1, k, ie) = f1 * arr(i, nx - l1, k, ie) &
                     + f2 * arr(i, nx - l2, k, ie)
             end do
          end do
          do j = i0, i1
             do i = i0, i1
                arr(i, j, l1 + 1, ie) = f1 * arr(i, j, l1 + 1, ie) &
                     + f2 * arr(i, j, l2 + 1, ie)
                arr(i, j, nx - l1, ie) = f1 * arr(i, j, nx - l1, ie) &
                     + f2 * arr(i, j, nx - l2, ie)
             end do
          end do
       end do
       !$omp end parallel do
    end if
  end subroutine schwarz_extrude_single
 
  subroutine schwarz_compute(this, e, r)
    class(schwarz_t), intent(inout) :: this
    real(kind=rp), dimension(this%dof%size()), intent(inout) :: e, r
    integer :: n, enx, eny, enz, ns
    real(kind=rp), parameter :: zero = 0.0_rp
    real(kind=rp), parameter :: one = 1.0_rp
    type(c_ptr) :: e_d, r_d
    associate(work1 => this%work1, work1_d => this%work1_d, &
         work2 => this%work2, work2_d => this%work2_d)
 
      n = this%dof%size()
      enx = this%Xh_schwarz%lx
      eny = this%Xh_schwarz%ly
      enz = this%Xh_schwarz%lz
      if (.not. this%msh%gdim .eq. 3) enz = 1
      ns = enx * eny * enz * this%msh%nelv
      if (neko_bcknd_device .eq. 1) then
         r_d = device_get_ptr(r)
         e_d = device_get_ptr(e)
         call device_event_record(this%event, glb_cmd_queue)
         call device_stream_wait_event(aux_cmd_queue, this%event, 0)
         call device_schwarz_toext3d(work1_d, r_d, this%Xh%lx, &
              this%msh%nelv, aux_cmd_queue)
         call device_schwarz_extrude(work1_d, 0, zero, work1_d, 2, one, &
              enx, eny, enz, this%msh%nelv, aux_cmd_queue)
 
         this%gs_schwarz%bcknd%gs_stream = aux_cmd_queue
         call this%gs_schwarz%op(work1, ns, gs_op_add, this%event)
         call device_event_sync(this%event)
         call device_schwarz_extrude(work1_d, 0, one, work1_d, 2, -one, &
              enx, eny, enz, this%msh%nelv, aux_cmd_queue)
 
         call this%fdm%compute(work2, work1, aux_cmd_queue) ! do local solves
 
         call device_schwarz_extrude(work1_d, 0, zero, work2_d, 0, one, &
              enx, eny, enz, this%msh%nelv, aux_cmd_queue)
         call this%gs_schwarz%op(work2, ns, gs_op_add, this%event)
         call device_event_sync(this%event)
 
         call device_schwarz_extrude(work2_d, 0, one, work1_d, 0, -one, &
              enx, eny, enz, this%msh%nelv, aux_cmd_queue)
         call device_schwarz_extrude(work2_d, 2, one, work2_d, 0, one, &
              enx, eny, enz, this%msh%nelv, aux_cmd_queue)
         call device_schwarz_toreg3d(e_d, work2_d, this%Xh%lx, &
              this%msh%nelv, aux_cmd_queue)
 
         this%gs_h%bcknd%gs_stream = aux_cmd_queue
         call this%gs_h%op(e, n, gs_op_add, this%event)
 
         call this%bclst%apply_scalar(e, n, strm = aux_cmd_queue)
         call device_col2(e_d, this%wt_d, n, aux_cmd_queue)
 
         ! switch back to the default stream on the shared gs
         if (.not. this%local_gs) then
            call device_event_sync(this%event)
            this%gs_h%bcknd%gs_stream = glb_cmd_queue
         end if
      else
         call this%bclst%apply_scalar(r, n)
         call schwarz_toext3d(work1, r, this%Xh%lx, this%msh%nelv)
 
         !  exchange interior nodes
         call schwarz_extrude_single(work1, 0, zero, 2, one, &
              enx, eny, enz, this%msh%nelv)
         call this%gs_schwarz%op(work1, ns, gs_op_add)
         call schwarz_extrude_single(work1, 0, one, 2, -one, &
              enx, eny, enz, this%msh%nelv)
 
         call this%fdm%compute(work2, work1) ! do local solves
 
         !   Sum overlap region (border excluded)
         call schwarz_extrude(work1, 0, zero, work2, 0, one, &
              enx, eny, enz, this%msh%nelv)
         call this%gs_schwarz%op(work2, ns, gs_op_add)
         call schwarz_extrude(work2, 0, one, work1, 0, -one, &
              enx, eny, enz, this%msh%nelv)
         call schwarz_extrude_single(work2, 2, one, 0, one, &
              enx, eny, enz, this%msh%nelv)
 
         call schwarz_toreg3d(e, work2, this%Xh%lx, this%msh%nelv)
 
         ! sum border nodes
         call this%gs_h%op(e, n, gs_op_add)
         call this%bclst%apply_scalar(e, n)
 
         call schwarz_wt3d(e, this%wt, this%Xh%lx, this%msh%nelv)
      end if
    end associate
  end subroutine schwarz_compute
 
  !Apply schwarz weights along the boundary of each element.
  subroutine schwarz_wt3d(e, wt, n, nelv)
    integer, intent(in) :: n, nelv
    real(kind=rp), intent(inout) :: e(n, n, n, nelv)
    real(kind=rp), intent(inout) :: wt(n, n, 4, 3, nelv)
    integer :: ie, i, j, k
 
    !$omp parallel do private(ie, i, j, k)
    do ie = 1, nelv
       do k = 1, n
          do j = 1, n
             e(1, j, k, ie) = e(1, j, k, ie) * wt(j, k, 1, 1, ie)
             e(2, j, k, ie) = e(2, j, k, ie) * wt(j, k, 2, 1, ie)
             e(n - 1, j, k, ie) = e(n - 1, j, k, ie) * wt(j, k, 3, 1, ie)
             e(n, j, k, ie) = e(n, j, k, ie) * wt(j, k, 4, 1, ie)
          end do
       end do
       do k = 1, n
          do i = 3, n-2
             e(i, 1, k, ie) = e(i, 1, k, ie) * wt(i, k, 1, 2, ie)
             e(i, 2, k, ie) = e(i, 2, k, ie) * wt(i, k, 2, 2, ie)
             e(i, n - 1, k, ie) = e(i, n - 1, k, ie) * wt(i, k, 3, 2, ie)
             e(i, n, k, ie) = e(i, n, k, ie) * wt(i, k, 4, 2, ie)
          end do
       end do
       do j = 3, n-2
          do i = 3, n-2
             e(i, j, 1, ie) = e(i, j, 1, ie) * wt(i, j, 1, 3, ie)
             e(i, j, 2, ie) = e(i, j, 2, ie) * wt(i, j, 2, 3, ie)
             e(i, j, n - 1, ie) = e(i, j, n - 1, ie) * wt(i, j, 3, 3, ie)
             e(i, j, n, ie) = e(i, j, n, ie) * wt(i, j, 4, 3, ie)
          end do
       end do
    end do
    !$omp end parallel do
  end subroutine schwarz_wt3d
end module schwarz