halo_exchange_end Subroutine

   subroutine halo_exchange_end(ha, h, hu, hv, b_fld)
      !! Wait for MPI to complete and unpack received ghost cells
      !!
      !! Buffers are module-level arrays (ha_buf_*) — see
      !! halo_exchange_begin for explanation.
      type(halo_async_t), intent(inout) :: ha
      real(wp), intent(inout) :: h(:, :), hu(:, :), hv(:, :), b_fld(:, :)

      type(MPI_Status) :: stats(MAX_REQS)
      integer :: i, j, k, base
      integer :: ng, nxl, nyl, nxt, nyt

      ng = ha%nghost
      nxl = ha%nx_local
      nyl = ha%ny_local
      nxt = ha%nx_total
      nyt = ha%ny_total

      if (ha%nreq > 0) then
         call waitall(ha%reqs(1:ha%nreq), stats(1:ha%nreq))
      end if

      ! --- Unpack all 4 fields from combined buffers on device ---
      if (.not. ha%decomp%has_west) then
         !$acc parallel loop collapse(2)
         do j = 1, nyt
            do k = 1, ng
               base = (j - 1)*ng + k
               h(k, j) = ha_buf_recv_west(base)
               hu(k, j) = ha_buf_recv_west(base + ng*nyt)
               hv(k, j) = ha_buf_recv_west(base + 2*ng*nyt)
               b_fld(k, j) = ha_buf_recv_west(base + 3*ng*nyt)
            end do
         end do
      end if

      if (.not. ha%decomp%has_east) then
         !$acc parallel loop collapse(2)
         do j = 1, nyt
            do k = 1, ng
               base = (j - 1)*ng + k
               h(ng + nxl + k, j) = ha_buf_recv_east(base)
               hu(ng + nxl + k, j) = ha_buf_recv_east(base + ng*nyt)
               hv(ng + nxl + k, j) = ha_buf_recv_east(base + 2*ng*nyt)
               b_fld(ng + nxl + k, j) = ha_buf_recv_east(base + 3*ng*nyt)
            end do
         end do
      end if

      if (.not. ha%decomp%has_south) then
         !$acc parallel loop collapse(2)
         do k = 1, ng
            do i = 1, nxt
               base = (k - 1)*nxt + i
               h(i, k) = ha_buf_recv_south(base)
               hu(i, k) = ha_buf_recv_south(base + nxt*ng)
               hv(i, k) = ha_buf_recv_south(base + 2*nxt*ng)
               b_fld(i, k) = ha_buf_recv_south(base + 3*nxt*ng)
            end do
         end do
      end if

      if (.not. ha%decomp%has_north) then
         !$acc parallel loop collapse(2)
         do k = 1, ng
            do i = 1, nxt
               base = (k - 1)*nxt + i
               h(i, ng + nyl + k) = ha_buf_recv_north(base)
               hu(i, ng + nyl + k) = ha_buf_recv_north(base + nxt*ng)
               hv(i, ng + nyl + k) = ha_buf_recv_north(base + 2*nxt*ng)
               b_fld(i, ng + nyl + k) = ha_buf_recv_north(base + 3*nxt*ng)
            end do
         end do
      end if

   end subroutine halo_exchange_end