pad2d Module Function

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

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

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

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

    output_shape = [ &
         output_size_h, output_size_w, 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(3 + 2 + sum( facets(:)%num )))
    output%indices(1) = imethod
    output%indices(2) = pad_size(1)
    output%indices(3) = pad_size(2)
    output%indices(4) = facets(1)%num
    output%indices(5) = facets(2)%num
    output%indices(6:5 + facets(1)%num) = [(facets(1)%dim(i), i=1, facets(1)%num)]
    output%indices(6 + facets(1)%num:5 + facets(1)%num + facets(2)%num) = &
         [(facets(2)%dim(i), i=1, facets(2)%num)]
    allocate(output%adj_ja(2, 4 * ( facets(1)%num + facets(2)%num )))
    ! Edges (1D faces)
    do i = 1, facets(1)%num
       output%adj_ja(1,(i-1)*4 + 1) = facets(1)%orig_bound(1,1,i)
       output%adj_ja(2,(i-1)*4 + 1) = facets(1)%dest_bound(1,1,i)
       output%adj_ja(1,(i-1)*4 + 2) = facets(1)%orig_bound(2,1,i)
       output%adj_ja(2,(i-1)*4 + 2) = facets(1)%dest_bound(2,1,i)
       output%adj_ja(1,(i-1)*4 + 3) = facets(1)%orig_bound(1,2,i)
       output%adj_ja(2,(i-1)*4 + 3) = facets(1)%dest_bound(1,2,i)
       output%adj_ja(1,(i-1)*4 + 4) = facets(1)%orig_bound(2,2,i)
       output%adj_ja(2,(i-1)*4 + 4) = facets(1)%dest_bound(2,2,i)
    end do
    idx_shift = facets(1)%num * 4
    ! Corners (2D edges)
    do i = 1, facets(2)%num
       output%adj_ja(1,(i-1)*4 + 1 + idx_shift) = facets(2)%orig_bound(1,1,i)
       output%adj_ja(2,(i-1)*4 + 1 + idx_shift) = facets(2)%dest_bound(1,1,i)
       output%adj_ja(1,(i-1)*4 + 2 + idx_shift) = facets(2)%orig_bound(2,1,i)
       output%adj_ja(2,(i-1)*4 + 2 + idx_shift) = facets(2)%dest_bound(2,1,i)
       output%adj_ja(1,(i-1)*4 + 3 + idx_shift) = facets(2)%orig_bound(1,2,i)
       output%adj_ja(2,(i-1)*4 + 3 + idx_shift) = facets(2)%dest_bound(1,2,i)
       output%adj_ja(1,(i-1)*4 + 4 + idx_shift) = facets(2)%orig_bound(2,2,i)
       output%adj_ja(2,(i-1)*4 + 4 + idx_shift) = facets(2)%dest_bound(2,2,i)
    end do

    ! Initialise with zero
    output%val = 0._real32

    ! Copy input into the correct location in output
    do concurrent( &
         s = 1:output_shape(4), &
         m = 1:num_channels, &
         j = 1:input_size_w, &
         i = 1:input_size_h)
       idx_in = i + (j-1) * input_size_h + (m-1) * input_size_h * input_size_w
       idx_out = (i + pad_size(1)) + (j + pad_size(2) - 1) * output_size_h + &
            (m-1) * output_size_h * output_size_w
       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_2d( input, output )
       call fill_edge_region_2d( input, output )
    end if

    output%get_partial_left => get_partial_pad2d
    output%get_partial_left_val => get_partial_pad2d_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 pad2d