BaseArrayManager

class cubie.batchsolving.arrays.BaseArrayManager.BaseArrayManager(precision: type = <class 'numpy.float32'>, sizes: ArraySizingClass | None = None, device: ArrayContainer = NOTHING, host: ArrayContainer = NOTHING, chunks: int = 0, stream_group: str = 'default', memory_proportion: float | None = None, memory_manager: MemoryManager = MemoryManager(totalmem=8589934592, registry={}, stream_groups=StreamGroups(groups={}, streams={}), _mode='passive', _allocator=<class 'cubie.cuda_simsafe.FakeNumbaCUDAMemoryManager'>, _auto_pool=[], _manual_pool=[], _queued_allocations={}), num_runs: int = 1)

Bases: ABC

Coordinate allocation and transfer for batch host and device arrays.

Parameters:

• _precision (type) – Precision factory used to create new arrays.

• _sizes (cubie.outputhandling.output_sizes.ArraySizingClass | None) – Size specifications for arrays managed by this instance.

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

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

• _chunks (int) – Number of chunks for memory management. Chunking is always performed along the run axis.

• _stream_group (str) – Stream group identifier for CUDA operations.

• _memory_proportion (float | None) – Proportion of available memory to use.

• _needs_reallocation (list[str]) – Array names that require device reallocation.

• _needs_overwrite (list[str]) – Array names that require host overwrite.

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

Notes

Subclasses must implement update(), finalise(), and initialise() to wire batching behaviour into host and device execution paths.

_arrays_equal(arr1: ndarray[Any, dtype[_ScalarType_co]] | None, arr2: ndarray[Any, dtype[_ScalarType_co]] | None, check_type: bool = True, shape_only: bool = False) → bool

Check if two arrays are equal in shape and optionally content.

Parameters:

• arr1 – First array or None.

• arr2 – Second array or None.

• check_type – Check dtype equality. Defaults to True.

• shape_only – Skip element comparison; only check shape and optionally dtype. Faster for output arrays that will be overwritten. Defaults to False.

Returns:

True if arrays are equal, False otherwise.

Return type:

bool

_convert_host_to_numpy() → None

Convert pinned host arrays to regular numpy for chunked mode.

When chunking is active, host arrays should be regular numpy to limit pinned memory usage. Per-chunk pinned buffers are used for staging during transfers.

_convert_host_to_pinned() → None

Convert regular numpy host arrays to pinned for non-chunked mode.

When a run is not chunked, the host arrays should be pinned to enable asynchronous transfers.

_invalidate_hook() → None

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.

Returns:

Nothing is returned.

Return type:

None

_on_allocation_complete(response: ArrayResponse) → None

Callback for when the allocation response is received.

Parameters:

response – Response object containing allocated arrays and metadata.

Warns:

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

Returns:

Nothing is returned.

Return type:

None

Notes

Warnings are only issued if the response contains some arrays but not the expected one, indicating a potential allocation mismatch.

Stores chunk parameters from response in ManagedArray objects for both host and device containers.

_update_host_array(new_array: ndarray[Any, dtype[_ScalarType_co]], current_array: ndarray[Any, dtype[_ScalarType_co]] | None, label: str, shape_only: bool = False) → None

Mark host arrays for overwrite or reallocation based on updates.

Parameters:

• new_array – Updated array that should replace the stored host array.

• current_array – Previously stored host array or None.

• label – Array name used to index tracking lists.

• shape_only – Only check shape equality when comparing arrays. Faster for output arrays that will be overwritten. Defaults to False.

Raises:

ValueError – If new_array is None.

Returns:

Nothing is returned.

Return type:

None

allocate() → None

Queue allocation requests for arrays that need reallocation.

Notes

Builds ArrayRequest objects for arrays marked for reallocation and sets the unchunkable hint based on host metadata.

Chunking is always performed along the run axis by convention. The specific axis index is determined by each array’s chunk_axis_index.

Returns:

Nothing is returned.

Return type:

None

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

Attach existing arrays to a host or device container.

