cubie.batchsolving.arrays.BaseArrayManager

Base Array Manager Module.

This module provides abstract base classes for managing arrays between host and device memory in batch operations. It includes container classes for storing arrays and manager classes for handling memory allocation, transfer, and synchronization.

Classes

ArrayContainer([stride_order, memory_type, ...])	Base class for storing arrays in CUDA array managers.
BaseArrayManager(precision, sizes, device, ...)	Common base class for managing arrays between host and device.

class cubie.batchsolving.arrays.BaseArrayManager.ArrayContainer(stride_order: tuple[str, ...] | None = None, memory_type: str = 'device', unchunkable: tuple[str] = NOTHING)

Bases: ABC

Base class for storing arrays in CUDA array managers.

Any CUDA array manager should have one subclass of this for both host and device arrays.

Parameters:

• _stride_order (tuple[str, ...], optional) – Order of array dimensions (e.g., (“time”, “run”, “variable”)).

• _memory_type (str, default="device") – Type of memory allocation. Must be one of “device”, “mapped”, “pinned”, “managed”, or “host”.

• _unchunkable (tuple[str], default=()) – Names of arrays that cannot be chunked during memory management.

Notes

Underscored attributes are not really private, but we want to be able to filter them in methods which use __dict__ to get attributes, so we prefix them and add getters/setters.

attach(label, array)

Attach an array to this container.

Parameters:

• label (str) – The name/label for the array.

• array (array_like) – The array to attach.

Warns:

UserWarning – If the specified label does not exist as an attribute.

delete(label)

Delete reference to an array.

Parameters:

label (str) – The name/label of the array to delete.

Warns:

UserWarning – If the specified label does not exist as an attribute.

delete_all()

Delete all array references.

Notes

This method removes all non-private, non-callable attributes, effectively clearing all stored arrays.

property memory_type

Get the memory type.

Returns:

The type of memory allocation.

Return type:

str

property stride_order

Get the stride order.

Returns:

The order of array dimensions.

Return type:

tuple[str, …] or None

class cubie.batchsolving.arrays.BaseArrayManager.BaseArrayManager(precision: type = <class 'numpy.float32'>, sizes: ~cubie.outputhandling.output_sizes.ArraySizingClass | None = None, device: ~cubie.batchsolving.arrays.BaseArrayManager.ArrayContainer = NOTHING, host: ~cubie.batchsolving.arrays.BaseArrayManager.ArrayContainer = NOTHING, chunks: int = 0, chunk_axis: str = 'run', stream_group: str = 'default', memory_proportion: float | None = None, memory_manager: ~cubie.memory.mem_manager.MemoryManager = MemoryManager(totalmem=8589934592, registry={}, stream_groups=StreamGroups(groups={}, streams={}), _mode='passive', _allocator=<class 'cubie.cudasim_utils.FakeNumbaCUDAMemoryManager'>, _auto_pool=[], _manual_pool=[], _stride_order=('time', 'run', 'variable'), _queued_allocations={}))

Bases: ABC

Common base class for managing arrays between host and device.

This class provides a unified interface for MemoryManager integration, allocation/deallocation patterns, stream management, change detection and caching, and queued allocation support.

Parameters:

• _precision (type, default=float32) – Numerical precision type for arrays.

• _sizes (ArraySizingClass, optional) – Size specifications for arrays managed by this instance.

• device (ArrayContainer) – Container for device-side arrays.

• host (ArrayContainer) – Container for host-side arrays.

• _chunks (int, default=0) – Number of chunks for memory management.

• _chunk_axis (str, default="run") – Axis along which to perform chunking. Must be one of “run”, “variable”, or “time”.

• _stream_group (str, default="default") – Stream group identifier for CUDA operations.

• _memory_proportion (float, optional) – Proportion of available memory to use.

• _needs_reallocation (list[str]) – List of array names that need reallocation.

• _needs_overwrite (list[str]) – List of array names that need data overwriting.

• _memory_manager (MemoryManager) – Memory manager instance for handling GPU memory.

Notes

This is an abstract base class that provides common functionality for managing arrays between host and device memory. Subclasses must implement the abstract methods: update, finalise, and initialise.

_arrays_equal(arr1: ndarray[Any, dtype[_ScalarType_co]] | None, arr2: ndarray[Any, dtype[_ScalarType_co]] | None) → bool

Check if two arrays are equal in shape and content.

Parameters:

• arr1 (NDArray or None) – First array to compare.

• arr2 (NDArray or None) – Second array to compare.

Returns:

True if arrays are equal, False otherwise.

Return type:

bool

_invalidate_hook()

Drop all references and assign all arrays for reallocation.

Notes

This method is called when the memory cache needs to be invalidated. It clears all device array references and marks them for reallocation.

_on_allocation_complete(response: ArrayResponse)

Callback for when the allocation response is received.

