cuDecomp Python API

These are all the types and functions available in the cuDecomp Python API. See Python API (experimental) for array, datatype, stream, and resource conventions. Methods raise exceptions on failure instead of returning C API result codes.

Types

Internal types

class cudecomp.Handle(mpi_comm)

Collectively initialize cuDecomp from an existing MPI communicator.

Parameters:: mpi_comm – Initialized mpi4py.MPI.Comm containing participating ranks.

Select the CUDA device before construction. Keep the communicator valid until finalize().

class cudecomp.GridDesc: Grid descriptor returned by Handle.create_grid(); not directly constructible. Retains its Python handle until released. Library resources require explicit destroy().

Grid Descriptor Configuration

class cudecomp.GridDescConfig

A data structure defining configuration options for grid descriptor creation. The constructor initializes its fields to defaults. Array properties are writable NumPy views; whole-field assignment accepts tuples, lists, or NumPy arrays. See Basic Usage Guide for examples and cudecompGridDescConfig_t for configuration constraints.

Variables:

gdims – Dimensions of global data grid in X/Y/Z order, shape (3,) (default: all zeros).
gdims_dist – Dimensions of global data grid to use for distribution, shape (3,) (default: all zeros, use gdims).
pdims – Dimensions of process grid, P_rows and P_cols, shape (2,) (default: both zero, request autotuning).
rank_order – Process grid rank assignment order (default: RankOrder.DEFAULT, row-major).
transpose_comm_backend – Communication backend to use for transpose communication (default: TransposeCommBackend.MPI_P2P).
transpose_axis_contiguous – Flags by axis indicating if memory should be contiguous along the pencil axis, shape (3,) (default: all False).
transpose_mem_order – Memory orders by pencil axis, shape (3, 3). Each row specifies a permutation of (0, 1, 2) from fastest to slowest and overrides transpose_axis_contiguous. All rows must be set when used (default: all -1, unset).
halo_comm_backend – Communication backend to use for halo communication (default: HaloCommBackend.MPI).

class cudecomp.GridDescAutotuneOptions

A data structure defining autotuning options for grid descriptor creation. The constructor initializes its fields to defaults. Array properties have the same assignment behavior as GridDescConfig. See Autotuning for usage and cudecompGridDescAutotuneOptions_t for configuration constraints.

Variables:

n_warmup_trials – Number of warmup trials for each tested configuration during autotuning (default: 3).
n_trials – Number of timed trials for each tested configuration during autotuning (default: 5).
grid_mode – Communication type (transpose/halo) to use to autotune the process grid (default: AutotuneGridMode.TRANSPOSE).
dtype – Data type to use during autotuning; accepts a DataType member (default: DataType.DOUBLE).
allow_uneven_decompositions – Allow process grids that result in uneven distributions of elements across processes (default: True).
disable_mpi_backends – Disable MPI backend options during autotuning (default: False).
disable_nccl_backends – Disable NCCL backend options during autotuning (default: False).
disable_nvshmem_backends – Disable NVSHMEM backend options during autotuning (default: False).
skip_threshold – Skip a configuration if skip_threshold * t > t_best, where t is its first trial time and t_best the best average time (default: 0.0, disabled).
autotune_transpose_backend – Enable transpose backend tuning (default: False).
transpose_use_inplace_buffers – Flags controlling in-place or out-of-place transpose buffers by operation, shape (4,) (default: all False).
transpose_op_weights – Multiplicative weights for trial time contributions by transpose operation, shape (4,) (default: all 1.0).
transpose_input_halo_extents – Input halo extent arguments to use during autotuning by transpose operation, shape (4, 3) (default: all zeros).
transpose_output_halo_extents – Output halo extent arguments to use during autotuning by transpose operation, shape (4, 3) (default: all zeros).
transpose_input_padding – Input padding arguments to use during autotuning by transpose operation, shape (4, 3) (default: all zeros).
transpose_output_padding – Output padding arguments to use during autotuning by transpose operation, shape (4, 3) (default: all zeros).
autotune_halo_backend – Enable halo backend tuning (default: False).
halo_extents – Symmetric halo extents in global X/Y/Z order, shape (3,) (default: all zeros).
halo_periods – Periodicity in global X/Y/Z order, shape (3,) (default: all False).
halo_axis – Pencil axis for halo tuning: 0=X, 1=Y, 2=Z (default: 0).
halo_padding – Padding in global X/Y/Z order, shape (3,) (default: all zeros).

Transpose array entries use operation order X-to-Y, Y-to-Z, Z-to-Y, Y-to-X. Halo and padding rows use global X/Y/Z order.

Pencil Information

class cudecomp.PencilInfo

A data structure containing geometry information about a pencil data buffer, returned by GridDesc.get_pencil_info(). Properties are read-only. Array properties are tuples; shapes and bounds use Python C-order. See Working with pencil data for local/global indexing conventions.

