pure subroutine fill_ke_total_impl(h_layer, u_face, v_face, buf)
! assumed-shape-ok: diag fill — fires once per output frame (cadence-bounded);
! face-sized arrays have nx+1/ny+1 dims; size() min-clips at call.
real(wp), intent(in) :: h_layer(:, :, :)
real(wp), intent(in) :: u_face(:, :, :), v_face(:, :, :) ! assumed-shape-ok: diag fill — cadence-bounded; face-sized dims
real(wp), intent(inout) :: buf(:, :, :) ! assumed-shape-ok: diag fill — cadence-bounded
integer :: i, j, k, nx, ny, nz
real(wp) :: uc, vc, col_ke
nx = min(size(buf, 1), size(u_face, 1) - 1, size(v_face, 1))
ny = min(size(buf, 2), size(u_face, 2), size(v_face, 2) - 1)
nz = min(size(u_face, 3), size(v_face, 3), size(h_layer, 3))
! Per-cell column sum: outer DC over (j, i), inner serial over k.
do concurrent(j=1:ny, i=1:nx) &
local(uc, vc, col_ke, k)
col_ke = 0.0_wp
do k = 1, nz
uc = 0.5_wp*(u_face(i, j, k) + u_face(i + 1, j, k))
vc = 0.5_wp*(v_face(i, j, k) + v_face(i, j + 1, k))
col_ke = col_ke + 0.5_wp*h_layer(i, j, k)*(uc*uc + vc*vc)
end do
buf(i, j, 1) = col_ke
end do
end subroutine fill_ke_total_impl