InputArrays

class cubie.batchsolving.arrays.BatchInputArrays.InputArrays(precision: type = <class 'numpy.float32'>, 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, sizes: BatchInputSizes | None = NOTHING, host: InputArrayContainer = NOTHING)[source]

Bases: BaseArrayManager

Manage allocation and transfer of batch input arrays.

Parameters:

Notes

Instances are configured from BatchSolverKernel metadata. Updates request memory through the shared manager, ensure array heights match solver expectations, and attach received buffers prior to device transfers.

_attach_device_inputs(device_arrays: Dict[str, object]) None[source]

Attach caller-supplied device arrays as kernel inputs.

Parameters:

device_arrays – Mapping of array names to device arrays.

Raises:
  • ValueError – If a device array’s shape differs from the expected size.

  • TypeError – If a device array’s dtype differs from the slot dtype.

_convert_host_to_numpy() None[source]

Input slots hold caller-supplied arrays verbatim.

Chunked transfers stage slices straight from the attached array, so no conversion is needed.

_convert_host_to_pinned() None[source]

Input slots hold caller-supplied arrays verbatim.

Their backing is classified at attach time and never converted; non-pinned sources stage through the bounded pinned pool instead.

_fast_path_update(initial_values: NDArray, parameters: NDArray, driver_coefficients: NDArray | None) bool[source]

Queue overwrites when the attached inputs are re-supplied.

_invalidate_hook() None[source]

Drop device inputs; attached host data refills new buffers.

_matches_slot(name: str, array: NDArray) bool[source]

Return whether array is the attached, dtype-current slot.

_stage_array(array_name, host_array, device_array, stream) None[source]

Stage one non-pinned input through pooled pinned buffers.

The host source may be a strided view (a chunk slice or a memmap) whose run extent is smaller than the device array’s on the final chunk, so each block is copied into its buffer with shape-aware indexing; a flat copy would misalign the runs. Blocks come from staging_blocks(). Each block’s buffer is handed to the transfer watcher with its own event, so it returns to the pool as soon as its copy lands on the device.

_teardown_cleanups()[source]

Return transfer cleanup calls.

device: InputArrayContainer
property device_driver_coefficients: NDArray | FakeCUDAArray | None

Device driver coefficients array.

property device_initial_values: NDArray | FakeCUDAArray | None

Device initial values array.

property device_parameters: NDArray | FakeCUDAArray | None

Device parameters array.

property driver_coefficients: NDArray | FakeCUDAArray | None

Host driver coefficients array.

finalise(chunk_index: int, stream=None) None[source]

Finish input handling for a chunk.

classmethod from_solver(solver_instance: BatchSolverKernel) InputArrays[source]

Create an InputArrays instance from a solver.

Creates an empty instance from a solver instance, importing the heights of the parameters, initial values, and driver arrays from the ODE system for checking inputs against. Does not allocate host or device arrays.

Parameters:

solver_instance – The solver instance to extract configuration from.

Returns:

A new InputArrays instance configured for the solver.

Return type:

InputArrays

property has_device_inputs: bool

Whether any input was supplied as a device array.

host: InputArrayContainer
property initial_values: NDArray | FakeCUDAArray | None

Initial values used in the last run.

Returns the caller’s device array when initial values were supplied on device; otherwise the host array.

initialise(chunk_index: int, stream=None) None[source]

Copy a batch chunk of host data to device buffers.

Parameters:

chunk_index – Indices for the chunk being initialized.

Notes

Pinned host arrays transfer directly and asynchronously. Everything else stages through pooled pinned buffers whose releases are event-driven via the transfer watcher, so this method never blocks on the stream: with the pool deep enough, the CPU stages the next chunk while the previous kernel runs.

property parameters: NDArray | FakeCUDAArray | None

Parameters used in the last run.

Returns the caller’s device array when parameters were supplied on device; otherwise the host array.

reset() None[source]

Clear all cached arrays and reset allocation tracking.

update(solver_instance: BatchSolverKernel, initial_values: NDArray, parameters: NDArray, driver_coefficients: NDArray | None) None[source]

Set host arrays and request device allocations.

Parameters:
  • solver_instance – The solver instance providing configuration and sizing information.

  • initial_values – Initial state values for each integration run.

  • parameters – Parameter values for each integration run.

  • driver_coefficients – Driver coefficient table; None keeps the attached one.

Notes

Inputs supplied as device arrays are attached directly as the kernel’s device inputs: no host staging or host-to-device transfer occurs, and no managed device buffer is allocated for them. They must already match the expected shape and dtype. A slot’s own device buffer supplied back is not re-uploaded.

update_from_solver(solver_instance: BatchSolverKernel) None[source]

Refresh size, precision, and chunk axis from the solver.

Parameters:

solver_instance – The solver instance to update from.

wait_pending(timeout: float | None = None) None[source]

Wait for pending staging-buffer releases.