data_transformer.py

from __future__ import annotations

import abc
from collections.abc import Sequence
from dataclasses import dataclass
from enum import Enum
from types import ModuleType
from typing import no_type_check
from uuid import UUID, uuid1

import numpy as np

from ndsl.buffer import Buffer
from ndsl.halo.cuda_kernels import (
    pack_scalar_f32_kernel,
    pack_scalar_f64_kernel,
    pack_vector_f32_kernel,
    pack_vector_f64_kernel,
    unpack_scalar_f32_kernel,
    unpack_scalar_f64_kernel,
    unpack_vector_f32_kernel,
    unpack_vector_f64_kernel,
)
from ndsl.halo.rotate import rotate_scalar_data, rotate_vector_data
from ndsl.optional_imports import cupy as cp
from ndsl.quantity import Quantity, QuantityHaloSpec
from ndsl.utils import device_synchronize


# ------------------------------------------------------------------------
# Simple pool of streams to lower the driver pressure
# Use _pop/_push_stream to manipulate the pool

STREAM_POOL: list["cp.cuda.Stream"] = []


def _pop_stream() -> "cp.cuda.Stream":
    if len(STREAM_POOL) == 0:
        return cp.cuda.Stream(non_blocking=True)
    return STREAM_POOL.pop()


def _push_stream(stream: "cp.cuda.Stream") -> None:
    STREAM_POOL.append(stream)


# ------------------------------------------------------------------------
# Indices array

# Keyed cached - key is a str at the moment to go around the fact that
# a slice is not hashable. getting a string from
# Tuple(slices, rotation, shape, strides, itemsize) e.g. # noqa
# str(tuple[Any, int, tuple[int], tuple[int], int]) # noqa
INDICES_CACHE: dict[str, "cp.ndarray"] = {}


# `array_value[...] = xxx` is failing mypy because of bad inference
# of the type. We can't type ignore, because mypy also thinks that it
# no needed (but if removed, it will fail...)
@no_type_check
def _build_flatten_indices(
    key,
    shape,
    slices: tuple[slice, ...],
    dims,
    strides,
    itemsize: int,
    rotate: bool,
    rotation: int,
) -> "cp.ndarray":
    """Build an array of indexing from a slice & memory description to
    build an indexation into the "flatten" memory.

    Go from a memory layout (strides, itemsize, shape) and slices into it to a
    single array of indices. We leverage numpy iterator and calculate from
    the multi_index using memory layout the index into the original memory buffer.
    """

    # Have to go down to numpy to leverage indices calculation
    arr_indices = np.zeros(shape, dtype=np.int32, order="C")[slices]

    # Get offset from first index
    offset_dims = []
    for s in slices:
        offset_dims.append(s.start)
    offset_to_slice = sum(np.array(offset_dims) * strides) // itemsize

    # Flatten the index into an indices array
    with np.nditer(
        arr_indices,
        flags=["multi_index"],
        op_flags=["writeonly"],
        order="K",
    ) as it:
        for array_value in it:
            offset = sum(np.array(it.multi_index) * strides) // itemsize
            array_value[...] = offset_to_slice + offset

    if rotate:
        # sending data across the boundary will rotate the data
        # n_clockwise_rotations times, due to the difference in axis orientation.
        # Thus we rotate that number of times counterclockwise before sending,
        # to get the right final orientation. We apply those rotations to the
        # indices here to prepare for a straightforward copy in cu kernel
        arr_indices = rotate_scalar_data(arr_indices, dims, cp, -rotation)
    return cp.asarray(arr_indices.flatten(order="C"))


# ------------------------------------------------------------------------
# HaloDataTransformer helpers


def _slices_size(slices: tuple[slice, ...]) -> int:
    """Compute linear size from slices."""
    length = 1
    for s in slices:
        assert s.step is None
        length *= abs(s.start - s.stop)
    return length


@dataclass
class HaloExchangeSpec:
    """Memory description of the data exchanged.

    The data stored here target a single exchange, with an optional
    rotation to give prior to pack. Slices are tupled following the
    convention of one slice per dimension

    Args:
        specification: memory layout of the data
        pack_slices: indexing to pack, one slice per dimension
        pack_clockwise_rotation:  number of 90-degree rotations to perform
            before packing
        unpack_slices: indexing to unpack, one slice per dimension
    """

    specification: QuantityHaloSpec
    pack_slices: tuple[slice, ...]
    pack_clockwise_rotation: int
    unpack_slices: tuple[slice, ...]

    def __post_init__(self) -> None:
        self._id = uuid1()
        self.pack_buffer_size = _slices_size(self.pack_slices)
        self._unpack_buffer_size = _slices_size(self.unpack_slices)


