Skip to content

Commit 4bd8cab

Browse files
committed
adds graph capture support for cuszp variants
1 parent fcf67d1 commit 4bd8cab

6 files changed

Lines changed: 283 additions & 72 deletions

File tree

‎CHANGELOG.md‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,11 @@ Version numbers follow [Semantic Versioning](https://semver.org/).
1717
- `modules/coders/huffman/phf/hf_kernels.cu` `KERNEL_CUHIP_HF_decode` — guarded the post-symbol bit fetch in `single_thread_inflate`: after emitting a symbol the decoder unconditionally fetched the next bit, which when `total_bw % 32 == 0` triggered an `input[i / 32]` read one 32-bit word past the partition's bitstream slice (and past the encoded-bitstream allocation entirely on the final partition `gid == pardeg - 1`). compute-sanitizer flagged this as `Invalid __global__ read of size 4 bytes ... 1 byte after the nearest allocation` on every Huffman round-trip test where the last partition's bit count happened to land on a 32-bit boundary. Fix: only fetch when `i < total_bw`. The outer `while (i < total_bw)` catches the terminating case on the next iteration, so the fetch is unnecessary. `test_huffman` went from 91 → 0 sanitizer errors. Upstream PHF (origin/v1.1.0_dev) has the same bug; this fix lives in our adapter-deltas comment block.
1818
- `tests/helpers/stage_harness.h` `pipeline_round_trip<T>()` was unconditionally calling `cudaFree(d_dec)` on the decompress output. When the pipeline uses `pool_managed_decomp_=true` (the default), `d_dec` is owned by `MemoryPool::graph_allocations_` and freed by `Pipeline::~Pipeline()` via `cudaFreeAsync`. The harness's extra `cudaFree` was a double-free that surfaced under compute-sanitizer as `cudaErrorInvalidValue` from `cudaFreeAsync` at teardown (one error per affected test), and the sticky error propagated to the next test's `cudaGetLastError()` checks in `CompressionDAG::addStage`. Harness now branches on `pipeline.isPoolManagedDecompOutput()` — only `cudaFree`s when the caller opted into caller-owned mode via `setPoolManagedDecompOutput(false)`. `pipeline_file_round_trip<T>()` unchanged (it uses the static `Pipeline::decompressFromFile`, which always returns a `cudaMalloc`-owned pointer). Also fixed two callers in `tests/stages/test_ginterp.cpp` (GI27, GI29) that replicated the same pattern inline. Result: `test_ginterp`, `test_lorenzo`, `test_quantizer`, `test_bitpack`, `test_ans`, `test_graph_capture` all report `ERROR SUMMARY: 0 errors` under `compute-sanitizer --tool=memcheck` (was 1+ pre-existing errors on most). `test_huffman` still has a pre-existing kernel-side issue unrelated to ownership.
1919

20+
### Changed
21+
- `AdaptiveBitpackStage<T>` forward (compress) path is now **graph-compatible** (`isGraphCompatible()` returns `!is_inverse_`, was unconditional `false`). cuSZp's data-dependent compressed-size readback (`cmpSize` D2H + the `cudaStreamSynchronize` it required) was moved out of `execute()` into a new `postStreamSync()` override, which the pipeline calls after the launch and a full stream sync — outside any CUDA-graph capture window. To support reading the scanned payload length post-sync (and to avoid any allocation inside a captured graph replay), the per-block cost/offset arrays are now **persistent** pool scratch (`allocate(..., persistent=true)`, grown lazily when a larger input is seen, freed in the destructor) instead of per-execute alloc/free; the CUB scan temp stays transient. `execute()` sets a worst-case provisional `actual_output_size_` (graph *recording* does not invoke `postStreamSync`), which `postStreamSync()` then refines to `meta_region + scanned_payload`. Mirrors the `RZEStage` forward precedent. The inverse path is unchanged and remains out of graph capture. This makes the full modular cuSZp/cuSZp2/cuSZp3 compress pipelines (`Quantizer(linear) → Lorenzo(block)/TiledLorenzo → AdaptiveBitpack`) capturable and replayable as a CUDA graph. `tests/stages/test_adaptive_bitpack.cpp`: AB10 retargeted (`GraphCompatForwardOnly`) and new AB20 `GraphCaptureRoundTrip` (capture + 3× replay of the cuSZp compress graph, size-stability + within-bound round-trip); suite 20/20, compute-sanitizer clean. Docs: `docs/stages/adaptive_bitpack.md` graph-compatibility section.
22+
2023
### Added
21-
- `examples/cuszp_variants.cpp` — benchmarks the three modular cuSZp variants (cuSZp plain, cuSZp2 outlier, cuSZp3 2-D tiled) under PREALLOCATE, reporting compression ratio, round-trip max error, and peak compress/decompress throughput, then attempts CUDA graph capture per variant and reports that graph mode is unavailable because the terminal `AdaptiveBitpackStage` performs a forward D2H size readback (the same per-compress readback real cuSZp does) — so PREALLOCATE is the fastest capturable path. Registered as the `cuszp_variants` example target.
24+
- `examples/cuszp_variants.cpp` — benchmarks the three modular cuSZp variants (cuSZp plain, cuSZp2 outlier, cuSZp3 2-D tiled), reporting compression ratio, round-trip max error, and peak compress/decompress throughput under PREALLOCATE, then **captures and replays each variant's compress pipeline as a CUDA graph** and validates the round-trip through the graph-produced buffer, printing a PREALLOCATE-vs-GRAPH comparison. Registered as the `cuszp_variants` example target.
2225
- `tests/stages/test_cuszp_block_sizes.cpp` — verifies the block/tile size is a freely tunable knob for the modular cuSZp-family pipelines (not locked to cuSZp's default 32 / 8×8). 6 tests: cuSZp plain block sweep ({8,16,32,64,128,256,1024}, each compresses + round-trips within eb), cuSZp2 outlier sweep, mismatched predictor-vs-coder block sizes (independent for correctness), cuSZp3 2-D tile sweep (TiledLorenzo tile + AdaptiveBitpack `block_size = tile_elems`), cuSZp3 coder-block ≠ tile_elems (still lossless within eb), and a 3-D tile sweep. All pass.
2326
- `tests/stages/test_large_data.cpp` — **optional** large-data stress tests (env-gated: each test SKIPs unless `FZ_LARGE_DATA_TEST=1`, and self-SKIPs if the GPU can't hold the working set). Motivated by a published comparison reporting cuSZ-Hi failing to produce output on large tensors (a suspected 32-bit index/offset overflow). Each test pushes one stage/pipeline to a default 2 GiB input — crossing the 2^31-byte offset boundary — (overridable via `FZ_LARGE_DATA_BYTES`) with device-side tiled fill + 64 MiB-chunk verification (host RAM stays small); lossless stages checked for an exact byte round-trip, lossy within the error bound; OOM is caught and downgraded to SKIP. Covers Lorenzo, TiledLorenzo, AdaptiveBitpack, BitplaneRLE, RZE, Bitshuffle, Difference (lossless), Quantizer(linear), LorenzoQuant, GInterp (lossy), and the cuSZp/cuSZp3 pipelines. **All 12 pass at 2 GiB on a 12 GB RTX 3080 Ti** — including GInterp on an ~808³ float volume (the cuSZ-Hi-family stage the comparison flagged), confirming our stages handle the large-tensor size that reportedly broke cuSZ-Hi. Registered with a `large` CTest label. Entropy coders (Huffman/ANS/ADM) are intentionally excluded (raw stress patterns don't meet their symbol/codebook assumptions).
2427
- `TiledLorenzoStage<T>` (`modules/predictors/tiled_lorenzo/`, `StageType::TILED_LORENZO = 25`) — the cuSZp3 (SC'25) **dimension-aware (tiled separable) delta** predictor as a modular stage, completing the modular cuSZp3 pipeline. Within each fixed tile (default 8×8 in 2-D, 4×4×4 in 3-D; `setTileShape(tx,ty,tz)`, product ≤ 1024) it applies a separable integer Lorenzo — the leading x-column predicts down y (`prevQuant_y`), the leading y/z edges predict down z (`prevQuant_z`), and every other element predicts from its x neighbour — exactly mirroring the cuSZp3 plain/outlier kernels (verified against `compressors/cuSZp3/src/cuSZp_kernels_{2D,3D}_f32.cu`). The forward pass emits **tile-major** output (each tile's `tile_elems` elements contiguous, edge tiles zero-padded), so a downstream `AdaptiveBitpackStage(block_size = tile_elems)` packs each tile as exactly one block — reproducing cuSZp3's per-tile fixed-rate layout with the coder unchanged. The inverse un-tiles back to natural row-major order via a storage-free separable prefix-sum (3 running scalars per tile). Signed `int16_t`/`int32_t`; forward is parallel one-thread-per-output-element, inverse one-thread-per-tile (correctness-first; fusion is a downstream-compiler concern). Made a **separate stage** rather than a `LorenzoStage` mode because the tile-major reshape + padding break LorenzoStage's size/layout-preserving contract and the separable formula differs from its N-D inclusion-exclusion delta (same reasoning as `AdaptiveBitpackStage` vs `BitpackStage`). **Graph-compatible** (pure kernels, deterministic sizes; no D2H). **Independent reimplementation — no cuSZp source copied.** New 20-byte `TiledLorenzoConfig` (dtype, ndim, tile_x/y/z, dims) in the FZM header. Registered in `stage_factory.h`, `fzgpumodules.h`, and TOML (`type = "TiledLorenzo"`, `data_type`, `tile_x`/`tile_y`/`tile_z`). The three cuSZp3 modes map to: fixed = `Quantizer(linear) → AdaptiveBitpack`; plain = `… → TiledLorenzo → AdaptiveBitpack`; outlier = plain + `AdaptiveBitpack.setOutlierSelection(true)`. `tests/stages/test_tiled_lorenzo.cpp` — 14 tests (2-D/3-D int16/int32 round-trips, partial edge tiles in 2-D and 3-D, custom tile shape, constant field, `setTileShape` bounds, header round-trip, port/type/graph contract, stage-type id, the full `Quantizer(linear)→TiledLorenzo→AdaptiveBitpack(64)` plain and outlier float pipelines, and an FZM file round-trip); all pass, compute-sanitizer-clean (0 errors). Docs: `docs/stages/tiled_lorenzo.md`; design notes in `memory/cuszp_stages.md` Part 8.

‎docs/stages/adaptive_bitpack.md‎

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -100,8 +100,19 @@ Single input → single output.
100100
101101
## Graph compatibility
102102
103-
`isGraphCompatible()` is **false** — the forward path does a host-blocking D2H to
104-
read the scanned total payload length (same pattern as `BitplaneRLEStage`).
103+
`isGraphCompatible()` is **true for the forward (compress) path, false for the
104+
inverse**. The archive length is still data-dependent (cuSZp's `cmpSize`), but
105+
the host readback of the scanned total payload is deferred to `postStreamSync()`
106+
— which the pipeline calls after the launch and a full stream sync, outside any
107+
capture window — so `execute()` enqueues only stream-ordered device work. The
108+
per-block cost/offset scratch is kept persistent (grown lazily, freed in the
109+
destructor) so the readback can happen post-sync and no allocation occurs inside
110+
a captured graph replay. This mirrors `RZEStage`'s forward path. The inverse
111+
keeps a per-execute layout and is left out of graph capture.
112+
113+
This means the full cuSZp pipelines (`Quantizer(linear) → Lorenzo(block)/TiledLorenzo
114+
→ AdaptiveBitpack`) can be captured and replayed as a CUDA graph on the compress
115+
side. See `examples/cuszp_variants.cpp` for a PREALLOCATE-vs-GRAPH benchmark.
105116
106117
---
107118

0 commit comments

Comments
 (0)