Skip to content

gpu_utils#

Includes utility functions for GPU operations.

Functions:

get_array_module_name #

get_array_module_name(arr: NDArray) -> str

Given an array, returns the array's module name.

This works for numpy even when cupy is not available.

Parameters:

  • arr (NDArray) –

    The array to check.

Returns:

  • str –

    The array module name used by the array.

get_host #

get_host(arr: NDArray, out: None | NDArray = None) -> NDArray

Returns the host array of the given array.

Note: special behaviour if numpy is used: If out is not set, then the returned array is the same as the input and the pointers alias.

Parameters:

  • arr (NDArray) –

    The array to convert.

  • out (NDArray, default: None ) –

    The output array.

Returns:

  • ndarray –

    The equivalent numpy array.

get_device #

get_device(arr: NDArray, out: None | NDArray = None) -> NDArray

Returns the device array of the given array.

Note: special behaviour if cupy is used: If out is not set, then the returned array is the same as the input and the pointers alias.

Parameters:

  • arr (NDArray) –

    The array to convert.

  • out (NDArray, default: None ) –

    The output array.

Returns:

  • NDArray –

    The equivalent cupy array.

get_any_location #

get_any_location(arr: NDArray, output_module: str, use_pinned_memory: bool = False)

Returns the array in the desired location.

Parameters:

  • arr (NDArray) –

    The array to convert.

  • output_module (str) –

    The desired location. The location can be either "numpy" or "cupy".

  • use_pinned_memory (bool, default: False ) –

    Whether to use pinnend memory if cupy is used. Default is True.

Returns:

  • NDArray –

    The equivalent array in the desired location

empty_pinned #

empty_pinned(shape: int | tuple[int, ...], dtype: dtype = float, order: str = 'C')

Returns a new, uninitialized NumPy array with the given shape and dtype. The array is allocated in pinned memory if using cupy.

Parameters:

  • shape (int or tuple of ints) –

    Shape of the empty array.

  • dtype (data - type, default: float ) –

    Desired data-type for the array. Default is float.

  • order ((C, F), default: 'C' ) –

    Whether to store multi-dimensional data in row-major (C-style) or column-major (Fortran-style) order in memory. Default is 'C'.

Returns:

  • NDArray –

    The empty array.

  • .. seealso:: :func:`numpy.empty` :func:`cupy.empty` :func:`cupyx.empty_pinned` –

zeros_pinned #

zeros_pinned(shape: int | tuple[int, ...], dtype: dtype = float, order: str = 'C')

Returns a new array of given shape and type, filled with zeros. The array is allocated in pinned memory if using cupy.

Parameters:

  • shape (int or tuple of ints) –

    Shape of the new array.

  • dtype (data - type, default: float ) –

    The desired data-type for the array. Default is float.

  • order ((C, F), default: 'C' ) –

    Whether to store multi-dimensional data in row-major (C-style) or column-major (Fortran-style) order in memory. Default is 'C'.

Returns:

  • NDArray –

    The array of zeros.

  • .. seealso:: :func:`numpy.zeros` :func:`cupy.zeros` :func:`cupyx.zeros_pinned` –

empty_like_pinned #

empty_like_pinned(a: NDArray, dtype: dtype = None, order: str = 'K', shape: int | tuple[int, ...] = None)

Returns a new array with the same shape and type as a given array. The array is allocated in pinned memory if using cupy.

Parameters:

  • a (NDArray) –

    The shape and data-type of a define these same attributes of the returned array.

  • dtype (data - type, default: None ) –

    Overrides the data type of the result.

  • order ((C, F, A, K), default: 'C' ) –

    Overrides the memory layout of the result. 'C' means C-order, 'F' means F-order, 'A' means 'F' if a is Fortran contiguous, 'C' otherwise. 'K' means match the layout of a as closely as possible.

  • shape (int or tuple of ints, default: None ) –

    Overrides the shape of the result.

Returns:

  • NDArray –

    The empty array.

  • .. seealso:: :func:`numpy.empty_like` :func:`cupy.empty_like` :func:`cupyx.empty_like_pinned` –

zeros_like_pinned #

zeros_like_pinned(a: NDArray, dtype: dtype = None, order: str = 'K', shape: int | tuple[int, ...] = None)

Returns an array of zeros with the same shape and type as a given array. The array is allocated in pinned memory if using cupy.

Parameters:

  • a (NDArray) –

    The shape and data-type of a define these same attributes of the returned array.

  • dtype (data - type, default: None ) –

    Overrides the data type of the result.

  • order ((C, F, A, K), default: 'C' ) –

    Overrides the memory layout of the result. 'C' means C-order, 'F' means F-order, 'A' means 'F' if a is Fortran contiguous, 'C' otherwise. 'K' means match the layout of a as closely as possible.

  • shape (int or tuple of ints, default: None ) –

    Overrides the shape of the result.

Returns:

  • NDArray –

    The array of zeros.

  • .. seealso:: :func:`numpy.zeros_like` :func:`cupy.zeros_like` :func:`cupyx.zeros_like_pinned` –

synchronize_current_stream #

synchronize_current_stream()

Synchronizes the current stream if using cupy.

Does nothing if using numpy.

synchronize_device #

synchronize_device()

Synchronizes the device if using cupy.

Does nothing if using numpy.

free_mempool #

free_mempool()

Frees the memory pool if using cupy.

Does nothing if using numpy.