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', _auto_pool=[], _manual_pool=[], _queued_allocations={}, _group_chunk_parameters={}, _usage_clock=0, _pending_teardowns=[], pinned_max_bytes=8589934592, allocation_granule_bytes=33554432, _pinned_lock=<unlocked _thread.lock object>, _pinned_live_bytes=0, _pinned_retained_bytes=0, _pinned_releases=deque([]), _device_probe_error=None), memory_owner: object | None = None, 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 | None, arr2: NDArray | 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

static _base_memory_type(memory_type: str) → str

Return the type to request when replacing a slot’s array.

A spilled slot re-requests pinned backing; the replacement spills again only if its size still exceeds the policy.

_convert_host_to_numpy() → None

Move chunk-staged kernel-written slots to pageable backing.

When chunking is active, full-size host buffers stay pageable and per-chunk transfers stage through the bounded pinned pool. Slot content is not preserved: this applies only to buffers each chunk writes back into, and input managers override it as a no-op because their slots hold caller-supplied arrays verbatim.

_convert_host_to_pinned() → None

Repin unchunked kernel-written slots for direct transfers.

Runs after the chunk decision, so pinning never happens for a chunked solve. Slot content is not preserved: this applies only to buffers the kernel’s device transfers overwrite, and input managers override it as a no-op because their slots hold caller-supplied arrays verbatim. A slot whose replacement the pinned budget refuses keeps its pageable buffer.

static _host_memory_type(array: NDArray) → str

Classify a host array’s actual backing.

"pinned" requires C-contiguous page-locked memory, which transfers directly and asynchronously. Strided or pageable arrays stage through bounded pinned blocks, and memmaps are disk-backed.

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

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

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.

static _requires_staging(array: NDArray, memory_type: str) → bool

Return whether a host array needs pinned staging.

_teardown_cleanups() → List[Callable[[], None]]

Return cleanup calls that do not capture this manager.

_update_host_array(new_array: NDArray, current_array: NDArray | None, label: str, shape_only: bool = False) → None

Attach one incoming host array and record allocation needs.

Parameters:

• new_array – Replacement array for the slot.

• current_array – Previously stored host array or None.

• label – Array name used to index tracking lists.

• shape_only – The stored buffer only needs to match new_array’s shape; values are ignored. Used for output arrays, which the kernel overwrites.

Raises:

ValueError – If new_array is None.

Notes

Incoming arrays are attached verbatim: the slot records the array’s actual backing (pinned, pageable, or memmap) and transfers read from it directly, so no copy is made. The one exception is a dtype mismatch, which is cast once into a fresh buffer. The solve reads input arrays when it runs — mutating a submitted array between solves changes what the next solve integrates.

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.

check_incoming_arrays(arrays: Dict[str, NDArray], 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], 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]) → 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]

close() → None

Release this manager’s resources.

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

from_device(from_arrays: List[object], to_arrays: List[object], stream: Any | None = None) → None

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

Parameters:

• from_arrays – Device arrays to copy.

• to_arrays – Destination host arrays.

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.

property is_chunked: bool

Return True if arrays are being processed in multiple chunks.

loan_host_arrays(owner: object) → None

Hand every host buffer to owner, emptying the slots.

owner (a result object that already references the arrays) keeps the data for as long as it lives. If it has been garbage collected by the next solve, the buffers return to their slots and are reused; otherwise the next solve allocates fresh backing.

num_runs: int

reclaim_or_release_loan() → None

Recover loaned host buffers if their owner was collected.

A live owner keeps its buffers: the loan record is dropped so the arrays belong solely to the owner, and the next allocation builds fresh backing. A collected owner cannot be holding views, so the buffers return to their slots for reuse.

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.

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

reset() → None

Clear cached arrays and allocation tracking.

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.

to_device(from_arrays: List[object], to_arrays: List[object], stream: Any | None = None) → None

Copy host arrays to the device using the memory manager.

Parameters:

• from_arrays – Host arrays to copy.

• to_arrays – Destination device arrays.

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.

update_host_arrays(new_arrays: Dict[str, NDArray], 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 – Stored buffers only need to match the new arrays’ shapes; values are ignored. Used for output arrays, which the kernel overwrites. Defaults to False.

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.