Parameters:

• arrays – Dictionary mapping array names to arrays.

• location – "host" or "device" indicating the target container.

Returns:

True if arrays pass validation, False otherwise.

Return type:

bool

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

Validate shape and precision for incoming arrays.

Parameters:

• arrays – Dictionary mapping array names to arrays.

• location – "host" or "device" indicating the target container.

Returns:

Dictionary indicating whether each array is ready for attachment.

Return type:

Dict[str, bool]

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

Check whether arrays match configured sizes and stride order.

Parameters:

• new_arrays – Dictionary mapping array names to arrays.

• location – "host" or "device" indicating which container to inspect.

Returns:

Dictionary indicating whether each array matches its expected shape.

Return type:

Dict[str, bool]

Raises:

AttributeError – If the location is neither "host" nor "device".

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

Check if the dtype of arrays matches their stored dtype.

Parameters:

arrays – Dictionary mapping array names to arrays.

Returns:

Dictionary indicating whether each array matches the expected precision.

Return type:

Dict[str, bool]

device: ArrayContainer

abstractmethod finalise(chunk_index: int) → None

Execute post-chunk behaviour for device outputs.

Parameters:

chunk_index – Chunk index about to run on the device

Returns:

Nothing is returned.

Return type:

None

from_device(from_arrays: List[object], to_arrays: List[object]) → None

Copy device arrays back to the host using the memory manager.

Parameters:

• from_arrays – Device arrays to copy.

• to_arrays – Destination host arrays.

Returns:

Nothing is returned.

Return type:

None

host: ArrayContainer

abstractmethod initialise(chunk_index: int) → None

Execute pre-chunk behaviour for device inputs.

Parameters:

chunk_index – Chunk index about to run on the device.

Returns:

Nothing is returned.

Return type:

None

initialize_device_zeros() → None

Initialize device arrays to zero values.

Returns:

Nothing is returned.

Return type:

None

property is_chunked: bool

Return True if arrays are being processed in multiple chunks.

num_runs: int

register_with_memory_manager() → None

Register this instance with the MemoryManager.

Notes

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

Returns:

Nothing is returned.

Return type:

None

request_allocation(request: dict[str, ArrayRequest]) → None

Send a request for allocation of device arrays.

Parameters:

request – Dictionary mapping array names to allocation requests.

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”.

Returns:

Nothing is returned.

Return type:

None

reset() → None

Clear all cached arrays and reset allocation tracking.

Returns:

Nothing is returned.

Return type:

None

set_array_runs(num_runs: int) → None

Update num_runs in all ManagedArray instances.

This method sets the num_runs attribute to specify the total number of runs in the batch. This value is used during allocation to determine chunking behavior.

Parameters:

num_runs (int) – Total number of runs in the batch. Must be >= 1.

Returns:

Nothing is returned.

Return type:

None

to_device(from_arrays: List[object], to_arrays: List[object]) → None

Copy host arrays to the device using the memory manager.

Parameters:

• from_arrays – Host arrays to copy.

• to_arrays – Destination device arrays.

Returns:

Nothing is returned.

Return type:

None

abstractmethod update(*args: object, **kwargs: object) → None

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 – Positional arguments passed by subclasses.

• **kwargs – Keyword arguments passed by subclasses.

Notes

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

Returns:

Nothing is returned.

Return type:

None

update_host_arrays(new_arrays: Dict[str, ndarray[Any, dtype[_ScalarType_co]]], shape_only: bool = False) → None

Update host arrays and record allocation requirements.

Parameters:

• new_arrays – Dictionary mapping array names to new host arrays.

• shape_only – Only check shape equality when comparing arrays. Faster for output arrays that will be overwritten. Defaults to False.

Returns:

Nothing is returned.

Return type:

None

update_sizes(sizes: ArraySizingClass) → None

Update the expected sizes for arrays in this manager.

Parameters:

sizes – Array sizing configuration with new dimensions.

Raises:

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

Returns:

Nothing is returned.

Return type:

None