pad3d Module Function

  module function pad3d(input, facets, pad_size, imethod) result(output)
    !! 3D padding operation
    implicit none

    ! Arguments
    type(array_type), intent(in), target :: input
    type(facets_type), dimension(3), intent(in) :: facets
    integer, dimension(3), intent(in) :: pad_size
    integer, intent(in) :: imethod
    type(array_type), pointer :: output

    ! Local variables
    integer :: i, j, k, m, s
    integer :: idx_in, idx_out, idx_shift
    integer :: input_size_h, input_size_w, input_size_d, num_channels
    integer :: output_size_h, output_size_w, output_size_d
    integer, dimension(5) :: output_shape

    input_size_h = input%shape(1)
    input_size_w = input%shape(2)
    input_size_d = input%shape(3)
    num_channels = input%shape(4)
    output_size_h = input_size_h + 2 * pad_size(1)
    output_size_w = input_size_w + 2 * pad_size(2)
    output_size_d = input_size_d + 2 * pad_size(3)

    output_shape = [ output_size_h, output_size_w, output_size_d, num_channels, &
         size(input%val, dim=2) ]
    output => input%create_result(array_shape = output_shape)

    ! save the facet values to indices and adj_ja
    allocate(output%indices(4 + 3 + sum( facets(:)%num )))
    output%indices(1) = imethod
    output%indices(2) = pad_size(1)
    output%indices(3) = pad_size(2)
    output%indices(4) = pad_size(3)
    output%indices(5) = facets(1)%num
    output%indices(6) = facets(2)%num
    output%indices(7) = facets(3)%num
    output%indices(8:7 + facets(1)%num) = [(facets(1)%dim(i), i=1, facets(1)%num)]
    output%indices(8 + facets(1)%num:7 + facets(1)%num + facets(2)%num) = &
         [(facets(2)%dim(i), i=1, facets(2)%num)]
    output%indices(8 + facets(1)%num + facets(2)%num:7 + &
         facets(1)%num + facets(2)%num + facets(3)%num) = &
         [(facets(3)%dim(i), i=1, facets(3)%num)]
    allocate(output%adj_ja(2, 6 * (facets(1)%num + facets(2)%num + facets(3)%num)))
    ! Edges (1D edges)
    do i = 1, facets(1)%num
       output%adj_ja(1,(i-1)*6 + 1 : (i-1)*6 + 2) = facets(1)%orig_bound(1:2,1,i)
       output%adj_ja(1,(i-1)*6 + 3 : (i-1)*6 + 4) = facets(1)%orig_bound(1:2,2,i)
       output%adj_ja(1,(i-1)*6 + 5 : (i-1)*6 + 6) = facets(1)%orig_bound(1:2,3,i)
       output%adj_ja(2,(i-1)*6 + 1 : (i-1)*6 + 2) = facets(1)%dest_bound(1:2,1,i)
       output%adj_ja(2,(i-1)*6 + 3 : (i-1)*6 + 4) = facets(1)%dest_bound(1:2,2,i)
       output%adj_ja(2,(i-1)*6 + 5 : (i-1)*6 + 6) = facets(1)%dest_bound(1:2,3,i)
    end do
    idx_shift = facets(1)%num * 6
    ! Faces (2D faces)
    do i = 1, facets(2)%num
       output%adj_ja(1,(i-1)*6 + 1 + idx_shift : (i-1)*6 + 2 + idx_shift) = &
            facets(2)%orig_bound(1:2,1,i)
       output%adj_ja(1,(i-1)*6 + 3 + idx_shift : (i-1)*6 + 4 + idx_shift) = &
            facets(2)%orig_bound(1:2,2,i)
       output%adj_ja(1,(i-1)*6 + 5 + idx_shift : (i-1)*6 + 6 + idx_shift) = &
            facets(2)%orig_bound(1:2,3,i)
       output%adj_ja(2,(i-1)*6 + 1 + idx_shift : (i-1)*6 + 2 + idx_shift) = &
            facets(2)%dest_bound(1:2,1,i)
       output%adj_ja(2,(i-1)*6 + 3 + idx_shift : (i-1)*6 + 4 + idx_shift) = &
            facets(2)%dest_bound(1:2,2,i)
       output%adj_ja(2,(i-1)*6 + 5 + idx_shift : (i-1)*6 + 6 + idx_shift) = &
            facets(2)%dest_bound(1:2,3,i)
    end do
    idx_shift = idx_shift + facets(2)%num * 6
    ! Corners (3D corners)
    do i = 1, facets(3)%num
       output%adj_ja(1,(i-1)*6 + 1 + idx_shift : (i-1)*6 + 2 + idx_shift) = &
            facets(3)%orig_bound(1:2,1,i)
       output%adj_ja(1,(i-1)*6 + 3 + idx_shift : (i-1)*6 + 4 + idx_shift) = &
            facets(3)%orig_bound(1:2,2,i)
       output%adj_ja(1,(i-1)*6 + 5 + idx_shift : (i-1)*6 + 6 + idx_shift) = &
            facets(3)%orig_bound(1:2,3,i)
       output%adj_ja(2,(i-1)*6 + 1 + idx_shift : (i-1)*6 + 2 + idx_shift) = &
            facets(3)%dest_bound(1:2,1,i)
       output%adj_ja(2,(i-1)*6 + 3 + idx_shift : (i-1)*6 + 4 + idx_shift) = &
            facets(3)%dest_bound(1:2,2,i)
       output%adj_ja(2,(i-1)*6 + 5 + idx_shift : (i-1)*6 + 6 + idx_shift) = &
            facets(3)%dest_bound(1:2,3,i)
    end do

    ! Initialise with zero
    output%val = 0._real32

    ! Copy input into the correct location in output
    do concurrent( &
         s = 1:output_shape(5), &
         m = 1:num_channels, &
         k = 1:input_size_d, &
         j = 1:input_size_w, &
         i = 1:input_size_h)
       idx_in = i + (j-1) * input_size_h + (k-1) * input_size_h * input_size_w + &
            (m-1) * input_size_h * input_size_w * input_size_d
       idx_out = (i + pad_size(1)) + &
            (j + pad_size(2) - 1) * output_size_h + &
            (k + pad_size(3) - 1) * output_size_h * output_size_w + &
            (m-1) * output_size_h * output_size_w * output_size_d
       output%val(idx_out, s) = input%val(idx_in, s)
    end do

    if(output%indices(1) .ge. 3 .and. output%indices(1) .le. 5)then
       call fill_corner_region_3d( input, output )
       call fill_edge_region_3d( input, output )
       call fill_face_region_3d( input, output )
    end if

    output%get_partial_left => get_partial_pad3d
    output%get_partial_left_val => get_partial_pad3d_val
    if(input%requires_grad)then
       output%requires_grad = .true.
       output%is_forward = input%is_forward
       output%operation = 'pad'
       output%left_operand => input
    end if

  end function pad3d