tensor3.f90

Go to the documentation of this file.

! Copyright (c) 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 tensor3
  use neko_config, only : neko_bcknd_device
  use num_types, only : rp
  use device, only : device_map, device_unmap, device_memcpy, device_sync
  use math, only : cfill
  use device_math, only : device_copy, device_cfill
  use utils, only : neko_error, neko_varname_len
  use, intrinsic :: iso_c_binding
  implicit none
  private
 
  type, public :: tensor3_t
     real(kind=rp), allocatable :: x(:,:,:) 
     character(len=NEKO_VARNAME_LEN) :: name = ""
     type(c_ptr) :: x_d = c_null_ptr 
     integer, private :: n1 = 0
     integer, private :: n2 = 0
     integer, private :: n3 = 0
     integer, private :: n = 0
   contains
     procedure, pass(t) :: init => tensor3_init
     procedure, pass(t) :: free => tensor3_free
     procedure, pass(t) :: size => tensor3_size
     procedure, pass(t) :: copy_from => tensor3_copy_from
     procedure, pass(t) :: get_n1 => tensor3_n1
     procedure, pass(t) :: get_n2 => tensor3_n2
     procedure, pass(t) :: get_n3 => tensor3_n3
     procedure, pass(t) :: get_dims => tensor3_get_dims
     procedure, pass(t) :: tensor3_assign_tensor3
     procedure, pass(t) :: tensor3_assign_scalar
 
     generic :: assignment(=) => tensor3_assign_tensor3, &
          tensor3_assign_scalar
 
     procedure, pass(t), private :: alloc => tensor3_allocate
  end type tensor3_t
 
  type, public :: tensor3_ptr_t
     type(tensor3_t), pointer :: ptr => null()
   contains
     procedure, pass(this) :: init => tensor3_ptr_init
     procedure, pass(this) :: free => tensor3_ptr_free
  end type tensor3_ptr_t
 
contains
 
  subroutine tensor3_init(t, n1, n2, n3, name)
    class(tensor3_t), intent(inout) :: t
    integer, intent(in) :: n1
    integer, intent(in) :: n2
    integer, intent(in) :: n3
    character(len=*), intent(in), optional :: name
 
    ! t%alloc zeroes the device side (and synchronizes) before any
    ! host-side touch: under zero-copy the device then faults the
    ! pages first (device first touch), which gives contiguous
    ! physical mappings and thus better GPU TLB utilisation; rewriting
    ! the zeros on the host afterwards is benign.
    call t%alloc(n1, n2, n3)
 
    if (present(name)) then
       t%name = name
    end if
 
  end subroutine tensor3_init
 
  subroutine tensor3_allocate(t, n1, n2, n3)
    class(tensor3_t), intent(inout) :: t
    integer, intent(in) :: n1
    integer, intent(in) :: n2
    integer, intent(in) :: n3
 
    call t%free()
 
    allocate(t%x(n1, n2, n3))
    t%n1 = n1
    t%n2 = n2
    t%n3 = n3
    t%n = n1*n2*n3
 
    if (neko_bcknd_device .eq. 1) then
       call device_map(t%x, t%x_d, t%n)
       call device_cfill(t%x_d, 0.0_rp, t%n)
       call device_sync()
    end if
    call cfill(t%x, 0.0_rp, t%n)
 
  end subroutine tensor3_allocate
 
  subroutine tensor3_free(t)
    class(tensor3_t), intent(inout) :: t
 
    if (allocated(t%x)) then
       if (neko_bcknd_device .eq. 1) then
          call device_unmap(t%x, t%x_d)
       end if
       deallocate(t%x)
    end if
 
    t%n1 = 0
    t%n2 = 0
    t%n3 = 0
    t%n = 0
    t%name = ""
 
  end subroutine tensor3_free
 
  pure function tensor3_size(t) result(s)
    class(tensor3_t), intent(in) :: t
    integer :: s
    s = t%n
  end function tensor3_size
 
  subroutine tensor3_copy_from(t, memdir, sync)
    class(tensor3_t), intent(inout) :: t
    integer, intent(in) :: memdir
    logical, intent(in) :: sync
 
    if (neko_bcknd_device .eq. 1) then
       call device_memcpy(t%x, t%x_d, t%n, memdir, sync)
    end if
 
  end subroutine tensor3_copy_from
 
  pure function tensor3_n1(t) result(n1)
    class(tensor3_t), intent(in) :: t
    integer :: n1
    n1 = t%n1
  end function tensor3_n1
 
  pure function tensor3_n2(t) result(n2)
    class(tensor3_t), intent(in) :: t
    integer :: n2
    n2 = t%n2
  end function tensor3_n2
 
  pure function tensor3_n3(t) result(n3)
    class(tensor3_t), intent(in) :: t
    integer :: n3
    n3 = t%n3
  end function tensor3_n3
 
  pure function tensor3_get_dims(t) result(dims)
    class(tensor3_t), intent(in) :: t
    integer :: dims(3)
    dims = [t%n1, t%n2, t%n3]
  end function tensor3_get_dims
 
  subroutine tensor3_assign_tensor3(t, w)
    class(tensor3_t), intent(inout) :: t
    type(tensor3_t), intent(in) :: w
 
    if (allocated(t%x)) then
       call t%free()
    end if
 
    if (.not. allocated(t%x)) then
 
       t%n1 = w%n1
       t%n2 = w%n2
       t%n3 = w%n3
       t%n = w%n
       allocate(t%x(t%n1, t%n2, t%n3))
 
       if (neko_bcknd_device .eq. 1) then
          call device_map(t%x, t%x_d, t%n)
       end if
 
    end if
 
    if (neko_bcknd_device .eq. 1) then
       call device_copy(t%x_d, w%x_d, t%n)
    else
       t%x = w%x
    end if
 
    t%name = w%name
 
  end subroutine tensor3_assign_tensor3
 
  subroutine tensor3_assign_scalar(t, s)
    class(tensor3_t), intent(inout) :: t
    real(kind=rp), intent(in) :: s
 
    if (.not. allocated(t%x)) then
       call neko_error('tensor3 not allocated')
    end if
 
    if (neko_bcknd_device .eq. 1) then
       call device_cfill(t%x_d, s, t%n)
    else
       t%x = s
    end if
 
  end subroutine tensor3_assign_scalar
 
  ! ========================================================================== !
  ! tensor3 pointer type subroutines
 
  subroutine tensor3_ptr_init(this, ptr)
    class(tensor3_ptr_t), intent(inout) :: this
    type(tensor3_t), target, intent(in) :: ptr
 
    call this%free()
    this%ptr => ptr
  end subroutine tensor3_ptr_init
 
  subroutine tensor3_ptr_free(this)
    class(tensor3_ptr_t), intent(inout) :: this
 
    if (associated(this%ptr)) then
       nullify(this%ptr)
    end if
 
  end subroutine tensor3_ptr_free
 
end module tensor3