class _HaloDataTransformerType(Enum):
    """Dimensionality of the data in the packed buffer."""

    UNKNOWN = 0
    SCALAR = 1
    VECTOR = 2


# ------------------------------------------------------------------------
# HaloDataTransformer classes


class HaloDataTransformer(abc.ABC):
    """Transform data to exchange in a format optimized for network communication.

    Current strategy: pack/unpack multiple nD array into/from a single buffer.
    Offers a pack and an unpack buffer to use for communicating data.

    The class is responsible for packing & unpacking, not communication.
    Order of operations:
    - get HaloDataTransformer via get() with N transformation
      with the proper halo specifications.
      At the end of get() a _compile() will be triggered, reading
      the internal buffers.
    - call async_pack(quantities) to start packing the quantities in the
      internal buffer.
    - synchronize() to make sure all operations are finished or use get_pack_buffer()
      when ready to communicate which will internally call synchronize.
    [... user should communicate the buffers...]
    - call async_unpack(quantities) to start unpacking
    - call synchronize() to finish all the unpacking operations and make sure
      the quantities passed in async_unpack have been updated.

    The class will hold onto the buffers up until deletion, where they will be
    returned to an internal buffer pool.
    """

    _pack_buffer: Buffer | None
    _unpack_buffer: Buffer | None

    _infos_x: tuple[HaloExchangeSpec, ...]
    _infos_y: tuple[HaloExchangeSpec, ...]

    def __init__(
        self,
        np_module: ModuleType,
        exchange_descriptors_x: Sequence[HaloExchangeSpec],
        exchange_descriptors_y: Sequence[HaloExchangeSpec] | None = None,
    ) -> None:
        """
        Args:
            np_module: numpy-like module for allocation
            exchange_descriptors_x: list of memory information describing an exchange.
                Used for scalar data and the x-component of vectors.
            exchange_descriptors_y: list of memory information describing an exchange.
                Optional, used for the y-component of vectors only. If `none` the
                data will packed as a scalar.
        """
        self._type = (
            _HaloDataTransformerType.SCALAR
            if exchange_descriptors_y is None
            else _HaloDataTransformerType.VECTOR
        )
        if exchange_descriptors_y is not None and len(exchange_descriptors_y) != len(
            exchange_descriptors_x
        ):
            raise RuntimeError(
                "Vector halo exchange must have same exchange data for X and Y"
            )
        self._np_module = np_module
        self._infos_x = tuple(exchange_descriptors_x)
        self._infos_y = (
            tuple(exchange_descriptors_y)
            if exchange_descriptors_y is not None
            else tuple()
        )
        self._pack_buffer = None
        self._unpack_buffer = None
        self._compile()

    def finalize(self) -> None:
        """Deletion routine, making sure all buffers were inserted back into cache."""
        # Synchronize all work
        self.synchronize()

        # Push the buffers back in the cache
        if self._pack_buffer is not None:
            Buffer.push_to_cache(self._pack_buffer)
        self._pack_buffer = None

        if self._unpack_buffer is not None:
            Buffer.push_to_cache(self._unpack_buffer)
        self._unpack_buffer = None

    @staticmethod
    def get(
        np_module: ModuleType,
        exchange_descriptors_x: Sequence[HaloExchangeSpec],
        exchange_descriptors_y: Sequence[HaloExchangeSpec] | None = None,
    ) -> HaloDataTransformer:
        """Construct a module from a numpy-like module.

        Args:
            np_module: numpy-like module to determine child transformer type.
            exchange_descriptors_x: list of memory information describing an exchange.
                Used for scalar data and the x-component of vectors.
            exchange_descriptors_y: list of memory information describing an exchange.
                Optional, used for the y-component of vectors only. If `none` the data
                will packed as a scalar.

        Returns:
            an initialized packed buffer.
        """
        if len(exchange_descriptors_x) == 0:
            raise RuntimeError("Attempting to init an empty halo exchange")

        dtype = exchange_descriptors_x[0].specification.dtype
        for desc in exchange_descriptors_x:
            if dtype != desc.specification.dtype:
                raise NotImplementedError("Halo exchange process mixed precision")
        if exchange_descriptors_y:
            for desc in exchange_descriptors_y:
                if dtype != desc.specification.dtype:
                    raise NotImplementedError("Halo exchange process mixed precision")

        if np_module is np:
            return HaloDataTransformerCPU(
                np,
                exchange_descriptors_x,
                exchange_descriptors_y=exchange_descriptors_y,
            )
        elif np_module is cp:
            return HaloDataTransformerGPU(
                cp,
                exchange_descriptors_x,
                exchange_descriptors_y=exchange_descriptors_y,
            )

        raise NotImplementedError(
            f"Quantity module {np_module} has no HaloDataTransformer implemented"
        )

    def get_unpack_buffer(self) -> Buffer:
        """Retrieve unpack buffer.

        Synchronizes operations.
        """
        if self._unpack_buffer is None:
            raise RuntimeError("Recv buffer can't be retrieved before allocate()")
        self.synchronize()
        return self._unpack_buffer

    def get_pack_buffer(self) -> Buffer:
        """Retrieve pack buffer.

        Synchronizes operations.
        """
        if self._pack_buffer is None:
            raise RuntimeError("Send buffer can't be retrieved before allocate()")
        self.synchronize()
        return self._pack_buffer

    def _compile(self) -> None:
        """Allocate contiguous memory buffers from description queued."""

        # Compute required size
        buffer_size = 0
        dtype = np.float32  # default that will be overriden or not used
        for edge_x in self._infos_x:
            buffer_size += edge_x.pack_buffer_size
            dtype = edge_x.specification.dtype
        if self._type is _HaloDataTransformerType.VECTOR:
            for edge_y in self._infos_y:
                buffer_size += edge_y.pack_buffer_size

        # Retrieve two properly sized buffers
        self._pack_buffer = Buffer.pop_from_cache(
            self._np_module.zeros,
            (buffer_size,),
            dtype,
        )
        self._unpack_buffer = Buffer.pop_from_cache(
            self._np_module.zeros,
            (buffer_size,),
            dtype,
        )

    def ready(self) -> bool:
        """Check if the buffers are ready for communication."""
        return self._pack_buffer is not None and self._unpack_buffer is not None

    @abc.abstractmethod
    def async_pack(
        self,
        quantities_x: list[Quantity],
        quantities_y: list[Quantity] | None = None,
    ) -> None:
        """Pack all given quantities into a single send Buffer.

        Does not guarantee the buffer returned by `get_unpack_buffer` has
        received data, doing so requires calling `synchronize`.
        Reaching for the buffer via get_pack_buffer() will call synchronize().

        Args:
            quantities_x: scalar or vector x-component quantities to pack,
                if one is vector they must all be vector

            quantities_y: if quantities are vector, the y-component
                quantities.
        """
        pass

    @abc.abstractmethod
    def async_unpack(
        self,
        quantities_x: Sequence[Quantity],
        quantities_y: Sequence[Quantity] | None = None,
    ) -> None:
        """Unpack the buffer into destination quantities.

        Does not guarantee the buffer returned by `get_unpack_buffer` has
        received data, doing so requires calling `synchronize`.
        Reaching for the buffer via get_unpack_buffer() will call synchronize().

        Args:
            quantities_x: scalar or vector x-component quantities to be unpacked into,
                if one is vector they must all be vector
            quantities_y: if quantities are vector, the y-component
                quantities.
        """
        pass

    @abc.abstractmethod
    def synchronize(self) -> None:
        """Synchronize all operations.

        Guarantees all memory is now safe to access.
        """
        pass