Variables:

shape – Pencil shape in local C-order, including halos and padding.
lo – Zero-based inclusive lower interior bounds in local order.
hi – Zero-based inclusive upper interior bounds in local order.
order – Global axis corresponding to each local array dimension.
halo_extents – Symmetric halo extents in global X/Y/Z order.
padding – Upper-end padding extents in global X/Y/Z order.
size – Number of elements in the pencil, including halos and padding.

Communication Backends

class cudecomp.TransposeCommBackend: Members: MPI_P2P, MPI_P2P_PL, MPI_A2A, NCCL, NCCL_PL, NVSHMEM, NVSHMEM_PL, NVSHMEM_SM. See cudecompTransposeCommBackend_t for backend behavior and restrictions. Enum availability does not imply the backend was built.

class cudecomp.HaloCommBackend: Members: MPI, MPI_BLOCKING, NCCL, NVSHMEM, NVSHMEM_BLOCKING. See cudecompHaloCommBackend_t.

Additional Enumerators

class cudecomp.DataType: Members: FLOAT (float32), DOUBLE (float64), FLOAT_COMPLEX (complex64), DOUBLE_COMPLEX (complex128).

class cudecomp.AutotuneGridMode: Members: TRANSPOSE and HALO. See cudecompAutotuneGridMode_t.

class cudecomp.RankOrder: Members: DEFAULT, ROW_MAJOR, COL_MAJOR. See cudecompRankOrder_t.

class cudecomp.Result: C API result constants: SUCCESS, INVALID_USAGE, NOT_SUPPORTED, INTERNAL_ERROR, CUDA_ERROR, CUTENSOR_ERROR, MPI_ERROR, NCCL_ERROR, NVSHMEM_ERROR, NVML_ERROR. Public Python operations return objects or None and raise exceptions instead of returning these codes.

Functions

Library Initialization/Finalization

Construct Handle with an MPI communicator to initialize cuDecomp.

Handle.finalize()

Collectively finalize the handle and release its resources. Destroy grids before finalizing the handle. Repeated calls are no-ops.

Returns:: None.

Grid Descriptor Management

Handle.create_grid(config, autotune_options=None)

Collectively create and return a GridDesc.

Parameters:

config – Initialized GridDescConfig. Grid creation may update its properties with the selected configuration.
autotune_options – GridDescAutotuneOptions, or None to disable autotuning (equivalent to passing NULL in C). If config.pdims is (0, 0), options are required; otherwise InvalidUsageError is raised.

Returns:

The new grid descriptor.

GridDesc.destroy()

Collectively destroy the descriptor. Repeated calls are no-ops; other methods on a destroyed grid raise InvalidUsageError.

Returns:: None.

Helper Functions

GridDesc.get_config(): Return an independent GridDescConfig containing the resolved configuration. Mutating this copy does not reconfigure the grid.

GridDesc.get_pencil_info(axis, halo_extents=None, padding=None)

Collect geometry information about the pencil assigned to the calling process for the given axis. Returns a PencilInfo object.

Parameters:

axis – Pencil orientation: 0=X, 1=Y, 2=Z.
halo_extents – Three nonnegative symmetric extents in global X/Y/Z order. None means no halos.
padding – Three nonnegative upper-end padding extents in global X/Y/Z order. None means no padding.

GridDesc.get_shifted_rank(axis, dim, displacement, periodic)

Return the neighbor’s rank in the handle communicator, or -1 if the neighbor does not exist at a nonperiodic boundary.

Parameters:

axis – Pencil orientation: 0=X, 1=Y, 2=Z.
dim – Global direction to shift in: 0=X, 1=Y, 2=Z.
displacement – Signed integer neighbor displacement.
periodic – Boolean selecting periodic boundaries in dim.

cudecomp.get_data_type_size(dtype): Get the size (in bytes) of a cuDecomp DataType.

cudecomp.transpose_comm_backend_to_string(backend): Return the display name of a TransposeCommBackend member (for example "NCCL (pipelined)"), not the enum member name.

cudecomp.halo_comm_backend_to_string(backend): Return the display name of a HaloCommBackend member, not the enum member name.

Transposition Methods

All transpose methods are collective and use automatically managed workspaces. Input and output arrays must be C-contiguous and match their respective pencil query shapes, including halos and padding. These requirements also apply to in-place operations.

GridDesc.transpose_x_to_y(input, out=None, dtype=None, input_halo_extents=None, output_halo_extents=None, input_padding=None, output_padding=None, stream=None)

Transpose data from X-axis aligned pencils to Y-axis aligned pencils.

Parameters:

