Storage Allocation and Interfaces

# Storage Allocation and Interfaces ###### tags: `storage` GT4Py does not provide its own data container class, but supports established python standards for exposing N-dimensional buffers. There is a minimalistic interface allowing to specify the correspondence of buffer dimensions to the semantic dimensions assumed in the stencil. This correspondence needs not be specified since the stencils specify a default ordering. GT4Py provides utilities to allocate buffers that have optimal layout and alignment for a given backend. In this document, we describe the interfaces for * supported buffer interfaces * exposing dimension labels and the behavior for default values * performance-optimal allocation ## Interfaces ### Stencils Calls #### Supported Buffer Interfaces The user is free to choose a buffer interface (or multiple) as long as it is supported by `numpy.asarray` in case a CPU backend is choosen or `cupy.asarray` for a GPU backend respectively. If multiple buffer interfaces are implemented the provided information needs to agree otherwise the behaviour is undefined. Similarly the backend is also free to choose what buffer interface to use in order to retrieve the required information (e.g. pointer, strides, etc.) as long as it is provided by `numpy.ndarray` or `cupy.ndarray`. Consequently we support the following interfaces to expose a buffer: * [`__array_interface__`](https://omz-software.com/pythonista/numpy/reference/arrays.interface.html) (CPU backends) * [`__cuda_array_interface__`](https://numba.pydata.org/numba-doc/dev/cuda/cuda_array_interface.html) (GPU backends) * [python buffer protocol](https://docs.python.org/3/c-api/buffer.html) (CPU backends)  Allowing multiple interfaces both with respect to what can be provided by the user and what can used in the backend/bindings might appear as an additional complication, but this is absorbed by leveraging numpy and cupy that provide us with additional flexibility for free. Agreement between cartesian and declarative is a non-issue as using `numpy.asarray` and `cupy.asarray` on the buffer provided by the user automatically fullfils the requirements of the backend. In other words there is a single, well-defined way to bridge between the user provided buffer and the backend. We will provide utilities `gt4py.utils.as_numpy` and `gt4py.utils.as_cupy` that are initially aliases for the respective `asarray` function, but can be extended to support more interfaces such as dlpack. GT4Py developers are advised to always use those utilities. #### Dimension Mapping The user can optionally implement a `__gt_dims__(self)` method in the object implementing any of the supported buffer interfaces. As a fallback if neither is specified the dimensions given in the annotations (by means of `gtscript.Field` for cartesian) are assumed. The returned object should be a tuple (or any sequence) of strings labeling the dimensions in index order. In cartesian GT4Py, valid strings are ``"I"``, ``"J"``, ``"K"`` as well as decimal string representations of integer numbers to denote data dimensions. We will provide the utility `gt4py.utils.get_dims(storage, annotation)` that implements this lookup. Note: Support for xarray can be added manually by the user by means of the mechanism described [here](https://xarray.pydata.org/en/stable/internals/extending-xarray.html). Declarative: TBD #### Default Origin Only in cartesian GT4Py, an object can optionally implement a `default_origin` attribute which is used as the origin value unless overwritten by the `origin` keyword argument to the stencil call. ### Allocation For creation, allocation of arrays that perform well in GT4Py, we propose the following set of functions which closely resemble their NumPy counterparts (meaning of the common parameters is explained below): * `empty(shape: Sequence[int], dtype: dtype_like = np.float64, **kwargs) -> ndarray` Allocate a storage with uninitialized (undefined) values. * Parameters: + `shape: Sequence[int]` Sequence of length `ndim` (`ndim` = number of dimensions) with the shape of the storage. + `dtype: dtype_like` The dtype of the storage (NumPy dtype or accepted by `np.dtype()`). It defaults to `np.float64`. For common keyword-only arguments, please see below. * `empty_like(data: Storage, dtype: dtype_like = np.float64, **kwargs) -> ndarray` Allocate a storage with uninitialized (undefined) values, while choosing the not explicitly overridden parameters according to `data`. * Parameters: + `data: Storage` Not explicitly overridden parameters are chosen as the value used in this. `Storage` + `dtype: dtype_like` The dtype of the storage (NumPy dtype or accepted by `np.dtype()`). It defaults to `data.dtype` The common keyword-only arguments can also be overridden. Please see below for their description. Note that `shape` is not a parameter and can not be overridden. * `zeros(shape: Sequence[int], dtype: dtype_like = np.float64, **kwargs) -> ndarray` Allocate a storage with values initialized to 0. * Parameters: + `shape: Sequence[int]` Sequence of length `ndim` (`ndim` = number of dimensions) with the shape of the storage. + `dtype: dtype_like` The dtype of the storage (NumPy dtype or accepted by `np.dtype()`). It defaults to `np.float64`. For common keyword-only arguments, please see below. * `zeros_like(data: Storage, dtype: dtype_like = np.float64, **kwargs) -> ndarray` Allocate a storage with values initialized to 0, while choosing the not explicitly overridden parameters according to `data`. * Parameters: + `data: Storage` Not explicitly overridden parameters are chosen as the value used in this `Storage` + `dtype: dtype_like` The dtype of the storage (NumPy dtype or accepted by `np.dtype()`). It defaults to `data.dtype` The common keyword-only arguments can also be overridden. Please see below for their description. Note that `shape` is not a parameter and can not be overridden. * `ones(shape: Sequence[int], dtype: dtype_like = np.float64, **kwargs) -> ndarray` Allocate a storage with values initialized to 1. * Parameters: + `shape: Sequence[int]` Sequence of length `ndim` (`ndim` = number of dimensions) with the shape of the storage. + `dtype: dtype_like` The dtype of the storage (NumPy dtype or accepted by `np.dtype()`). It defaults to `np.float64`. For common keyword-only arguments, please see below. * `ones_like(data: Storage, dtype: dtype_like = np.float64, **kwargs) -> ndarray` Allocate a storage with values initialized to 1, while choosing the not explicitly overridden parameters according to `data`. * Parameters: + `data: Storage` Not explicitly overridden parameters are chosen as the value used in this `Storage` + `dtype: dtype_like` The dtype of the storage (NumPy dtype or accepted by `np.dtype()`). It defaults to `data.dtype` The common keyword-only arguments can also be overridden. Please see below for their description. Note that `shape` is not a parameter and can not be overridden. * `full(shape: Sequence[int], fill_value: Number, dtype=np.float64, **kwargs) -> ndarray` Allocate a storage with values initialized to the scalar given in `fill_value`. * Parameters: + `shape: Sequence[int]` Sequence of length `ndim` (`ndim` = number of dimensions) with the shape of the storage. + `fill_value: Number`. The number to which the storage is initialized. + `dtype: dtype_like` The dtype of the storage (NumPy dtype or accepted by `np.dtype()`). It defaults to `np.float64`. For common keyword-only arguments, please see below. * `full_like(shape: Sequence[int], fill_value: Number, dtype=np.float64, **kwargs) -> ndarray` Allocate a storage with values initialized to the scalar given in `fill_value`, while choosing the not explicitly overridden parameters according to `data`. * Parameters: + `data: Storage` Not explicitly overridden parameters are chosen as the value used in this `Storage` + `fill_value: Number`. The number to which the storage is initialized. + `dtype: dtype_like` The dtype of the storage (NumPy dtype or accepted by `np.dtype()`). It defaults to `data.dtype` The common keyword-only arguments can also be overridden. Please see below for their description. Note that `shape` is not a parameter and can not be overridden. ##### Optional Keyword-Only Parameters Additionally, these **optional** keyword-only parameters are accepted: * `aligned_index: Sequence[int]` The index of the grid point to which the memory is aligned. For performance optimality, this in general corresponds with the value in the origin used in cartesian gt4py. TODO: Recommendation for declarative? * `alignment_size: Optional[int]` The buffers are allocated such that `mod(aligned_addr, alignment_size) == 0`, where `aligned_addr` is the memory address of the grid point denoted by `aligned_index`. It defaults to `1`, which indicates no alignment. * `backend: Optional[str]` For each backend, a preset of suitable values for `alignment_size`, `device` and `layout` is provided. Explicit definitions of parameters possible and they override its default value from the preset. * `device: Optional[str]` Indicates whether the storage should contain a buffer on an accelerator device. Currently it only accepts `"gpu"` or `None`. Defaults to `None`. * `dims: Optional[Sequence[str]` Sequence indicating the semantic meaning of the dimensions of this storage. This is used to determine the default layout for the storage. For cartesian gt4py, the values `"I"`, `"J"`, `"K"` and additional dimensions as string representations of integers, starting at `"0"` will be supported. TODO: declarative * `layout: Optional[Sequence[int]]` A permutation of integers in `[0 .. ndim-1]`. It indicates the order of strides in decreasing order. I.e. "0" indicates that the stride in that dimension is the largest, while the largest entry in the layout sequence corresponds to the dimension with the smallest stride, which typically is contiguous in memory. Default values as indicated by the `backend` parameter depend on the dimensions. E.g. if `backend` is any of the compiled GridTools backends, the default value is defined according to the semantic meaning of each dimension. For example for the `"gt:cpu_kfirst"` backend, the smallest stride is always in the K dimension, independently of which index corresponds to the K dimension. The use of `backend` and any of the parameters `alignment_size`, `device` or `layout` is exclusive.