Skip to content

[Feature] Add estimate_memory() for pipeline GPU footprint #1429

Description

@0lai0

What

Add a pure estimate_memory() function that, given an encoding configuration, returns the predicted CPU prefetch-pool and GPU state-buffer memory the pipeline will allocate — so callers can know the footprint before the first batch runs. No allocation, no CUDA calls, no side effects.

Why

There is no upfront memory check anywhere in the pipeline today. Users only discover they are over budget when CUDA throws OOM mid-run — sometimes after 90+ minutes of encoding (see #1262). Because the failure surfaces deep inside encode_batch, the error gives no hint about which knob (batch size, qubit count, prefetch depth) caused it. A cheap, deterministic estimate computed from config alone lets us fail fast (B2) and gives users an actionable number to tune against.

How

  • Add a pure function in qdp-core (e.g. qdp-core/src/estimate.rs — note qdp-core/src/gpu/memory.rs already exists, so avoid a bare memory.rs at the crate root to prevent confusion):
pub fn estimate_memory(
    encoding: Encoding,
    num_qubits: u32,
    batch_size: usize,
    dtype: Dtype,          // existing enum in qdp-core/src/types.rs (Float32/Float64)
    prefetch_depth: usize,
) -> MemoryEstimate
  • Introduce a MemoryEstimate struct exposing at least cpu_prefetch_bytes, gpu_state_bytes, and a total() helper.

Memory model. These formulas are the source of truth for this issue:

cpu_prefetch_bytes = prefetch_depth × batch_size × sample_size × bytes_per_elem
gpu_state_bytes    = 2 × batch_size × state_len × complex_bytes
total()            = cpu_prefetch_bytes + gpu_state_bytes

where:

term definition
bytes_per_elem from Dtype::bytes() — 4 for Float32, 8 for Float64. Do not hardcode.
complex_bytes 2 × bytes_per_elem (8 for f32 → complex64, 16 for f64 → complex128)
sample_size input vector length; for amplitude, 2^num_qubits
state_len output state vector length; for amplitude, 2^num_qubits
  • Derive sample_size / state_len from encoding + num_qubits so the function stays correct across encodings.

Modeling notes — please preserve these semantics and document them in rustdoc:

  1. The leading on gpu_state_bytes is a double-buffering allowance (state vector + recycle buffer), not the complex real/imaginary factor — the real/imag pair is already inside complex_bytes. The current batch path (GpuStateVector::new_batch) allocates a single batch_size × state_len × complex_bytes buffer, so the is a deliberate conservative upper bound.
  2. cpu_prefetch_bytes covers exactly prefetch_depth batches, matching the formula. The true instantaneous peak can reach prefetch_depth + 2 batches (one being produced, one being consumed, alongside a full channel). Document this in rustdoc but do not fold the +2 into the returned value — the returned number stays aligned with the formula above.
  3. prefetch_depth is a parameter here, not something this function computes. The pipeline auto-computes it at runtime (pipeline_runner.rs:153; it yields 16 at 16 qubits and 1 at 20 qubits). estimate_memory must not replicate that logic — callers pass the resolved value.
  4. Known exclusion: the host-side streaming chunk buffer (Parquet reader) is not included in this estimate. Document this explicitly in rustdoc so B2 does not silently under-count for streaming sources. Folding it in requires reader config beyond this signature — tracked as a follow-up.

Out of scope: VRAM querying and fail-fast wiring (B2), Python bindings, the streaming chunk buffer term (note 4), non-amplitude encodings beyond the size formula already needed here.

Acceptance criteria

  • estimate_memory is pure: no allocation, no CUDA dependency
  • Unit tests for the two reference cases below assert the exact byte values:
case encoding qubits dtype batch_size prefetch_depth cpu_prefetch_bytes gpu_state_bytes total()
A amplitude 16 f32 64 16 268,435,456 (256 MiB) 67,108,864 (64 MiB) 335,544,320 (320 MiB)
B amplitude 20 f32 64 1 268,435,456 (256 MiB) 1,073,741,824 (1 GiB) 1,342,177,280 (1.25 GiB)

The prefetch_depth values match what the pipeline auto-computes at those qubit counts, but the tests should pass them explicitly rather than depending on the auto-compute path.

  • Unit test covering the f64 vs f32 bytes_per_elem difference via Dtype::bytes() — case A with dtype = Float64 must double both terms: cpu = 536,870,912 (512 MiB), gpu = 134,217,728 (128 MiB), total = 671,088,640 (640 MiB)
  • MemoryEstimate documented in rustdoc with the formula and a worked example
  • Rustdoc states the three documented caveats: the is double-buffering (not the complex factor), the +2 instantaneous peak is excluded, and the streaming chunk buffer is excluded
  • cargo test -p qdp-core passes on Linux

Metadata

Metadata

Assignees

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions