[Feature][QDP] One encode path across Python, Rust and CUDA (tracking)

[ryankert01]

TL;DR

Adding one encoding to QDP today means writing it three times: a CUDA host launcher plus extern "C" declaration plus no-CUDA stub, an 8-method Rust encoder (f32/f64 × single/batch × host/device), and a hand-matched PyO3 branch. #1268 added one f32 variant of angle encoding and touched 9 files across 2 languages for 1,135 lines.

I refactored this so an encoding is one device-only .cu file, one ~80-line Rust descriptor, and one PyTorch reference function. The data-movement engine (dual-stream pinned upload, prefetch, buffer pool, DLPack) is untouched, the public Python API is unchanged, and the repo's standard benchmarks show no measurable throughput or latency change.

What

Six commits on a branch, each independently green (Rust tests, clippy, no-CUDA build, Python suites including a new parity grid). I'll open them as PRs against this issue:

• Parity grid + baseline tool — one test over encoding × precision × shape × input location against torch_ref.py (which now covers phase); benchmark/baseline.py captures and compares throughput/latency with a threshold.
• Runtime kernel loading — build.rs compiles each .cu to a device-only fatbin and embeds it; qdp_kernels::registry loads through the CUDA driver API on first use and exposes one generic launch + kernel_args!. A test proves every embedded symbol resolves.
• Encoder collapse — a Kernel trait with a single launch on device-resident data, Input (host slice or device pointer on a stream), and QdpEngine::encode(input, shape, num_qubits, encoding) replacing 24 named methods (old names kept as one-line shorthands in compat.rs). Bindings become PyInput::from_py + one call.
• Host-path fixes found while benchmarking (amplitude CPU finite pass, basis device cast).
• Remove launchers — 26 host launchers, 26 declarations, 26 stubs, the static nvcc compile and the cc dependency are gone.
• Docs — DEVELOPMENT.md architecture section, a "which files do I touch" table, the three-file recipe, make check-kernels, make parity.

Surface	Before	After
CUDA host launchers / extern "C" decls / no-CUDA stubs	26 / 26 / 26	0
qdp-kernels/src/lib.rs	726 lines	57 lines
CUDA sources (device code only after)	2,633 lines	1,434 lines
Encoder trait methods per encoding	8	1
Angle encoder	868 lines	89 lines
Public encode* methods on QdpEngine	24	1 (+ shorthands)
PyO3 engine methods	11	6
Encodings that stream from Parquet	3	6
Encodings that read a CUDA tensor in place	5	6

Net: 62 files, +4,334 / −9,974 lines. Bonus behaviour: phase now encodes from CUDA tensors, basis accepts float64 CUDA tensors, streaming Parquet works for all six encodings, NumPy float32 input is accepted.

Why

Contributors should be able to add or tune an encoding knowing CUDA and a little Rust, add a reader knowing Rust, and add a loader feature knowing Python. Today every encoder change needs all three, and the FFI signatures are matched by eye (a mismatch is UB, not a compile error).

How

Performance. Old (pre-refactor commit) vs new, built into separate venvs, run with benchmark_throughput.py, benchmark_latency.py and benchmark_e2e.py at their default arguments, 5 alternating rounds on an idle RTX 6000 Ada. Medians, with new/old ratio:

Encoding	Throughput old → new (vec/s)	Latency old → new (ms/vec)
amplitude	12,660 → 5,667 (0.45, see below)	0.096 → 0.097 (1.01)
angle	23,110 → 22,993 (0.99)	0.043 → 0.044 (1.02)
basis	68,651 → 77,569 (1.13)	0.013 → 0.015 (1.15)
iqp	22,140 → 22,045 (1.00)	0.046 → 0.045 (0.98)
iqp-z	49,597 → 51,432 (1.04)	0.020 → 0.020 (1.00)

The amplitude throughput row is benchmark noise, not code: that benchmark alternates between ~11k and ~5.5k vec/s on consecutive runs of the same build (old: 10,563 5,665 11,367 5,694 11,434 5,606; new: 5,973 12,592 5,761 10,385 5,569 5,390). Amplitude latency on the same workload is identical, and an nsys trace shows the three amplitude kernels at the same grid and the same time on both builds (69.0 / 29.7 / 2.5 µs per batch), with the batch dominated by the pageable 16 MB host copy. Per-call amplitude encode, interleaved: 3.73 vs 3.76 ms and 3.60 vs 3.71 ms.

Behaviour changes to note in review

• Parquet streaming now honours engine precision (previously always complex128).
• Error message wording was preserved where tests asserted it.

Follow-ups (separate issues once this lands)

• Delete the compat.rs shorthands once callers use encode.
• Delete the now-unused single-sample kernels from the .cu files.
• Consider pinned producer buffers: the amplitude pipeline is bound by a pageable host copy per batch, not by kernels.

[ryankert01]

Update: review fixes landed, re-measured

A /code-review pass on the branch found ten issues; all are fixed in three follow-up commits (1546980, bbc1a27, 2acf683):

