diff --git a/.github/workflows/python-testing.yml b/.github/workflows/python-testing.yml index 9f025647ad..fe348f1f30 100644 --- a/.github/workflows/python-testing.yml +++ b/.github/workflows/python-testing.yml @@ -62,12 +62,25 @@ jobs: # GPU fidelity cases self-skip when the probe finds no functional GPU # (stub build: kernel launch returns Err; no driver: cudarc panics and # catch_unwind treats it as skip). Reader-level assertions still run. - # Scoped to these two binaries — other gpu_*.rs tests lack probe-skip. + # pipeline_vram_guard (issue #1430) uses the same catch_unwind probe via + # common::qdp_engine_probed, so it skips rather than fails here. + # Scoped to these binaries — other gpu_*.rs tests lack probe-skip. - name: Cargo test (f32 Parquet CPU smoke) working-directory: qdp env: QDP_NO_CUDA: "1" - run: cargo test -p qdp-core --test parquet_f32 --test parquet_f32_fidelity + run: cargo test -p qdp-core --test parquet_f32 --test parquet_f32_fidelity --test pipeline_vram_guard + + # The pipeline_vram_guard cases above all self-skip without a device, so the + # memory-guard logic (issue #1430) is covered here instead: these unit tests + # are pure and need no GPU. Scoped to the module rather than a bare `--lib` + # because other qdp-core unit tests (e.g. encoding::amplitude's stream + # end-to-end case) do require a device and fail on the stub build. + - name: Cargo test (pipeline_runner CPU unit tests) + working-directory: qdp + env: + QDP_NO_CUDA: "1" + run: cargo test -p qdp-core --lib pipeline_runner::tests test: needs: rust-check diff --git a/docs/qdp/python-api.md b/docs/qdp/python-api.md index a9d8b2ec38..9f88107a33 100644 --- a/docs/qdp/python-api.md +++ b/docs/qdp/python-api.md @@ -21,6 +21,7 @@ from qumat_qdp import ( BACKEND, Backend, LatencyResult, + MemoryEstimate, NativeQuantumTensor, QdpBenchmark, QdpEngine, @@ -30,6 +31,7 @@ from qumat_qdp import ( RustQdpEngine, ThroughputResult, TritonAmdEngine, + estimate_memory, force_backend, is_triton_amd_available, run_throughput_pipeline_py, @@ -42,6 +44,7 @@ Key exports: - `QdpBenchmark`: benchmark builder for the Rust pipeline or explicit PyTorch reference backend. - `QuantumDataLoader`: batch iterator builder for synthetic or file-backed inputs. - `QdpTensor` / `QuantumTensor`: thin DLPack facade type. +- `estimate_memory` / `MemoryEstimate`: upfront host and device memory sizing for a pipeline configuration. - `Backend`, `BACKEND`, `force_backend`: backend detection and override helpers for the `_qdp` / PyTorch-reference selection layer. - `TritonAmdEngine`, `is_triton_amd_available`: direct AMD-route entry points. - `RustQdpEngine`, `NativeQuantumTensor`: native `_qdp` exports when the extension is available. @@ -246,6 +249,18 @@ File source notes: - The Rust backend supports broader file formats and streaming loaders. - The explicit PyTorch backend supports synthetic data plus `.npy`, `.pt`, and `.pth` files only. +Memory check notes: + +- The Rust backend estimates how much device memory a configuration needs and rejects it when the estimate exceeds free GPU memory, rather than failing part-way through a run. The check runs when iteration starts, before any input file is read — so a `QuantumDataLoader(...)` call returns normally and the rejection surfaces at the `for` statement. +- The error names the encoding, batch size, and qubit count alongside the requested and available memory, and the remedy is to lower `batch_size` or `qubits`. +- The rejection is raised as `RuntimeError`, not `ValueError`, because the loader wraps every backend error in `RuntimeError` — catch that and match on the message. `estimate_memory()` raises `ValueError` for the same oversized configuration. +- The estimate budgets two concurrent batch state buffers, because a batch stays resident on the device until the consumer releases its tensor. A configuration that fits one buffer but not two is rejected. +- The buffer is sized by the wider of the loader's input dtype and the engine's precision, since the input dtype sets the encode path's working precision while every path converts its result to the engine precision. With the defaults — file loaders read float64, `QdpEngine` runs float32 — it is the input dtype that sizes the budget. The error reports both values. +- Basis input read from a file is always budgeted as float64, because basis values are integer state indices that the file readers load as float64 whatever `dtype` is requested. Synthetic basis data is budgeted at the requested `dtype`. +- The check is skipped whenever the CUDA runtime reports no usable device — a build without the CUDA toolkit, a host with no driver, or an empty `CUDA_VISIBLE_DEVICES` — so CPU-only environments are unaffected. +- The comparison is against free memory sampled when iteration starts; nothing is reserved. Another process or a second loader can claim that memory before the first batch is allocated, so passing the check is a fast sanity check rather than a guarantee. +- To size a configuration before building the loader, call `estimate_memory()` (see [Memory Estimation API](#memory-estimation-api)). + Iteration behavior depends on backend: | Loader backend | Iteration yields | @@ -300,6 +315,38 @@ for qt in loader: batch = torch.from_dlpack(qt) ``` +## Memory Estimation API + +`estimate_memory()` answers the same question the loader's memory check answers when iteration starts, but without building anything — use it to size `qubits` and `batch_size` against a memory budget instead of discovering the ceiling from a rejection. + +Signature: + +`estimate_memory(num_qubits, batch_size, encoding_method="amplitude", dtype="f64", prefetch_depth=16)` + +Returns a `MemoryEstimate` with three fields, all in bytes: + +| Field | Meaning | +|-------|---------| +| `cpu_prefetch_bytes` | Host prefetch pool: `prefetch_depth` batches of raw, unencoded input | +| `gpu_state_bytes` | Device state-vector buffer, including the two-buffer allowance the memory check applies | +| `total_bytes` | Sum of the two | + +```python +from qumat_qdp import estimate_memory + +est = estimate_memory(num_qubits=20, batch_size=64, dtype="f32") +print(f"{est.gpu_state_bytes / 1024**2:.0f} MiB of device state") +``` + +Notes: + +- The function is pure configuration arithmetic: it allocates nothing and opens no device, so it works on a stub build or a host with no GPU — which is the point, since it exists to size configurations that the current machine may not be able to run. +- The device state vector holds `2**num_qubits` complex amplitudes per sample for every encoding, so `gpu_state_bytes` is identical for `amplitude` and `angle` at equal qubit counts even though their input widths differ. Only `cpu_prefetch_bytes` reflects the input width. +- `prefetch_depth` scales `cpu_prefetch_bytes` only. Device memory does not depend on it, which is why the memory check's error suggests lowering `batch_size` or `qubits` and never `prefetch_depth`. +- `gpu_state_bytes` is the figure the memory check compares against free VRAM, with one caveat: the check budgets the wider of this `dtype` and the engine's precision, and always float64 for `basis` read from a file. Pass the engine precision as `dtype` when the two differ, or the estimate will be half what the check applies. +- Every failure the estimator reports is an argument error — unknown encoding or dtype name, a `num_qubits` whose `2**n` state vector is not representable, or arithmetic overflow — and raises `ValueError`. A negative `num_qubits` or `batch_size` fails at the argument boundary first and raises `OverflowError`. `RuntimeError` is raised only when the `_qdp` extension is not installed. +- Encodings without an f32 batch path are estimated as float64 even when float32 is requested, mirroring what the pipeline does with the same request. + ## Low-level Rust Pipeline Helper `run_throughput_pipeline_py(...)` is the low-level native helper used by the Rust benchmark path. @@ -314,6 +361,8 @@ Returns a tuple: This helper is only available when `_qdp` is installed. +It builds its own engine and pipeline rather than going through `QuantumDataLoader`, so the memory check described under [Data Loader API](#data-loader-api) does not apply to it: an oversized configuration fails part-way through the run. Call `estimate_memory()` first if you are sizing a benchmark near the limits of the device. + ## Backward Compatibility `benchmark/api.py` and `benchmark/loader.py` continue to re-export the modern `qumat_qdp` API. Prefer importing from `qumat_qdp` directly. diff --git a/qdp/qdp-core/src/gpu/memory.rs b/qdp/qdp-core/src/gpu/memory.rs index f68461a6e2..99a6ba626d 100644 --- a/qdp/qdp-core/src/gpu/memory.rs +++ b/qdp/qdp-core/src/gpu/memory.rs @@ -92,18 +92,19 @@ fn build_oom_message( ) } -/// Guard that checks available GPU memory before attempting a large allocation. +/// The accept/reject decision itself, against caller-supplied `free`/`total` figures. /// -/// Returns a MemoryAllocation error with a helpful message when the request -/// exceeds the currently reported free memory. +/// Split out of [`ensure_device_memory_available`] so the decision can be exercised without a +/// device: every path that queries CUDA is unreachable on the stub build, which would otherwise +/// leave the rejection rule — the behavior issue #1430 is about — untested wherever CI runs. #[cfg(target_os = "linux")] -pub(crate) fn ensure_device_memory_available( +pub(crate) fn ensure_fits_in_free_memory( requested_bytes: usize, context: &str, qubits: Option, + free: usize, + total: usize, ) -> Result<()> { - let (free, total) = query_cuda_mem_info()?; - if (requested_bytes as u64) > (free as u64) { return Err(MahoutError::MemoryAllocation(build_oom_message( context, @@ -117,6 +118,20 @@ pub(crate) fn ensure_device_memory_available( Ok(()) } +/// Guard that checks available GPU memory before attempting a large allocation. +/// +/// Returns a MemoryAllocation error with a helpful message when the request +/// exceeds the currently reported free memory. +#[cfg(target_os = "linux")] +pub(crate) fn ensure_device_memory_available( + requested_bytes: usize, + context: &str, + qubits: Option, +) -> Result<()> { + let (free, total) = query_cuda_mem_info()?; + ensure_fits_in_free_memory(requested_bytes, context, qubits, free, total) +} + /// Wraps CUDA allocation errors with an OOM-aware MahoutError. #[cfg(target_os = "linux")] pub(crate) fn map_allocation_error( diff --git a/qdp/qdp-core/src/lib.rs b/qdp/qdp-core/src/lib.rs index 4454df9536..68bee37f7b 100644 --- a/qdp/qdp-core/src/lib.rs +++ b/qdp/qdp-core/src/lib.rs @@ -180,6 +180,15 @@ impl QdpEngine { &self.device } + /// Precision every encode path converts its result to before handing back a DLPack tensor. + /// + /// This — not the pipeline's host-input `dtype` — determines the element size of the state + /// buffer that stays resident on the device, because `encode*_for_pipeline` ends with + /// [`GpuStateVector::to_precision`](crate::gpu::memory::GpuStateVector::to_precision). + pub fn precision(&self) -> Precision { + self.precision + } + /// Block until all GPU work on the default stream has completed. /// Used by the generic pipeline and other callers that need to sync before timing. #[cfg(target_os = "linux")] diff --git a/qdp/qdp-core/src/pipeline_runner.rs b/qdp/qdp-core/src/pipeline_runner.rs index e9374dbf21..2d4ff5a7ff 100644 --- a/qdp/qdp-core/src/pipeline_runner.rs +++ b/qdp/qdp-core/src/pipeline_runner.rs @@ -24,7 +24,9 @@ use std::time::Instant; use crate::QdpEngine; use crate::dlpack::DLManagedTensor; use crate::error::{MahoutError, Result}; -use crate::gpu::memory::Precision; +use crate::estimate::estimate_memory; +use crate::gpu::cuda_ffi::cuda_runtime_available; +use crate::gpu::memory::{Precision, ensure_device_memory_available, ensure_fits_in_free_memory}; use crate::io; use crate::reader::{FloatElem, NullHandling, StreamingDataReader}; use crate::readers::ParquetStreamingReader; @@ -209,6 +211,168 @@ fn compute_optimal_prefetch_depth( } } +/// Element precision of the state buffer that stays resident on the device. +/// +/// Not simply `config.dtype`, which describes the *host* input: +/// +/// * Every encode path ends with `to_precision(engine.precision())`, so the resident buffer is +/// materialized at the engine's precision. The two are set independently — the Python synthetic +/// loader hardcodes an f32 config dtype while `QdpEngine(precision="float64")` is a documented +/// public option — and taking the narrower of the two would under-budget by 2x. +/// * Basis input is integer state indices, which the *file* loaders always read as f64 regardless +/// of `config.dtype` (see `read_file_by_extension` and the streaming loader's `use_f32`), so the +/// encoder produces an f64 state for basis whatever the config says. `basis_reads_f64` carries +/// that distinction: the synthetic producer fills basis batches at `config.dtype` (and +/// [`Encoding::supports_f32`] is true for basis, so `normalize` leaves an f32 request alone), +/// so forcing f64 there would over-budget by 2x and reject configurations that run fine. +/// +/// Taking the widest precision in play keeps the estimate on the conservative side of all three. +fn resident_device_precision( + config: &PipelineConfig, + engine_precision: Precision, + basis_reads_f64: bool, +) -> Precision { + let host_precision = if basis_reads_f64 && matches!(config.encoding, Encoding::Basis) { + Precision::Float64 + } else { + config.dtype + }; + match (host_precision, engine_precision) { + (Precision::Float32, Precision::Float32) => Precision::Float32, + _ => Precision::Float64, + } +} + +/// Names the configuration in the GPU-memory guard's error and log messages. +/// +/// Split out of [`ensure_config_fits_device`] so the wording can be asserted without a device. +/// `num_qubits` is deliberately absent: it is passed to `ensure_device_memory_available` +/// separately, which appends it as `(qubits=N)`. Both the requested host `dtype` and the +/// precision actually budgeted appear, because when they differ the second one explains the size. +fn vram_check_context(config: &PipelineConfig, device_precision: Precision) -> String { + format!( + "pipeline construction (encoding={}, batch_size={}, dtype={:?}, device_precision={:?}, \ + peak=2 concurrent batch buffers)", + config.encoding.as_str(), + config.batch_size, + config.dtype, + device_precision, + ) +} + +/// Reject a configuration whose GPU state buffers cannot fit in free device memory. +/// +/// Runs at iterator construction — before any host batch is allocated and before any input file is +/// read — so an oversized configuration fails in milliseconds with an actionable message instead of +/// running out of memory part-way through encoding. +/// +/// # What is compared +/// +/// [`estimate_memory`]'s `gpu_state_bytes`, which budgets **two** concurrent batch state buffers. +/// That factor is a real peak, not padding: [`to_dlpack`] clones the buffer `Arc`, +/// so a batch stays resident until the consumer releases the tensor — during a `for` loop's +/// `__next__` the previous batch is still alive while the next one is allocated — and +/// `encode_batch_for_pipeline`'s precision conversion holds source and destination buffers at once. +/// Comparing a single buffer would admit configurations that allocate successfully and then run out +/// of memory on the following batch, which is the late failure this guard exists to prevent. A +/// configuration that fits one buffer but not two is therefore rejected on purpose. +/// +/// The precision comes from [`resident_device_precision`], not `config.dtype`, so an engine and a +/// pipeline configured at different precisions are budgeted at the wider of the two. +/// +/// Two peaks are still not modeled, so this budget is a floor and not the true high-water mark: +/// +/// * The encoders upload their input batch to the device (`htod_sync_copy` in +/// `AmplitudeEncoder::encode_batch`, plus a per-sample norm buffer) and hold it alongside the +/// state buffers. For amplitude that input is half a state buffer, putting the real steady-state +/// peak near 2.5 buffers against the 2 budgeted here. This one applies to *every* configuration. +/// * While `to_precision` converts, the source buffer is alive alongside the destination, so a +/// mismatched precision pair briefly holds a third state buffer. +/// +/// Folding either in means teaching the estimator which encode path each config takes, which +/// belongs in the memory model — explicitly out of scope for issue #1430 — rather than in this +/// guard. The consequence is that a configuration sitting just under free memory can still OOM +/// mid-run; the guard narrows that window rather than closing it. +/// +/// `cpu_prefetch_bytes` is not compared against device memory: this guard covers device memory +/// only. Note that [`estimate_memory`] still overflow-checks the host figure first, so a config +/// whose device footprint fits can be rejected with [`MahoutError::InvalidInput`] for a host-side +/// overflow. Host limits themselves — including the Parquet streaming reader's chunk buffer, which +/// [`estimate_memory`] excludes from its model — are not checked anywhere yet. +/// +/// # When the check does not run +/// +/// `cudarc` links the CUDA *driver* API, while `cudaMemGetInfo` comes from the *runtime* API that +/// `build.rs` replaces with stubs when `nvcc` is absent. A driver-only host (the PyTorch-style +/// install called out in `build.rs`) can therefore hold a live [`QdpEngine`] while every runtime +/// call returns the unavailable sentinel. Probing with [`cuda_runtime_available`] first keeps +/// construction working there instead of failing on a query that cannot succeed. +/// +/// That probe also reports unavailable when the runtime is present but exposes no device — an +/// empty `CUDA_VISIBLE_DEVICES` — because it requires a device count above zero. Construction is +/// then allowed through unchecked, which is correct for CPU-only environments but means the +/// guard is silently inactive rather than merely permissive. +/// +/// The figures are sampled here and nothing is reserved, so a config that passes can still lose +/// the memory to another process before the first batch allocates. +/// +/// # Errors +/// +/// [`MahoutError::InvalidInput`] when the configuration overflows the memory model, +/// [`MahoutError::MemoryAllocation`] when the estimate exceeds free device memory, or +/// [`MahoutError::Cuda`] when the runtime reports itself available but the memory query then +/// fails — a broken runtime rather than an absent one, which every other allocation path in this +/// crate also surfaces rather than ignores. +/// +/// [`to_dlpack`]: crate::gpu::memory::GpuStateVector::to_dlpack +fn ensure_config_fits_device( + config: &PipelineConfig, + engine_precision: Precision, + basis_reads_f64: bool, +) -> Result<()> { + ensure_config_fits_device_with(config, engine_precision, basis_reads_f64, None) +} + +/// [`ensure_config_fits_device`] with the device memory figures injectable. +/// +/// `device_memory` is `None` on every production path, which probes the CUDA runtime and queries +/// it. `Some((free, total))` supplies the figures directly and skips the probe, so the +/// accept/reject rule — the behavior issue #1430 specifies — can be asserted on a build with no +/// CUDA runtime at all. Without this seam the rule is only reachable on a machine with a GPU, +/// which upstream CI is not: the estimate and the probe short-circuit would still be covered, but +/// deleting the comparison itself would not fail anything CI runs. +fn ensure_config_fits_device_with( + config: &PipelineConfig, + engine_precision: Precision, + basis_reads_f64: bool, + device_memory: Option<(usize, usize)>, +) -> Result<()> { + let device_precision = resident_device_precision(config, engine_precision, basis_reads_f64); + let estimate = estimate_memory( + config.encoding, + config.num_qubits, + config.batch_size, + device_precision, + config.prefetch_depth, + )?; + + if device_memory.is_none() && !cuda_runtime_available() { + log::debug!( + "CUDA runtime unavailable; skipping GPU memory check for {}", + vram_check_context(config, device_precision) + ); + return Ok(()); + } + + let requested = estimate.gpu_state_bytes as usize; + let context = vram_check_context(config, device_precision); + let qubits = Some(config.num_qubits as usize); + match device_memory { + Some((free, total)) => ensure_fits_in_free_memory(requested, &context, qubits, free, total), + None => ensure_device_memory_available(requested, &context, qubits), + } +} + pub struct SyntheticProducer { pub config: PipelineConfig, pub vector_len: usize, @@ -648,6 +812,11 @@ where impl PipelineIterator { pub fn new_synthetic(engine: QdpEngine, mut config: PipelineConfig) -> Result { config.normalize(); + // Before spawn_producer, which starts a thread that allocates batch_size * vector_len of + // host memory on its first call. Note the guard only covers the device side, so on builds + // where it short-circuits that host allocation is still unchecked. + // basis_reads_f64 = false: SyntheticProducer fills basis batches at config.dtype. + ensure_config_fits_device(&config, engine.precision(), false)?; let vector_len = vector_len(config.num_qubits, config.encoding); let producer = SyntheticProducer::new(config.clone(), vector_len); let prefetch_depth = config.prefetch_depth; @@ -673,6 +842,9 @@ impl PipelineIterator { batch_limit: usize, ) -> Result { config.normalize(); + // Before read_file_by_extension: that call loads the whole file, so checking afterwards + // would trade the sub-second rejection for a full read of data we are about to discard. + ensure_config_fits_device(&config, engine.precision(), true)?; let path = path.as_ref(); let (batch_data, num_samples, sample_size) = read_file_by_extension(path, config.null_handling, config.dtype, config.encoding)?; @@ -724,6 +896,8 @@ impl PipelineIterator { batch_limit: usize, ) -> Result { config.normalize(); + // Before the reader opens the file and reads its first chunk. + ensure_config_fits_device(&config, engine.precision(), true)?; let path = path.as_ref(); if path_extension_lower(path).as_deref() != Some("parquet") { return Err(MahoutError::InvalidInput(format!( @@ -1033,6 +1207,211 @@ pub fn run_latency_pipeline(config: &PipelineConfig) -> Result PipelineConfig { + PipelineConfig { + num_qubits: 20, + batch_size: 8, + encoding: Encoding::Amplitude, + dtype: Precision::Float32, + prefetch_depth: 1, + ..PipelineConfig::default() + } + } + + const REJECTION_TEST_REQUIRED_BYTES: usize = 2 * 8 * (1 << 20) * 8; + + /// The behavior issue #1430 specifies: a configuration whose estimate exceeds free VRAM is + /// rejected. Injecting the memory figures keeps this reachable on a build with no CUDA + /// runtime, so the rule is covered where CI actually runs rather than only on a GPU host. + #[test] + fn config_larger_than_free_memory_is_rejected() { + let config = rejection_test_config(); + let free = REJECTION_TEST_REQUIRED_BYTES - 1; + + let err = ensure_config_fits_device_with( + &config, + Precision::Float32, + true, + Some((free, REJECTION_TEST_REQUIRED_BYTES * 4)), + ) + .expect_err("estimate exceeds free memory by one byte, so this must be rejected"); + + assert!( + matches!(err, MahoutError::MemoryAllocation(_)), + "expected MemoryAllocation, got: {err}" + ); + // The remedy has to be actionable: #1430 asks for the offending values by name. + let message = err.to_string(); + for expected in ["amplitude", "batch_size=8", "qubits=20", "Reduce qubits"] { + assert!( + message.contains(expected), + "rejection message missing {expected:?}: {message}" + ); + } + } + + /// The other side of the same rule: exactly enough free memory is accepted. Pins the + /// comparison as `requested > free` rather than `>=`, so a config that fits precisely is not + /// turned away. + #[test] + fn config_that_exactly_fits_free_memory_is_accepted() { + let config = rejection_test_config(); + + assert!( + ensure_config_fits_device_with( + &config, + Precision::Float32, + true, + Some(( + REJECTION_TEST_REQUIRED_BYTES, + REJECTION_TEST_REQUIRED_BYTES * 4 + )), + ) + .is_ok(), + "an estimate equal to free memory fits and must be accepted" + ); + } + fn assert_generate_and_inplace_match(encoding_method: &str) { let config = PipelineConfig { num_qubits: 5, diff --git a/qdp/qdp-core/tests/common/mod.rs b/qdp/qdp-core/tests/common/mod.rs index 85b0183126..928a1af7fe 100644 --- a/qdp/qdp-core/tests/common/mod.rs +++ b/qdp/qdp-core/tests/common/mod.rs @@ -177,6 +177,36 @@ pub fn qdp_engine_with_precision(precision: Precision) -> Option { QdpEngine::new_with_precision(0, precision).ok() } +/// Returns a QDP engine, or `None` when one cannot be created — including on hosts with no +/// NVIDIA driver at all. +/// +/// [`qdp_engine`] is not enough for that case: when `libcuda` cannot be dlopened, `cudarc` +/// panics (`panic_no_lib_found`) instead of returning `Err`, so `.ok()` never yields `None` and +/// the test aborts rather than skipping. Probing inside `catch_unwind` with the default hook +/// suppressed borrows that much from `parquet_f32_fidelity.rs`'s helper for issue #1342. +/// +/// It stops short of what that helper does, and the difference matters. A stub build — toolkit +/// absent, `libcuda` present — creates an engine *successfully* and only fails when a kernel is +/// launched, so this probe returns `Some` there. That is sufficient for callers that never launch +/// a kernel, which is why the memory-guard tests can use it: they assert on the guard's decision +/// before any encode runs. A test that does launch one needs `parquet_f32_fidelity.rs`'s version, +/// which probes with a trivial 1-qubit encode; using this one instead reproduces #1342. +#[cfg(target_os = "linux")] +#[allow(dead_code)] +pub fn qdp_engine_probed() -> Option { + // The panic hook is process-global, so serialize probes against each other. This narrows the + // suppression window but does not close it: the harness runs tests on parallel threads, so a + // genuine panic elsewhere that lands inside this window still loses its message. Same hazard + // as the `parquet_f32_fidelity.rs` helper this borrows from. + static PROBE_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); + let _guard = PROBE_LOCK.lock().unwrap_or_else(|e| e.into_inner()); + let prev_hook = std::panic::take_hook(); + std::panic::set_hook(Box::new(|_| {})); + let probe = std::panic::catch_unwind(|| QdpEngine::new(0).ok()); + std::panic::set_hook(prev_hook); + probe.ok().flatten() +} + /// Copies f64 host data to the default CUDA device, or returns `None` when unavailable. #[cfg(target_os = "linux")] #[allow(dead_code)] diff --git a/qdp/qdp-core/tests/pipeline_vram_guard.rs b/qdp/qdp-core/tests/pipeline_vram_guard.rs new file mode 100644 index 0000000000..6ea78184b7 --- /dev/null +++ b/qdp/qdp-core/tests/pipeline_vram_guard.rs @@ -0,0 +1,205 @@ +// +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Construction-time GPU memory guard (issue #1430). +//! +//! Three probes are in play and they are not interchangeable: +//! +//! * `common::qdp_engine_probed()` — CUDA *driver* API, inside `catch_unwind` because `cudarc` +//! panics rather than returning `Err` when `libcuda` is absent. Answers "can an engine exist". +//! * `qdp_core::cuda_runtime_available()` — CUDA *runtime* API, stubbed out when `nvcc` is absent. +//! Answers "will the guard actually query device memory". +//! * Neither implies the other. A `QDP_NO_CUDA` build on a machine with a driver satisfies the +//! first and fails the second, which is exactly the configuration where the guard steps aside. + +mod common; + +#[cfg(target_os = "linux")] +use qdp_core::{Encoding, PipelineConfig, PipelineIterator, Precision, QdpEngine}; +#[cfg(target_os = "linux")] +use std::time::{Duration, Instant}; + +/// Amplitude, 30 qubits (the `MAX_QUBITS` ceiling), batch 64, f32 — far past any real card. +#[cfg(target_os = "linux")] +fn oversized_config() -> PipelineConfig { + PipelineConfig { + num_qubits: 30, + batch_size: 64, + total_batches: 1, + encoding: Encoding::Amplitude, + dtype: Precision::Float32, + prefetch_depth: 1, + ..PipelineConfig::default() + } +} + +/// An engine plus a working CUDA runtime, or `None` to skip. +/// +/// Both are required before a test may feed the guard an oversized config: without the runtime +/// API the guard steps aside by design, and the synthetic producer would then try to allocate +/// 2^30 * 64 f32 values on its own thread and abort the test process. +#[cfg(target_os = "linux")] +fn engine_with_live_runtime() -> Option { + let engine = common::qdp_engine_probed()?; + if !qdp_core::cuda_runtime_available() { + println!("SKIP: CUDA runtime unavailable (stub build)"); + return None; + } + Some(engine) +} + +#[test] +#[cfg(target_os = "linux")] +fn oversized_config_rejected_before_allocation() { + let Some(engine) = engine_with_live_runtime() else { + return; + }; + + let start = Instant::now(); + let result = PipelineIterator::new_synthetic(engine, oversized_config()); + let elapsed = start.elapsed(); + + let message = result + .err() + .expect("a 30-qubit batch-64 state buffer cannot fit in any current device") + .to_string(); + + assert!( + message.contains("Reduce qubits or batch size"), + "message must name a remedy that works, got: {message}" + ); + assert!( + !message.contains("prefetch_depth"), + "prefetch_depth does not affect gpu_state_bytes and must not be suggested, got: {message}" + ); + for expected in ["encoding=amplitude", "batch_size=64", "qubits=30"] { + assert!( + message.contains(expected), + "message must name {expected}, got: {message}" + ); + } + + // Pins the deliberately conservative budget: the guard compares two concurrent batch buffers, + // not the single buffer `GpuStateVector::new_batch` allocates, because a DLPack tensor held + // across `__next__` keeps the previous batch resident. One f32 batch buffer here is + // 64 * 2^30 complex64 (8 bytes) = 512 GiB, so the guard must request twice that. Asserting the + // exact figure needs no knowledge of the host's free memory. + let single_buffer_mib = 64.0 * (1u64 << 30) as f64 * 8.0 / (1024.0 * 1024.0); + let expected = format!("requested {:.2} MiB", 2.0 * single_buffer_mib); + assert!( + message.contains(&expected), + "expected the two-buffer budget ({expected}), got: {message}" + ); + + // The point of the guard is that rejection is immediate rather than discovered mid-run. + assert!( + elapsed < Duration::from_secs(1), + "rejection took {elapsed:?}, expected well under a second" + ); +} + +/// The file loaders must reject before touching the filesystem. Pointing them at a path that does +/// not exist separates the two outcomes: the guard's memory error proves it ran first, while a +/// missing-file error would mean the check had been skipped or wired in too late. +#[test] +#[cfg(target_os = "linux")] +fn file_loaders_reject_before_reading_the_file() { + let Some(engine) = engine_with_live_runtime() else { + return; + }; + let missing = std::path::Path::new("/nonexistent/qdp-vram-guard/never-created.parquet"); + + let full_read = + PipelineIterator::new_from_file(engine.clone(), missing, oversized_config(), usize::MAX) + .err() + .expect("new_from_file must reject the oversized config") + .to_string(); + assert!( + full_read.contains("Reduce qubits or batch size"), + "new_from_file must fail on memory before opening the file, got: {full_read}" + ); + + let streaming = + PipelineIterator::new_from_file_streaming(engine, missing, oversized_config(), usize::MAX) + .err() + .expect("new_from_file_streaming must reject the oversized config") + .to_string(); + assert!( + streaming.contains("Reduce qubits or batch size"), + "new_from_file_streaming must fail on memory before opening the file, got: {streaming}" + ); +} + +/// An f32 pipeline on an f64 engine holds complex128 buffers on the device, so the guard must +/// budget f64. Budgeting the config's dtype would halve the estimate and admit a config that then +/// runs out of memory mid-run. Reachable from Python: the synthetic loader hardcodes an f32 config +/// dtype while `QdpEngine(precision="float64")` is a documented option. +#[test] +#[cfg(target_os = "linux")] +fn engine_precision_widens_the_budget() { + if common::qdp_engine_probed().is_none() || !qdp_core::cuda_runtime_available() { + println!("SKIP: no engine or CUDA runtime unavailable"); + return; + } + let Some(engine_f64) = common::qdp_engine_with_precision(Precision::Float64) else { + println!("SKIP: no f64 engine available"); + return; + }; + + let message = PipelineIterator::new_synthetic(engine_f64, oversized_config()) + .err() + .expect("oversized config must be rejected") + .to_string(); + + assert!( + message.contains("dtype=Float32") && message.contains("device_precision=Float64"), + "message must show the f32 request budgeted as f64, got: {message}" + ); + // Twice the f32 figure asserted above: complex128 is 16 bytes per element. + let single_buffer_mib = 64.0 * (1u64 << 30) as f64 * 16.0 / (1024.0 * 1024.0); + let expected = format!("requested {:.2} MiB", 2.0 * single_buffer_mib); + assert!( + message.contains(&expected), + "expected the f64 budget ({expected}), got: {message}" + ); +} + +/// A configuration that fits must still construct. On a stub CUDA runtime this is the +/// graceful-degradation case: the engine exists via the driver API, the guard short-circuits, and +/// construction proceeds instead of failing on a memory query that cannot succeed. +#[test] +#[cfg(target_os = "linux")] +fn modest_config_constructs_on_any_build() { + let Some(engine) = common::qdp_engine_probed() else { + println!("SKIP: No GPU available"); + return; + }; + + let config = PipelineConfig { + num_qubits: 10, + batch_size: 8, + total_batches: 1, + encoding: Encoding::Amplitude, + dtype: Precision::Float32, + prefetch_depth: 1, + ..PipelineConfig::default() + }; + + assert!( + PipelineIterator::new_synthetic(engine, config).is_ok(), + "a 10-qubit batch-8 pipeline must not be rejected by the memory guard" + ); +} diff --git a/qdp/qdp-python/qumat_qdp/__init__.py b/qdp/qdp-python/qumat_qdp/__init__.py index 81ffd66347..e8dc8c0303 100644 --- a/qdp/qdp-python/qumat_qdp/__init__.py +++ b/qdp/qdp-python/qumat_qdp/__init__.py @@ -19,12 +19,14 @@ Public API: QdpEngine (unified router), QdpTensor/QuantumTensor (DLPack facade), QdpBenchmark, ThroughputResult, LatencyResult (benchmark API), -QuantumDataLoader (data loader iterator). +QuantumDataLoader (data loader iterator), +estimate_memory/MemoryEstimate (upfront memory sizing). Usage: from qumat_qdp import QdpEngine, QuantumTensor from qumat_qdp import QdpBenchmark, ThroughputResult, LatencyResult from qumat_qdp import QuantumDataLoader + from qumat_qdp import estimate_memory """ from __future__ import annotations @@ -71,6 +73,7 @@ def is_cuda_available() -> bool: ThroughputResult, ) from qumat_qdp.backend import QdpEngine +from qumat_qdp.estimate import MemoryEstimate, estimate_memory from qumat_qdp.loader import QuantumDataLoader from qumat_qdp.tensor import QdpTensor, QuantumTensor from qumat_qdp.triton_amd import TritonAmdEngine, is_triton_amd_available @@ -79,6 +82,7 @@ def is_cuda_available() -> bool: "BACKEND", "Backend", "LatencyResult", + "MemoryEstimate", "NativeQuantumTensor", "QdpBenchmark", "QdpEngine", @@ -88,6 +92,7 @@ def is_cuda_available() -> bool: "RustQdpEngine", "ThroughputResult", "TritonAmdEngine", + "estimate_memory", "force_backend", "is_cuda_available", "is_triton_amd_available", diff --git a/qdp/qdp-python/qumat_qdp/estimate.py b/qdp/qdp-python/qumat_qdp/estimate.py new file mode 100644 index 0000000000..3c062c14ba --- /dev/null +++ b/qdp/qdp-python/qumat_qdp/estimate.py @@ -0,0 +1,113 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Upfront memory estimation for QDP pipeline configurations. + +``QuantumDataLoader`` rejects a configuration whose device buffers cannot fit in +free VRAM when iteration starts. This module answers the same question *before* +building anything, so a caller can size ``num_qubits``/``batch_size`` against a +memory budget rather than discovering the ceiling from a rejection. + +Usage:: + + from qumat_qdp import estimate_memory + + est = estimate_memory(num_qubits=20, batch_size=64, dtype="f32") + print(est.gpu_state_bytes / 1024**2, "MiB of device state") +""" + +from __future__ import annotations + +from dataclasses import dataclass + +from qumat_qdp._backend import get_qdp + +_NO_EXTENSION = ( + "Memory estimation requires the native QDP extension (_qdp), which is not " + "available. Build it with: uv run --active maturin develop --manifest-path " + "qdp/qdp-python/Cargo.toml" +) + + +@dataclass(frozen=True) +class MemoryEstimate: + """Estimated memory footprint of a pipeline configuration. + + Returned by :func:`estimate_memory`. Every field is an upper-bound estimate + derived from configuration arithmetic alone -- nothing here reflects an actual + allocation, and no device is consulted. + """ + + cpu_prefetch_bytes: int + """Host prefetch pool: ``prefetch_depth`` batches of raw, unencoded input.""" + + gpu_state_bytes: int + """Device state-vector buffer, including a double-buffering allowance.""" + + total_bytes: int + """Combined host and device footprint.""" + + +def estimate_memory( + num_qubits: int, + batch_size: int, + encoding_method: str = "amplitude", + dtype: str = "f64", + prefetch_depth: int = 16, +) -> MemoryEstimate: + """Estimate the host and device memory a pipeline configuration would need. + + ``gpu_state_bytes`` is the figure the loader's guard compares against free VRAM + when iteration starts. The guard budgets the wider of ``dtype`` and the engine's + precision, and always f64 for ``"basis"`` read from a file (whose inputs are + integer state indices read as f64; synthetic basis data is budgeted at the + requested ``dtype``), so pass the engine precision as ``dtype`` when the two + differ -- otherwise this estimate is half the one the guard applies. + + :param num_qubits: Qubit count; the device state vector holds ``2**num_qubits`` + complex amplitudes per sample. + :param batch_size: Samples per batch. + :param encoding_method: ``"amplitude"``, ``"angle"``, ``"basis"``, ``"iqp"``, + ``"iqp-z"``, or ``"phase"`` (case-insensitive). + :param dtype: ``"float32"``/``"f32"`` or ``"float64"``/``"f64"`` + (case-insensitive). Encodings without an f32 batch path are estimated as + f64 regardless, mirroring what the pipeline does with the same request. + :param prefetch_depth: Host prefetch queue depth. Affects + ``cpu_prefetch_bytes`` only; device buffers do not scale with it. + :returns: The :class:`MemoryEstimate` for this configuration. + :raises ValueError: If a name is unrecognized, ``2**num_qubits`` is not + representable, or the arithmetic overflows. + :raises OverflowError: If ``num_qubits`` or ``batch_size`` is negative, which + fails at the argument boundary before the estimator runs. + :raises RuntimeError: If the native extension is not available. + """ + qdp = get_qdp() + native = getattr(qdp, "estimate_memory", None) if qdp is not None else None + if native is None: + raise RuntimeError(_NO_EXTENSION) + + cpu_prefetch_bytes, gpu_state_bytes, total_bytes = native( + num_qubits=num_qubits, + batch_size=batch_size, + encoding_method=encoding_method, + dtype=dtype, + prefetch_depth=prefetch_depth, + ) + return MemoryEstimate( + cpu_prefetch_bytes=cpu_prefetch_bytes, + gpu_state_bytes=gpu_state_bytes, + total_bytes=total_bytes, + ) diff --git a/qdp/qdp-python/src/lib.rs b/qdp/qdp-python/src/lib.rs index 0bfd77189e..bdbe87845c 100644 --- a/qdp/qdp-python/src/lib.rs +++ b/qdp/qdp-python/src/lib.rs @@ -22,7 +22,7 @@ mod pytorch; mod tensor; use engine::QdpEngine; -use pyo3::exceptions::PyRuntimeError; +use pyo3::exceptions::{PyRuntimeError, PyValueError}; use pyo3::prelude::*; use tensor::QuantumTensor; @@ -68,6 +68,49 @@ fn run_throughput_pipeline_py( )) } +/// Estimate the host and device memory a pipeline configuration would need. +/// +/// Returns ``(cpu_prefetch_bytes, gpu_state_bytes, total_bytes)``. Pure config +/// arithmetic: allocates nothing, touches no device, and is therefore callable on a +/// stub build or a host with no GPU -- which is the point, since the intended use is +/// sizing a configuration before building a loader that would reject it. +/// +/// ``gpu_state_bytes`` is the figure the loader's memory guard compares against free +/// VRAM when iteration starts, with one caveat worth stating here: the guard budgets +/// the *wider* of this ``dtype`` and the engine's precision (and always f64 for +/// ``basis`` read from a file, whose inputs are integer state indices read as f64; +/// synthetic basis data is budgeted at the requested ``dtype``). Pass the +/// engine's precision as ``dtype`` when the two differ, or the estimate will be half +/// what the guard uses. +/// +/// Every failure the estimator itself reports is a bad argument -- an unknown encoding +/// or dtype name, a ``num_qubits`` whose 2^n state vector is not representable, or a +/// product that overflows -- and raises ``ValueError``. Arguments that cannot be +/// converted at the boundary fail earlier and differently: a negative ``num_qubits`` or +/// ``batch_size`` raises ``OverflowError``, a non-integer raises ``TypeError``. +#[pyfunction] +#[pyo3(signature = (num_qubits, batch_size, encoding_method="amplitude", dtype="f64", prefetch_depth=16))] +fn estimate_memory( + num_qubits: u32, + batch_size: usize, + encoding_method: &str, + dtype: &str, + prefetch_depth: usize, +) -> PyResult<(u64, u64, u64)> { + let encoding = qdp_core::Encoding::from_str_ci(encoding_method) + .map_err(|e| PyValueError::new_err(format!("Invalid encoding_method: {e}")))?; + let dtype = qdp_core::Dtype::from_str_ci(dtype) + .map_err(|e| PyValueError::new_err(format!("Invalid dtype: {e}")))?; + let estimate = + qdp_core::estimate_memory(encoding, num_qubits, batch_size, dtype, prefetch_depth) + .map_err(|e| PyValueError::new_err(e.to_string()))?; + Ok(( + estimate.cpu_prefetch_bytes, + estimate.gpu_state_bytes, + estimate.total(), + )) +} + /// Returns ``True`` if a usable CUDA device is available to the native engine. /// /// This reflects whether GPU work can actually run -- it is ``False`` for a @@ -91,6 +134,7 @@ fn _qdp(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_function(wrap_pyfunction!(cuda_available, m)?)?; + m.add_function(wrap_pyfunction!(estimate_memory, m)?)?; #[cfg(target_os = "linux")] m.add_class::()?; #[cfg(target_os = "linux")] diff --git a/testing/qdp/test_estimate_memory.py b/testing/qdp/test_estimate_memory.py new file mode 100644 index 0000000000..ea44417055 --- /dev/null +++ b/testing/qdp/test_estimate_memory.py @@ -0,0 +1,162 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the Python memory-estimation API (issue #1430). + +Every case here is pure arithmetic on the Rust side: no device is opened and nothing +is allocated, so these run identically on a GPU host and on a stub CUDA build with no +GPU at all. That is the point of the API -- it exists to size a configuration on a +machine that could not run it. +""" + +from __future__ import annotations + +import pytest + +from .qdp_test_utils import requires_qdp + +pytestmark = requires_qdp + +MIB = 1024 * 1024 + + +@pytest.fixture +def estimate_memory(): + """The public ``qumat_qdp`` wrapper, which is what users are told to import.""" + from qumat_qdp import estimate_memory as fn + + return fn + + +def test_known_configuration_matches_the_documented_formulas(estimate_memory): + """16 qubits, batch 64, f32, depth 16 -- the worked example in the Rust docs. + + Pinning one exact configuration catches a change of formula, of unit, or of field + order in the binding's return tuple, none of which a relative assertion would. + """ + est = estimate_memory( + num_qubits=16, batch_size=64, encoding_method="amplitude", dtype="f32" + ) + + # 16 * 64 * 2**16 samples * 4 bytes + assert est.cpu_prefetch_bytes == 256 * MIB + # 2 concurrent buffers * 64 * 2**16 amplitudes * 8 bytes (complex64) + assert est.gpu_state_bytes == 64 * MIB + assert est.total_bytes == 320 * MIB + + +def test_f64_doubles_every_field(estimate_memory): + """Precision is the knob most likely to be mis-wired: f64 is twice f32 throughout.""" + f32 = estimate_memory(num_qubits=12, batch_size=16, dtype="f32") + f64 = estimate_memory(num_qubits=12, batch_size=16, dtype="f64") + + assert f64.cpu_prefetch_bytes == 2 * f32.cpu_prefetch_bytes + assert f64.gpu_state_bytes == 2 * f32.gpu_state_bytes + assert f64.total_bytes == 2 * f32.total_bytes + + +def test_prefetch_depth_moves_host_memory_only(estimate_memory): + """The guard's remedy names batch size and qubits, never prefetch_depth. + + That advice is only correct if device memory really is independent of depth. + """ + shallow = estimate_memory(num_qubits=10, batch_size=8, prefetch_depth=1) + deep = estimate_memory(num_qubits=10, batch_size=8, prefetch_depth=8) + + assert deep.cpu_prefetch_bytes == 8 * shallow.cpu_prefetch_bytes + assert deep.gpu_state_bytes == shallow.gpu_state_bytes + + +def test_batch_size_scales_device_memory(estimate_memory): + """Halving batch size halves the device budget -- the remedy the guard suggests.""" + big = estimate_memory(num_qubits=10, batch_size=64) + small = estimate_memory(num_qubits=10, batch_size=32) + + assert big.gpu_state_bytes == 2 * small.gpu_state_bytes + + +def test_device_state_is_2_to_the_n_whatever_the_input_width(estimate_memory): + """Angle input is n values per sample, but the device still holds 2**n amplitudes. + + So angle's host pool is far smaller than amplitude's while their device + footprints are identical -- a caller sizing only by input width would be wrong. + """ + amplitude = estimate_memory( + num_qubits=14, batch_size=8, encoding_method="amplitude" + ) + angle = estimate_memory(num_qubits=14, batch_size=8, encoding_method="angle") + + assert angle.gpu_state_bytes == amplitude.gpu_state_bytes + assert angle.cpu_prefetch_bytes < amplitude.cpu_prefetch_bytes + + +def test_the_estimate_grows_past_any_real_device(estimate_memory): + """The rejection case: 30 qubits at batch 64 is ~1 TiB of state, on purpose. + + This is the configuration the construction guard turns into an immediate error; + here it is just a number, which is what makes it checkable without a GPU. + """ + est = estimate_memory( + num_qubits=30, batch_size=64, encoding_method="amplitude", dtype="f32" + ) + + assert est.gpu_state_bytes == 2 * 64 * (1 << 30) * 8 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"encoding_method": "quantum-teapot"}, id="unknown-encoding"), + pytest.param({"dtype": "float128"}, id="unknown-dtype"), + pytest.param({"num_qubits": 64}, id="state-vector-not-representable"), + pytest.param({"batch_size": 2**62, "num_qubits": 30}, id="overflow"), + ], +) +def test_bad_arguments_raise_value_error(estimate_memory, kwargs): + """All four failure modes are argument errors, so all four must be ValueError. + + A RuntimeError here would read as "the device failed", which is never the cause: + nothing has been allocated at this point. + """ + args = {"num_qubits": 16, "batch_size": 64} + args.update(kwargs) + + with pytest.raises(ValueError): + estimate_memory(**args) + + +def test_names_are_case_insensitive_like_the_rest_of_the_api(estimate_memory): + """Loader and engine accept 'Float32'/'AMPLITUDE'; this must not be the exception.""" + upper = estimate_memory( + num_qubits=10, batch_size=4, encoding_method="AMPLITUDE", dtype="Float32" + ) + lower = estimate_memory( + num_qubits=10, batch_size=4, encoding_method="amplitude", dtype="f32" + ) + + assert upper == lower + + +def test_estimate_is_exported_from_the_package_root(): + """Users are told to ``from qumat_qdp import estimate_memory``; keep that true.""" + import qumat_qdp + + assert "estimate_memory" in qumat_qdp.__all__ + assert "MemoryEstimate" in qumat_qdp.__all__ + assert isinstance( + qumat_qdp.estimate_memory(num_qubits=4, batch_size=2), + qumat_qdp.MemoryEstimate, + )