Accumulate edge gradients for 1D padding - raw array version
| Type | Intent | Optional | 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 |
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