1. Stale kernel handles — the registry cached CUmodule/CUfunction handles for the life of the process, but they die with the primary context when the last CudaDevice for a GPU is dropped. Reproduced with cargo test -- --test-threads=1 (CUDA_ERROR_INVALID_HANDLE). Handles are now cached per primary context and reloaded transparently when stale. Both parallel and single-threaded Rust runs pass.
2. Launches are asynchronous again — Kernel::launch returns a Pending (buffers the queued kernels read + checks that need device results); validation flags are read once after a single stream sync. This restores copy/compute overlap on the Parquet streaming path for amplitude and basis.
3. Single-sample amplitude uses the wide single-sample reduction again, and small single host samples take their norm on the CPU with one kernel and one sync, as on main.
4. f32 host input under a float64 engine is widened on the CPU rather than encoded at f32 accuracy.
5. Parquet streaming encodes into the engine's precision through a bounded float64 staging buffer (no second full-size state); all-null/zero rows stay zero rows as on main.
6. build.rs fails loudly on kernel declarations it cannot parse and accepts __launch_bounds__; unknown symbols are reported by the driver.
7. Restored coverage for the misaligned odd-sample-length amplitude batch path; added finalize-flag and single-sample reduction tests.
8. Removed the per-encoding f32 device-pointer shorthands from compat.rs (no callers outside tests; tests ported to encode).
9. Encoding::vector_len delegates to Kernel::sample_size; stale docs removed.

Verification: Rust workspace tests pass parallel and --test-threads=1; clippy, rustfmt and the no-CUDA build are clean; 378 Python tests pass including the parity grid.

Performance, re-measured (old = b403d5e, new = branch head; idle RTX A6000, 5 alternating rounds of the repo scripts with default args, a guard that waits for an empty GPU before every measurement; 0 measurements dropped):

Encoding	Throughput old → new (vec/s, median)	Latency old → new (ms/vec, median)
amplitude	9,932 → 4,654 (mode lottery, see below)	0.223 → 0.187
angle	4,676 → 4,832	0.118 → 0.114
basis	24,585 → 42,180	0.029 → 0.018
iqp	3,715 → 3,806	0.271 → 0.266
iqp-z	13,330 → 10,501	0.097 → 0.103

Machine load was 17–48 on 96 cores during these runs, and the throughput script's synthetic producer is CPU-bound, so ranges are wide (e.g. angle old 4.4k–9.6k). The amplitude row is the two-mode behaviour noted above: six consecutive runs per build, CPU-pinned with taskset, give old 11,474 4,357 4,322 5,321 8,409 4,408 and new 7,712 4,522 4,152 4,483 11,472 4,046 — same modes, same peak (11.5k), which mode a run lands in is not build-dependent.

Paths the review flagged, per call, 3 interleaved rounds, medians:

Path	old	new
amplitude, one 2^24-element f32 CUDA tensor	1.82 ms	0.92 ms
amplitude, one 2^10-element f64 host sample	0.47 ms	0.96 ms (0.39–2.09; a load spike hit the third round of every new row)
amplitude, one 2^16-element f64 host sample	0.15 ms	0.15 ms
angle, 64×16 f32 CUDA tensor	1.14 ms	0.28 ms
basis, 64×1 int64 CUDA tensor	2.00 ms	0.17 ms

Net: no encoding is measurably slower on the repo's standard benchmarks; the CUDA-tensor paths are several times faster than main because validation no longer round-trips inside the launch.

[ryankert01]

PRs are up, one per commit, each cumulative against main (review in order, last commit of each is the new part):

1. [QDP] Add cross-backend parity grid and benchmark baseline tool (1/5) #1478 parity grid and benchmark baseline tool
2. [QDP] Load CUDA kernels at runtime from embedded fatbins (2/5) #1479 runtime kernel loading from embedded fatbins
3. [QDP] Collapse the per-encoder matrix into one generic encode path (3/5) #1480 encoder collapse (behaviour changes listed in the PR)
4. [QDP] Remove CUDA host launchers, FFI declarations and stubs (4/5) #1481 remove host launchers, FFI declarations and stubs
5. [QDP] Document the one-encode-path architecture for contributors (5/5) #1482 contributor docs and Makefile targets

[ryankert01]

The series is now a GitHub Stack (#1488) with branches in apache/mahout; the fork-based PRs #1478–#1482 are closed. Review in order:

1. [QDP] Add cross-backend parity grid and benchmark baseline tool #1483 parity grid and benchmark baseline tool
2. [QDP] Load CUDA kernels at runtime from embedded fatbins #1484 runtime kernel loading from embedded fatbins
3. [QDP] Collapse the per-encoder matrix into one generic encode path #1485 encoder collapse (behaviour changes listed in the PR)
4. [QDP] Remove CUDA host launchers, FFI declarations and stubs #1486 remove host launchers, FFI declarations and stubs
5. [QDP] Document the one-encode-path architecture for contributors #1487 contributor docs and Makefile targets

Each PR shows only its own commit; GitHub retargets the next base as each one merges.