class HaloDataTransformerCPU(HaloDataTransformer):
    """Pack/unpack data in a single buffer using numpy flattening & slicing.

    Default behavior, could be done with any numpy-like library.
    """

    def synchronize(self) -> None:
        if self._pack_buffer is not None:
            self._pack_buffer.finalize_memory_transfer()
        if self._unpack_buffer is not None:
            self._unpack_buffer.finalize_memory_transfer()

    def async_pack(
        self,
        quantities_x: Sequence[Quantity],
        quantities_y: Sequence[Quantity] | None = None,
    ) -> None:
        # Unpack per type
        if self._type == _HaloDataTransformerType.SCALAR:
            self._pack_scalar(quantities_x)
        elif self._type == _HaloDataTransformerType.VECTOR:
            assert quantities_y is not None
            self._pack_vector(quantities_x, quantities_y)
        else:
            raise RuntimeError(f"Unimplemented {self._type} pack")

        assert isinstance(self._pack_buffer, Buffer)  # e.g. allocate happened

    def _pack_scalar(self, quantities: Sequence[Quantity]) -> None:
        if __debug__:
            if len(quantities) != len(self._infos_x):
                raise RuntimeError(
                    f"Quantities count ({len(quantities)}"
                    f" is different that edges count {len(self._infos_x)}"
                )
            # TODO Per quantity check

        assert isinstance(self._pack_buffer, Buffer)  # e.g. allocate happened
        offset = 0
        for quantity, info_x in zip(quantities, self._infos_x):
            data_size = _slices_size(info_x.pack_slices)
            # sending data across the boundary will rotate the data
            # n_clockwise_rotations times, due to the difference in axis orientation.\
            # Thus we rotate that number of times counterclockwise before sending,
            # to get the right final orientation
            source_view = rotate_scalar_data(
                quantity.data[info_x.pack_slices],
                quantity.dims,
                quantity.np,
                -info_x.pack_clockwise_rotation,
            )
            self._pack_buffer.assign_from(
                source_view.flatten(),
                buffer_slice=np.index_exp[offset : offset + data_size],
            )
            offset += data_size

    def _pack_vector(
        self, quantities_x: Sequence[Quantity], quantities_y: Sequence[Quantity]
    ) -> None:
        if __debug__:
            if len(quantities_x) != len(self._infos_x) and len(quantities_y) != len(
                self._infos_y
            ):
                raise RuntimeError(
                    f"Quantities count (x: {len(quantities_x)}, y: {len(quantities_y)})"
                    " is different that specifications count "
                    f"(x: {len(self._infos_x)}, y: {len(self._infos_y)}"
                )
            # TODO Per quantity check

        assert isinstance(self._pack_buffer, Buffer)  # e.g. allocate happened
        assert len(quantities_y) == len(quantities_x)
        assert len(self._infos_x) == len(self._infos_y)
        offset = 0
        for (
            quantity_x,
            quantity_y,
            info_x,
            info_y,
        ) in zip(quantities_x, quantities_y, self._infos_x, self._infos_y):
            # sending data across the boundary will rotate the data
            # n_clockwise_rotations times, due to the difference in axis orientation
            # Thus we rotate that number of times counterclockwise before sending,
            # to get the right final orientation
            x_view, y_view = rotate_vector_data(
                quantity_x.data[info_x.pack_slices],
                quantity_y.data[info_y.pack_slices],
                -info_x.pack_clockwise_rotation,
                quantity_x.dims,
                quantity_x.np,
            )

            # Pack X/Y data slices in the buffer
            self._pack_buffer.assign_from(
                x_view.flatten(),
                buffer_slice=np.index_exp[offset : offset + x_view.size],
            )
            offset += x_view.size
            self._pack_buffer.assign_from(
                y_view.flatten(),
                buffer_slice=np.index_exp[offset : offset + y_view.size],
            )
            offset += y_view.size

    def async_unpack(
        self,
        quantities_x: Sequence[Quantity],
        quantities_y: Sequence[Quantity] | None = None,
    ) -> None:
        # Unpack per type
        if self._type == _HaloDataTransformerType.SCALAR:
            self._unpack_scalar(quantities_x)
        elif self._type == _HaloDataTransformerType.VECTOR:
            assert quantities_y is not None
            self._unpack_vector(quantities_x, quantities_y)
        else:
            raise RuntimeError(f"Unimplemented {self._type} unpack")

        assert isinstance(self._unpack_buffer, Buffer)  # e.g. allocate happened

    def _unpack_scalar(self, quantities: Sequence[Quantity]) -> None:
        if __debug__:
            if len(quantities) != len(self._infos_x):
                raise RuntimeError(
                    f"Quantities count ({len(quantities)}"
                    f" is different that specifications count {len(self._infos_x)}"
                )
            # TODO Per quantity check

        assert isinstance(self._unpack_buffer, Buffer)  # e.g. allocate happened
        offset = 0
        for quantity, info_x in zip(quantities, self._infos_x):
            quantity_view = quantity.data[info_x.unpack_slices]
            data_size = _slices_size(info_x.unpack_slices)
            self._unpack_buffer.assign_to(
                quantity_view,
                buffer_slice=np.index_exp[offset : offset + data_size],
                buffer_reshape=quantity_view.shape,
            )
            offset += data_size

    def _unpack_vector(
        self, quantities_x: Sequence[Quantity], quantities_y: Sequence[Quantity]
    ) -> None:
        if __debug__:
            if len(quantities_x) != len(self._infos_x) and len(quantities_y) != len(
                self._infos_y
            ):
                raise RuntimeError(
                    f"Quantities count (x: {len(quantities_x)}, y: {len(quantities_y)})"
                    " is different that specifications count "
                    f"(x: {len(self._infos_x)}, y: {len(self._infos_y)})"
                )
            # TODO Per quantity check

        assert isinstance(self._unpack_buffer, Buffer)  # e.g. allocate happened
        offset = 0
        for quantity_x, quantity_y, info_x, info_y in zip(
            quantities_x, quantities_y, self._infos_x, self._infos_y
        ):
            quantity_view = quantity_x.data[info_x.unpack_slices]
            data_size = _slices_size(info_x.unpack_slices)
            self._unpack_buffer.assign_to(
                quantity_view,
                buffer_slice=np.index_exp[offset : offset + data_size],
                buffer_reshape=quantity_view.shape,
            )
            offset += data_size
            quantity_view = quantity_y.data[info_y.unpack_slices]
            data_size = _slices_size(info_y.unpack_slices)
            self._unpack_buffer.assign_to(
                quantity_view,
                buffer_slice=np.index_exp[offset : offset + data_size],
                buffer_reshape=quantity_view.shape,
            )
            offset += data_size


