accumulate_edge_gradients_1d_val Subroutine

pure subroutine accumulate_edge_gradients_1d_val(upstream_grad, output, input_shape, indices, adj_ja)

Accumulate edge gradients for 1D padding - raw array version

Arguments

Type IntentOptional Attributes Name
real(kind=real32), intent(in), dimension(:,:) :: upstream_grad
real(kind=real32), intent(inout), dimension(:,:) :: output
integer, intent(in), dimension(3) :: input_shape
integer, intent(in), dimension(:) :: indices
integer, intent(in), dimension(:,:) :: adj_ja

Source Code

  pure subroutine accumulate_edge_gradients_1d_val(upstream_grad, output, &
       input_shape, indices, adj_ja)
    !! Accumulate edge gradients for 1D padding - raw array version
    implicit none

    ! Arguments
    real(real32), dimension(:,:), intent(in) :: upstream_grad
    real(real32), dimension(:,:), intent(inout) :: output
    integer, dimension(3), intent(in) :: input_shape
    integer, dimension(:), intent(in) :: indices
    integer, dimension(:,:), intent(in) :: adj_ja

    ! Local variables
    integer :: i, m, s, f
    integer :: idx_in, idx_out
    integer :: input_size, output_size
    integer :: num_facets
    real(real32) :: grad_sum

    input_size = input_shape(1)
    output_size = input_size + 2 * indices(2)
    num_facets = indices(3)

    if(num_facets .eq. 0) return

    select case(indices(1))
    case(3, 4) ! circular or reflection
       do f = 1, num_facets
          do s = 1, input_shape(3)
             do m = 1, input_shape(2)
                do i = adj_ja(2,(f-1)*2 + 1), adj_ja(2,(f-1)*2 + 2)
                   idx_out = i + (m-1) * output_size
                   if(indices(1) .eq. 3)then  ! circular
                      idx_in = adj_ja(1,(f-1)*2 + 1) + &
                           (i - adj_ja(2,(f-1)*2 + 1)) + (m-1) * input_size
                   else  ! reflection
                      idx_in = adj_ja(1,(f-1)*2 + 1) - &
                           (i - adj_ja(2,(f-1)*2 + 1)) + (m-1) * input_size
                   end if
                   output(idx_in, s) = output(idx_in, s) + &
                        upstream_grad(idx_out, s)
                end do
             end do
          end do
       end do
    case(5) ! replication
       do f = 1, num_facets
          do s = 1, input_shape(3)
             do m = 1, input_shape(2)
                grad_sum = 0._real32
                do i = adj_ja(2,(f-1)*2 + 1), adj_ja(2,(f-1)*2 + 2)
                   idx_out = i + (m-1) * output_size
                   grad_sum = grad_sum + upstream_grad(idx_out, s)
                end do
                idx_in = adj_ja(1,(f-1)*2 + 1) + (m-1) * input_size
                output(idx_in, s) = output(idx_in, s) + grad_sum
             end do
          end do
       end do
    end select

  end subroutine accumulate_edge_gradients_1d_val