Parameters:

response (ArrayResponse) – Response object containing allocated arrays and metadata.

Warns:

UserWarning – If a device array is not found in the allocation response.

Notes

WARNING - HERE BE DRAGONS This try/except is to catch case where tests were calling this method with an empty _needs_reallocation list. When the same tests were run one at a time, the error disappeared. I couldn’t trace it to a module-scope fixture or anything obvious. Adding the try/except seems to have suppressed even the warning, and the problem has stopped. If you get this warning, check for the possibility of two different classes calling allocate_queue in between “init” and “initialise”.

_update_host_array(new_array: ndarray[Any, dtype[_ScalarType_co]], current_array: ndarray[Any, dtype[_ScalarType_co]], label: str) → NoneType

Assign for reallocation or overwriting by shape/value change.

Check for equality and shape equality, append to reallocation or overwrite lists accordingly. Attaches changed array to host array container.

allocate()

Queue allocation requests for arrays that need reallocation.

attach_external_arrays(arrays: Dict[str, ndarray[Any, dtype[_ScalarType_co]]], location: str = 'host') → bool

Attach existing arrays to the specified container (host or device).

Parameters:

• arrays – Dictionary of array_name -> array to attach

• location – “host” or “device” - which container to attach to

Returns:

True if arrays were successfully attached, False if validation failed

check_incoming_arrays(arrays: Dict[str, ndarray[Any, dtype[_ScalarType_co]]], location: str = 'host') → Dict[str, bool]

Check dimensions and dtype of provided arrays match expected sizes and precision.

Parameters:

• arrays – Dictionary of array_name -> array to check

• location – “host” or “device” - which container to check against

Returns:

True if all arrays match expected sizes and precision, False otherwise

check_sizes(new_arrays: Dict[str, ndarray[Any, dtype[_ScalarType_co]]], location: str = 'host') → Dict[str, bool]

Check if provided arrays match the system along the “variable” axis.

Parameters:

• new_arrays – Dictionary of array_name -> array

• location – ‘host’ or ‘device’

Returns:

True if all arrays match their expected sizes, False otherwise

check_type(arrays: Dict[str, ndarray[Any, dtype[_ScalarType_co]]]) → Dict[str, bool]

Check if the precision of arrays matches the system precision.

Parameters:

arrays (Dict[str, NDArray]) – Dictionary of array_name -> array.

Returns:

Dictionary mapping array names to boolean values indicating whether each array matches the expected precision.

Return type:

Dict[str, bool]

device: ArrayContainer

abstractmethod finalise(indices)

Override with the desired behaviour after a chunk is executed.

For most output arrays, this is a copy back to the host, and potentially a remap if mapped. For input arrays, this is a typically no-op.

from_device(instance: object, from_arrays: list, to_arrays: list)

host: ArrayContainer

abstractmethod initialise(indices)

Override with the desired behaviour before a chunk is executed.

For most input arrays, this is a copy to device. For output arrays, this is typically a no-op.

initialize_device_zeros()

Initialize device arrays to zero values.

register_with_memory_manager()

Register this instance with the MemoryManager.

Notes

This method sets up the necessary hooks and callbacks for memory management integration.

request_allocation(request: dict[str, ArrayRequest], force_type: str | None = None)

Send a request for allocation of device arrays.

Parameters:

• request (dict[str, ArrayRequest]) – Dictionary mapping array names to allocation requests.

• force_type (str, optional) – Force request type to “single” or “group”. If None, the type is determined automatically based on stream group membership.

Notes

If the object is the only instance in its stream group, or is on the default group, then the request will be sent as a “single” request and be allocated immediately. If the object shares a stream group, then the response will be queued, and the allocation will be grouped with other requests in the same group, until one of the instances calls “process_queue” to process the queue. This behaviour can be overridden by setting force_type to “single” or “group”.

reset()

Clear all cached arrays and reset allocation tracking.

to_device(from_arrays: list, to_arrays: list)

abstractmethod update(*args, **kwargs)

Update arrays from external data.

This method should handle updating the manager’s arrays based on provided input data and trigger reallocation/allocation as needed.

Parameters:

• *args – Variable length argument list.

• **kwargs – Arbitrary keyword arguments.

Notes

This is an abstract method that must be implemented by subclasses with the desired behavior for updating arrays from external data.

update_host_arrays(new_arrays: Dict[str, ndarray[Any, dtype[_ScalarType_co]]])

Updates host arrays with new data, assigns for realloc or overwrite.

Parameters:

new_arrays – Dictionary of array_name -> new_array

Returns:

None

update_sizes(sizes: ArraySizingClass)

Update the expected sizes for arrays in this manager.

Parameters:

sizes (ArraySizingClass) – An ArraySizingClass instance with new sizes.

Raises:

TypeError – If the new sizes object is not the same size as the existing one.