class HaloDataTransformerGPU(HaloDataTransformer):
    """Pack/unpack data in a single buffer using CUDA Kernels.

    In order to efficiently pack/unpack on the GPU to a single GPU buffer
    we use streamed (e.g. async) kernels per quantity per edge to send. The
    kernels are store in `cuda_kernels.py`, they both follow the same simple pattern
    by reading the indices to the device memory of the data to pack/unpack.
    `_flatten_indices` is the routine that take the layout of the memory and
    the slice and compute an array of index into the original memory.
    """

    # Temporary "safe" code path
    #  _CODE_PATH_DEVICE_WIDE_SYNC: turns off streaming and issue a single
    #   device wide synchronization call instead
    _CODE_PATH_DEVICE_WIDE_SYNC = False

    @dataclass
    class _CuKernelArgs:
        """All arguments required for the CUDA kernels."""

        stream: cp.cuda.Stream
        x_send_indices: cp.ndarray
        x_recv_indices: cp.ndarray
        y_send_indices: cp.ndarray | None
        y_recv_indices: cp.ndarray | None

    def __init__(
        self,
        np_module: ModuleType,
        exchange_descriptors_x: Sequence[HaloExchangeSpec],
        exchange_descriptors_y: Sequence[HaloExchangeSpec] | None = None,
    ) -> None:
        self._cu_kernel_args: dict[UUID, HaloDataTransformerGPU._CuKernelArgs] = {}
        super().__init__(
            np_module,
            exchange_descriptors_x,
            exchange_descriptors_y=exchange_descriptors_y,
        )

    def _flatten_indices(
        self,
        exchange_data: HaloExchangeSpec,
        slices: tuple[slice, ...],
        rotate: bool,
    ) -> "cp.ndarray":
        """Extract a flat array of indices from the memory layout and the slice.

        Also take care of rotating the indices to account for axis orientation.
        """
        key = str(
            (
                slices,
                exchange_data.pack_clockwise_rotation,
                exchange_data.specification.shape,
                exchange_data.specification.strides,
                exchange_data.specification.itemsize,
            )
        )

        # We use a lazy caching mechanism here because in our use case
        # (halo exchange) there is a limited set of index patterns but a
        # large number of exchanges.
        if key not in INDICES_CACHE.keys():
            INDICES_CACHE[key] = _build_flatten_indices(
                key,
                exchange_data.specification.shape,
                slices,
                exchange_data.specification.dims,
                exchange_data.specification.strides,
                exchange_data.specification.itemsize,
                rotate,
                exchange_data.pack_clockwise_rotation,
            )

        # We don't return a copy since the indices are read-only in the algorithm
        return INDICES_CACHE[key]

    def _compile(self) -> None:
        # Super to get buffer allocation
        super()._compile()
        # Allocate the streams & build the indices arrays
        if self._type == _HaloDataTransformerType.SCALAR:
            for info_x in self._infos_x:
                self._cu_kernel_args[info_x._id] = HaloDataTransformerGPU._CuKernelArgs(
                    stream=_pop_stream(),
                    x_send_indices=self._flatten_indices(
                        info_x,
                        info_x.pack_slices,
                        True,
                    ),
                    x_recv_indices=self._flatten_indices(
                        info_x,
                        info_x.unpack_slices,
                        False,
                    ),
                    y_send_indices=None,
                    y_recv_indices=None,
                )
        else:
            assert self._type == _HaloDataTransformerType.VECTOR
            for info_x, info_y in zip(self._infos_x, self._infos_y):
                self._cu_kernel_args[info_x._id] = HaloDataTransformerGPU._CuKernelArgs(
                    stream=_pop_stream(),
                    x_send_indices=self._flatten_indices(
                        info_x,
                        info_x.pack_slices,
                        True,
                    ),
                    x_recv_indices=self._flatten_indices(
                        info_x,
                        info_x.unpack_slices,
                        False,
                    ),
                    y_send_indices=self._flatten_indices(
                        info_y,
                        info_y.pack_slices,
                        True,
                    ),
                    y_recv_indices=self._flatten_indices(
                        info_y,
                        info_y.unpack_slices,
                        False,
                    ),
                )

    def synchronize(self) -> None:
        if self._CODE_PATH_DEVICE_WIDE_SYNC:
            self._safe_synchronize()
        else:
            self._streamed_synchronize()

    def _streamed_synchronize(self) -> None:
        for cu_kernel in self._cu_kernel_args.values():
            cu_kernel.stream.synchronize()

    def _safe_synchronize(self) -> None:
        device_synchronize()

    def _get_stream(self, stream: "cp.cuda.stream") -> "cp.cuda.stream":
        if self._CODE_PATH_DEVICE_WIDE_SYNC:
            return cp.cuda.Stream.null
        else:
            return stream

    def async_pack(
        self,
        quantities_x: list[Quantity],
        quantities_y: list[Quantity] | None = None,
    ) -> None:
        """Pack the quantities into a single buffer via streamed cuda kernels

        Writes into self._pack_buffer using self._x_infos and self._y_infos
        to read the offsets and sizes per quantity.

        Args:
            quantities_x: list of quantities to pack. Must fit the specifications given
                at init time.
            quantities_y: Same as above but optional, used only for vector transfer.
        """

        # Unpack per type
        if self._type == _HaloDataTransformerType.SCALAR:
            self._opt_pack_scalar(quantities_x)
        elif self._type == _HaloDataTransformerType.VECTOR:
            assert quantities_y is not None
            self._opt_pack_vector(quantities_x, quantities_y)
        else:
            raise RuntimeError(f"Unimplemented {self._type} pack")

    def _opt_pack_scalar(self, quantities: list[Quantity]) -> None:
        """Specialized packing for scalar. See async_pack docs for usage."""
        if __debug__:
            if len(quantities) != len(self._infos_x):
                raise RuntimeError(
                    f"Quantities count ({len(quantities)}"
                    f" is different that specifications count {len(self._infos_x)}"
                )
            # TODO Per quantity check

        assert isinstance(self._pack_buffer, Buffer)  # e.g. allocate happened
        offset = 0
        for info_x, quantity in zip(self._infos_x, quantities):
            cu_kernel_args = self._cu_kernel_args[info_x._id]

            # Use private stream
            with self._get_stream(cu_kernel_args.stream):
                # Launch kernel
                blocks = 128
                grid_x = (info_x.pack_buffer_size // blocks) + 1
                # Pick a kernel looking at the precision set
                pack_kernel = None
                if info_x.specification.dtype == np.float32:
                    pack_kernel = pack_scalar_f32_kernel
                elif (
                    info_x.specification.dtype == np.float64
                    or info_x.specification.dtype == float
                ):
                    pack_kernel = pack_scalar_f64_kernel
                else:
                    RuntimeError(
                        "Halo exchange pack kernel for precision "
                        f" {info_x.specification.dtype} isn't implemented."
                    )

                # Check compile hasn't failed silently if this is the first
                # call to the kernel
                if pack_kernel is None:
                    RuntimeError("CUDA nvrtc failed")
                else:
                    pack_kernel(
                        (grid_x,),
                        (blocks,),
                        (
                            quantity.data[:],  # source_array
                            cu_kernel_args.x_send_indices,  # indices
                            info_x.pack_buffer_size,  # nIndex
                            offset,
                            self._pack_buffer.array,
                        ),
                    )

                # Next transformer offset into send buffer
                offset += info_x.pack_buffer_size

    def _opt_pack_vector(
        self, quantities_x: list[Quantity], quantities_y: list[Quantity]
    ) -> None:
        """Specialized packing for vectors. See async_pack docs for usage."""
        if __debug__:
            if len(quantities_x) != len(self._infos_x) and len(quantities_y) != len(
                self._infos_y
            ):
                raise RuntimeError(
                    f"Quantities count (x: {len(quantities_x)}, y: {len(quantities_y)}"
                    " is different that specifications count "
                    f"(x: {len(self._infos_x)}, y: {len(self._infos_y)}"
                )
            # TODO Per quantity check
        assert isinstance(self._pack_buffer, Buffer)  # e.g. allocate happened
        assert len(self._infos_x) == len(self._infos_y)
        assert len(quantities_x) == len(quantities_y)
        offset = 0
        for (
            quantity_x,
            quantity_y,
            info_x,
            info_y,
        ) in zip(quantities_x, quantities_y, self._infos_x, self._infos_y):
            cu_kernel_args = self._cu_kernel_args[info_x._id]

            # Use private stream
            with self._get_stream(cu_kernel_args.stream):
                # Buffer sizes
                transformer_size = info_x.pack_buffer_size + info_y.pack_buffer_size

                # Launch kernel
                blocks = 128
                grid_x = (transformer_size // blocks) + 1
                # Pick a kernel looking at the precision set
                pack_kernel = None
                if info_x.specification.dtype == np.float32:
                    pack_kernel = pack_vector_f32_kernel
                elif (
                    info_x.specification.dtype == np.float64
                    or info_x.specification.dtype == float
                ):
                    pack_kernel = pack_vector_f64_kernel
                else:
                    RuntimeError(
                        "Halo exchange pack kernel for precision "
                        f"{info_x.specification.dtype} isn't implemented."
                    )

                # Check compile hasn't failed silently if this is the first
                # call to the kernel
                if pack_kernel is None:
                    RuntimeError("CUDA nvrtc failed")
                else:
                    pack_kernel(
                        (grid_x,),
                        (blocks,),
                        (
                            quantity_x.data[:],  # source_array_x
                            quantity_y.data[:],  # source_array_y
                            cu_kernel_args.x_send_indices,  # indices_x
                            cu_kernel_args.y_send_indices,  # indices_y
                            info_x.pack_buffer_size,  # nIndex_x
                            info_y.pack_buffer_size,  # nIndex_y
                            offset,
                            (-info_x.pack_clockwise_rotation) % 4,  # rotation
                            self._pack_buffer.array,
                        ),
                    )

                # Next transformer offset into send buffer
                offset += transformer_size

    def async_unpack(
        self,
        quantities_x: Sequence[Quantity],
        quantities_y: Sequence[Quantity] | None = None,
    ) -> None:
        """Unpack the quantities from a single buffer via streamed cuda kernels

        Reads from self._unpack_buffer using self._x_infos and self._y_infos
        to read the offsets and sizes per quantity.

        Args:
            quantities_x: list of quantities to unpack. Must fit
                the specifications given at init time.
            quantities_y: Same as above but optional, used only for vector transfer.
        """
        # Unpack per type
        if self._type == _HaloDataTransformerType.SCALAR:
            self._opt_unpack_scalar(quantities_x)
        elif self._type == _HaloDataTransformerType.VECTOR:
            assert quantities_y is not None
            self._opt_unpack_vector(quantities_x, quantities_y)
        else:
            raise RuntimeError(f"Unimplemented {self._type} unpack")

    def _opt_unpack_scalar(self, quantities: Sequence[Quantity]) -> None:
        """Specialized unpacking for scalars. See async_unpack docs for usage."""
        if __debug__:
            if len(quantities) != len(self._infos_x):
                raise RuntimeError(
                    f"Quantities count ({len(quantities)})"
                    f" is different that specifications count ({len(self._infos_x)})"
                )
            # TODO Per quantity check
        assert isinstance(self._unpack_buffer, Buffer)  # e.g. allocate happened
        offset = 0
        for quantity, info_x in zip(quantities, self._infos_x):
            cu_kernel_args = self._cu_kernel_args[info_x._id]

            # Use private stream
            with self._get_stream(cu_kernel_args.stream):
                # Launch kernel
                blocks = 128
                grid_x = (info_x._unpack_buffer_size // blocks) + 1
                # Pick a kernel looking at the precision set
                unpack_kernel = None
                if info_x.specification.dtype == np.float32:
                    unpack_kernel = unpack_scalar_f32_kernel
                elif (
                    info_x.specification.dtype == np.float64
                    or info_x.specification.dtype == float
                ):
                    unpack_kernel = unpack_scalar_f64_kernel
                else:
                    RuntimeError(
                        "Halo exchange pack kernel for precision "
                        f"{info_x.specification.dtype} isn't implemented."
                    )

                # Check compile hasn't failed silently if this is the first
                # call to the kernel
                if unpack_kernel is None:
                    RuntimeError("CUDA nvrtc failed")
                else:
                    unpack_kernel(
                        (grid_x,),
                        (blocks,),
                        (
                            self._unpack_buffer.array,  # source_buffer
                            cu_kernel_args.x_recv_indices,  # indices
                            info_x._unpack_buffer_size,  # nIndex
                            offset,
                            quantity.data[:],  # destination_array
                        ),
                    )

                # Next transformer offset into recv buffer
                offset += info_x._unpack_buffer_size

    def _opt_unpack_vector(
        self, quantities_x: Sequence[Quantity], quantities_y: Sequence[Quantity]
    ) -> None:
        """Specialized unpacking for vectors. See async_unpack docs for usage."""
        if __debug__:
            if len(quantities_x) != len(self._infos_x) and len(quantities_y) != len(
                self._infos_y
            ):
                raise RuntimeError(
                    f"Quantities count (x: {len(quantities_x)}, y: {len(quantities_y)}"
                    " is different that specifications count "
                    f"(x: {len(self._infos_x)}, y: {len(self._infos_y)}"
                )
            # TODO Per quantity check
        assert isinstance(self._unpack_buffer, Buffer)  # e.g. allocate happened
        assert len(self._infos_x) == len(self._infos_y)
        assert len(quantities_x) == len(quantities_y)
        offset = 0
        for (
            quantity_x,
            quantity_y,
            info_x,
            info_y,
        ) in zip(quantities_x, quantities_y, self._infos_x, self._infos_y):
            cu_kernel_args = self._cu_kernel_args[info_x._id]

            # Use private stream
            with self._get_stream(cu_kernel_args.stream):
                # Buffer sizes
                edge_size = info_x._unpack_buffer_size + info_y._unpack_buffer_size

                # Launch kernel
                blocks = 128
                grid_x = (edge_size // blocks) + 1
                # Pick a kernel looking at the precision set
                unpack_kernel = None
                if info_x.specification.dtype == np.float32:
                    unpack_kernel = unpack_vector_f32_kernel
                elif (
                    info_x.specification.dtype == np.float64
                    or info_x.specification.dtype == float
                ):
                    unpack_kernel = unpack_vector_f64_kernel
                else:
                    RuntimeError(
                        "Halo exchange pack kernel for precision "
                        f"{info_x.specification.dtype} isn't implemented."
                    )

                # Check compile hasn't failed silently if this is the first
                # call to the kernel
                if unpack_kernel is None:
                    RuntimeError("CUDA nvrtc failed")
                else:
                    unpack_kernel(
                        (grid_x,),
                        (blocks,),
                        (
                            self._unpack_buffer.array,
                            cu_kernel_args.x_recv_indices,  # indices_x
                            cu_kernel_args.y_recv_indices,  # indices_y
                            info_x._unpack_buffer_size,  # nIndex_x
                            info_y._unpack_buffer_size,  # nIndex_y
                            offset,
                            quantity_x.data[:],  # destination_array_x
                            quantity_y.data[:],  # destination_array_y
                        ),
                    )

                # Next transformer offset into send buffer
                offset += edge_size

    def finalize(self) -> None:
        super().finalize()
        # Push the streams back in the pool
        for cu_info in self._cu_kernel_args.values():
            _push_stream(cu_info.stream)