input – Input pencil array, as described in Python API (experimental).
out – Writable output pencil array. None allocates a CuPy array and requires a CuPy input.
dtype – DataType member; None infers it from the input array. No conversion is performed.
input_halo_extents – Three input halo extents in global order; None means no halos.
output_halo_extents – Three output halo extents in global order; None means no halos.
input_padding – Three input padding extents in global order; None means no padding.
output_padding – Three output padding extents in global order; None means no padding.
stream – CUDA stream handle, object implementing __cuda_stream__, or object with integer ptr. None uses CuPy’s current stream for CuPy inputs and CUDA stream 0 otherwise. See Python API (experimental) for ordering requirements.

Returns:

The supplied out object or a new C-contiguous CuPy output.

Only interior values are transposed; output halos and padding are not initialized. Operations are asynchronous with respect to the host.

GridDesc.transpose_y_to_z(input, out=None, dtype=None, input_halo_extents=None, output_halo_extents=None, input_padding=None, output_padding=None, stream=None): Transpose data from Y-axis aligned pencils to Z-axis aligned pencils. Parameters and returns are as in GridDesc.transpose_x_to_y().

GridDesc.transpose_z_to_y(input, out=None, dtype=None, input_halo_extents=None, output_halo_extents=None, input_padding=None, output_padding=None, stream=None): Transpose data from Z-axis aligned pencils to Y-axis aligned pencils. Parameters and returns are as in GridDesc.transpose_x_to_y().

GridDesc.transpose_y_to_x(input, out=None, dtype=None, input_halo_extents=None, output_halo_extents=None, input_padding=None, output_padding=None, stream=None): Transpose data from Y-axis aligned pencils to X-axis aligned pencils. Parameters and returns are as in GridDesc.transpose_x_to_y().

Halo Exchange Methods

All halo methods are collective, update arrays in place, and use automatically managed workspaces. Input arrays must be writable, C-contiguous, and match the corresponding pencil query shape, including halos and padding.

GridDesc.update_halos_x(input, halo_extents, dim, dtype=None, halo_periods=None, padding=None, stream=None)

Perform halo communication of X-axis aligned pencil data.

Parameters:

input – Writable input pencil array matching get_pencil_info(0, halo_extents, padding).
halo_extents – Three nonnegative symmetric extents in global X/Y/Z order.
dim – Global update direction: 0=X, 1=Y, 2=Z.
dtype – DataType member; None infers it from the array. No conversion is performed.
halo_periods – Three booleans selecting periodic boundaries in global order. None means all nonperiodic.
padding – Three padding extents in global order; None means none.
stream – CUDA stream handle, object implementing __cuda_stream__, or object with integer ptr; None uses CuPy’s current stream for CuPy inputs and CUDA stream 0 otherwise. See Python API (experimental) for ordering requirements.

Returns:

The same input object. Operations are asynchronous with respect to the host.

Halos with no neighbor at a nonperiodic boundary are unchanged. Padding elements are not part of the halo update.

GridDesc.update_halos_y(input, halo_extents, dim, dtype=None, halo_periods=None, padding=None, stream=None): Perform halo communication of Y-axis aligned pencil data. Parameters and returns are as in GridDesc.update_halos_x(), with input matching the Y-pencil query.

GridDesc.update_halos_z(input, halo_extents, dim, dtype=None, halo_periods=None, padding=None, stream=None): Perform halo communication of Z-axis aligned pencil data. Parameters and returns are as in GridDesc.update_halos_x(), with input matching the Z-pencil query.

Errors and Version

Arrays that do not meet cuDecomp’s requirements (dtype, shape, contiguity, writability, byte order, or masks) raise InvalidUsageError. Arguments of the wrong kind and malformed CUDA array-interface metadata raise TypeError, ValueError, or OverflowError; errors raised by array providers may propagate as well. Missing required CUDA array-interface fields raise KeyError. With dtype=None, an input whose dtype is not a supported cuDecomp type raises TypeError. C API result codes are translated into the exceptions below.

Numeric arguments use pybind11’s standard conversions, including corresponding NumPy scalar types. Integer arguments do not accept floating-point values. Fixed-size sequence arguments must have the required length; argument conversion failures raise TypeError.

exception cudecomp.CuDecompError: Base class for cuDecomp errors; derives from RuntimeError.

exception cudecomp.InvalidUsageError: Invalid API usage, arrays that do not meet cuDecomp’s requirements, or use of a finalized handle or destroyed grid. Derives from CuDecompError.

exception cudecomp.NotSupportedError: Unsupported configuration. Derives from CuDecompError.

exception cudecomp.InternalError: Internal library failure. Derives from CuDecompError.

exception cudecomp.CudaError
exception cudecomp.CutensorError
exception cudecomp.MpiError
exception cudecomp.NcclError
exception cudecomp.NvshmemError
exception cudecomp.NvmlError: Errors reported by the corresponding dependency. Each derives from CuDecompError.

cudecomp.__version__: cuDecomp package version.