|
FZGPUModules 2.0
GPU-accelerated modular compression pipelines
|
Specialization is declaration-driven. A stage declares its fused identity; the planner walks the DAG matching those declarations by role; the runner builds the fused kernel from the declared op names. A correctly declared stage joins the fast path the same way an existing one does.
Every specializable stage spans two layers that must agree exactly:
Stage virtuals in the stage's .h) — what device op this stage maps to, its role, geometry, and runtime parameter bytes..cuh) — the how: the actual register/shared-memory kernel code. A policy may provide both forward and inverse methods, but the stage declares inverse eligibility separately and the inverse harness must support that role and chain shape.The host packs a parameter blob whose layout the device op reinterpret_casts, so the POD Params struct is shared verbatim between host and device (see modules/fused/fused_block/warp_op_params.h, modules/fused/chunk_fusion/chunk_op_params.h).
include/stage/fusion.h defines the vocabulary.
FusionAccess — a stage's dependency and codec role, which constrains how it can be composed:
| Role | Meaning | Example |
|---|---|---|
Elementwise | each output depends only on its corresponding input | linear Quantizer |
RegionLocal | values may depend within a fixed logical region, with no cross-region dependency | 1-D Lorenzo (region reset) |
SegmentCodec | encodes or decodes one fixed-size logical segment with a data-dependent encoded length | AdaptiveBitpack |
TileSelector | chooses one representation over a tile containing N codec segments | FSZ selector |
Unfusable (default) | opaque / global dependency — a fusion barrier | Huffman (global codebook) |
FusionStrategy — which execution model the op belongs to. A fused group is composed of ops that all share one strategy:
| Strategy | Execution model |
|---|---|
WarpRegister | one warp owns a ≤ 32·kMaxWarpElemsPerLane-element block, intermediates in registers + shuffles, no barriers (cuSZp / SZp) |
ChunkCooperative | one CTA owns a compatible size byte-chunk, intermediates in shared memory, __syncthreads between ops (LC / PFPL) |
The rest of this guide works the warp-register path end to end. The chunk-cooperative path uses the same declaration surface with a different harness. The current warp harness is generic over declared predictors, transforms, and coders within a float32 or float64 input / int32 code representation; those data types are checked as part of strategy matching. Float64 uses the warp-cooperative path and retains double arithmetic through quantization and reconstruction. The thread-independent path remains float32-only. The opt-in linear_high_precision quantization policy remains staged for both input types.
The registry also admits the two-stage LinearQuant → segment-coder chain and its inverse. LinearQuant1DPredictor is an internal identity policy: it emits quantized codes in natural order and applies no inverse prediction. The graph stays predictor-free; coder blocks supply the warp geometry (32/64/128 elements). Both f32 and f64 validate finite signed-int32 coordinates before compression.
Conventions:
include_header. For float64 warp input, the current generator selects op_name + "F64" for the predictor. A custom predictor used with float64 must provide that alias and a fromParams overload taking the separate double reciprocal. The shipped Lorenzo and tiled predictor policies provide both precision variants. The inverse harness deduces precision from the reconstruction step and output pointer, preserving the existing custom float predictor interface.inv2eb slot convention. Every warp predictor's Params begins with float inv2eb at offset 0. The predictor stage cannot know the error bound (the quantizer owns it), so it packs 0 there; the runner overwrites those 4 bytes from the Elementwise head's getFusedForwardQuantStep() contract after priming. The quantizer is absorbed into the predictor (it quantizes inline in delta()), which is why the Elementwise quant stage declares op "LinearQuant" with empty params. Float64 retains the same geometry POD layout but receives its double reciprocal as a separate kernel argument; it never reads the float reciprocal slot. Its quantized coordinates are checked for finite values and signed-int32 range before encoding, matching the staged quantizer's rejection of unrepresentable bins.block_size / 32) is the harness's compile-time template arg; n_ab is the padded block-covering element count (0 = 1-D, no padding).The coder (SegmentCodec, the group tail) declares similarly with op_name naming its coder policy ("PlainRateCoder" / "AdaptiveBitpackCoder").
The mirror surface, gated on isInverse(). This is what makes a stage fuse on decompress, not just compress. It reuses the same FusionSpec / FusedOpDecl structs and the same policy type names (the policy carries the inverse methods).
Two generic scalar hooks let the inverse runner pull what it needs without dynamic_casting to a concrete stage type — so a new coder/quant works too:
Keep forward and inverse eligibility in lockstep (reuse the same block-size gating), so a pipeline that fuses on compress also fuses on decompress. Guard anything the inverse harness can't undo — e.g. Lorenzo excludes centeringActive() because the undelta policy is plain first-difference.
A fused runner replaces the group's per-stage execute() calls with one kernel, so any state a stage would normally compute in execute() — and that its own inverse later reads — must be established explicitly. These Stage hooks (all default no-ops) exist for that:
| Hook | Who overrides it | Why |
|---|---|---|
primeFusedForwardState(ctx) | Quantizer | Run the value-range scan / bound resolution the fused kernel needs (and the inverse reads back). Called once per group member before codegen. |
getFusedForwardQuantStep() | linear quantizer | Provide the resolved quantization step to the generic warp runner without a concrete-stage cast. |
setFusedArchiveResult(archive, orig) | variable-length coders | Report archive + original sizes so the coder's inverse can size its output (else it falls back to the compressed size and overruns). |
setFusedInverseResult(bytes) | inverse tail (quant) | Publish the reconstructed byte count for output-size refinement. |
setFusedSideOutput(port, bytes) | outlier-producing quant | Report bytes written to an escaping side port (e.g. an outlier list) so serializeHeader matches the fused result. |
Without primeFusedForwardState, a fused quantizer's inverse read the default bound (1e-4) instead of the resolved one — producing a byte-identical archive that reconstructed 10× too small. If your stage computes anything in execute() that its inverse depends on, it must be primed. See docs/codebase_notes.md CN-FUSE-DRIVER / CN-CHUNK-WIRE.
Allocation is handled for you. When the group is installed, the DAG automatically (a) skips the device allocation for any intermediate buffer produced and consumed only inside the group — the fused kernel keeps it in registers/shared memory — and (b) collapses the group to one liveness point so PREALLOCATE buffer coloring can alias the rest of the DAG around it. Your stage does nothing for this; just make sure getFusedAuxOutputs() / the side-output hooks correctly identify any port that escapes the group (an outlier list, a means stream) so it stays materialized.
An escaping auxiliary port does not have to be a pipeline leaf. It may feed ordinary staged consumers outside the specialized region. The planner follows the group's linear port-0 main path while allowing declared auxiliary ports to cross its boundary; the executor materializes and sizes those buffers before their staged consumers run. On the inverse DAG, the corresponding staged branch runs first and its result enters the specialized inverse through side_inputs. Only declared auxiliary ports receive this treatment: an undeclared branch, or a fan-out of the main representation, remains a specialization barrier because the runner has no contract for reproducing that boundary value.
modules/fused/fused_block/warp_fusion.cuh. Each policy is a small POD with static device methods. The harness bodies (fused_rate_body / fused_pack_body / fused_unpack_body) call these; they are the only thing that changes per op.
Predictor (produces per-lane int codes, forward and inverse):
Coder (the SegmentCodec sink; forward cost/pack + inverse decode):
Transform (optional register→register map between predictor and coder): apply<EPL>() forward, invert<EPL>() inverse (applied in reverse order on decode). Current limitation: the inverse harness does not yet compose interior transforms, so a chain with an interior transform (e.g. Zigzag) fuses on compress but stays staged on decompress until invert() + applyInverseTransforms land.
The forward and inverse methods of one policy must be exact inverses that also match the staged stage's kernels bit-for-bit.
warp_fusion.cuh: add MyPredictor with fromParams, delta, and undelta. Add its Params POD to warp_op_params.h (leading float inv2eb).getFusionSpec → RegionLocal; getFusedOp → {WarpRegister, "MyPredictor", elems_per_lane, params}.getInverseFusionSpec → RegionLocal; getInverseFusedOp → {WarpRegister, "MyPredictor", elems_per_lane}.matchesWarpRegister / runWarpRegister and their inverse counterparts are role-based over the declarations — they build the WarpFusionSpec from your op_name, so a new predictor needs no registry, planner, matcher, or runner edit. The NVRTC codegen instantiates the harness with your policy name.Adding a new coder is the same, at the SegmentCodec role, with decode alongside cost/pack. Adding a new quantizer dequant is currently linear-only (the harness hard-codes code · 2·eb); a non-linear dequant would extend the harness with a dequant policy, the same way a predictor is added.
src/pipeline/fusion_registry.cpp — role-based, mirrored for forward and inverse:
matchesWarpRegister): every stage declares a WarpRegister op; front is Elementwise, one RegionLocal predictor, interior Elementwise/RegionLocal transforms, and back is SegmentCodec. No concrete types are named.runWarpRegister): primes each stage, then builds WarpFusionSpec{predictor = the RegionLocal op's name, coder = the SegmentCodec op's name, transforms = interior op names, elems_per_lane}, obtains the quantization step and reports coder state through generic Stage hooks, patches inv2eb, and calls launchNvrtcWarpFused. It does not downcast the forward stages.The inverse pair (matchesWarpRegisterInverse / runWarpRegisterInverse) is the same, over getInverseFusionSpec/getInverseFusedOp, with roles reversed (SegmentCodec coder → RegionLocal predictor → Elementwise quant) and the two scalar hooks for element count and dequant step.
The planner (planFusionGroups) only enumerates maximal fusable chains from getFusionSpec() roles — it never needs to know your op exists.
A specialization is only correct if it is indistinguishable from staged:
SpecializationPolicy::Off and ::Auto; the archives must be memcmp-equal. (For nondeterministic side outputs like an atomic-appended outlier list, compare reconstruction instead.)tests/pipeline/test_fusion_planner.cpp (see Warp1DGeneralEplFusesMatchesStaged for the pattern).Same declaration surface (getFusionSpec/getFusedOp with FusionStrategy::ChunkCooperative), different harness (modules/fused/chunk_fusion/). Device ops are templated on chunk size (Op<ChunkBytes>, one of {4096, 8192, 16384} — see chunk_geometry.h's Geom<Bytes>) with shared-memory ping-pong; the POD params live in chunk_op_params.h. The harness is NVRTC-composed from an op list, with the chunk size baked into the generated template args (ChunkFusionSpec::chunk_bytes), so a novel Elementwise → Transform* → SegmentCodec chunk chain fuses with zero new glue — proven for Quant → Diff → Bitshuffle → {RZE, RRE, RARE, RAZE} at all three sizes. All participating ops in a group must declare the same block_size (chunk size) or the planner rejects the group. The NVRTC surface needs the op headers to compile without the CUDA runtime; see CN-NVRTC-FUSE for the stub requirements — note this also means any host-only helper in a header pulled into the NVRTC translation unit must be #ifndef __CUDACC_RTC__-guarded, since NVRTC rejects an unannotated function merely being present, not just called (see chunk_geometry.h's isSupportedChunkBytes). Decompress generalization for this path beyond PFPL/RZE is future work.
The forward chunk path holds, beyond the caller's input, only the worst-case archive buffer, any split-outlier side buffers, and one look-back word per chunk:
FusedImpl::accepts_unpadded_input is set for chunk-coop. When every source consumer is such a group, finalize() skips the pad buffer and compress() passes the caller's buffer with a logical (chunk-aligned) length plus FusedRunContext::input_valid_bytes. The elementwise load ops read 0.0f past the valid length (loadOrZero), reproducing the staged zero padding bit for bit, and the NOA prime scan is clamped to the valid length.chunk_fused_direct_body encodes a chunk in shared memory, publishes its size, resolves its payload offset by decoupled look-back (warp 0 checks 32 predecessors at a time), and writes the size-table entry and payload straight into the archive with a funnel-shift word store (after PFPL's s2g). The archive is byte-identical to the scratch path, which remains available as FZ_CHUNK_SCRATCH=1 (full input-size scratch, then CUB scan + chunk_pack). One CTA per chunk with cid = blockIdx.x, relying on in-order block dispatch as CUB's look-back does.Measured on an H100 (PFPL preset, -b --runs 10, median compress device time, direct vs scratch): CESM-2D/CLDHGH (26 MB) 10% faster, a 0.5 MB synthetic input 15% faster, HACC/vx (1.12 GB) 2% faster, NYX/temperature (537 MB, very compressible) 9% slower. On NYX each chunk encodes quickly, so a block's wait on its predecessors' sizes is not hidden; the scratch path never waits. Rejected variants on the same inputs: a single-thread serial look-back (NYX +12% vs scratch); an atomic chunk ticket instead of blockIdx.x (NYX +5% vs the blockIdx.x version); a persistent grid with an atomic chunk counter, native PFPL's scheme (NYX +16%, HACC +3% vs scratch); __nanosleep backoff in the spin (no change). Both kernels use 32 registers and ~36.9 KB static shared memory, so occupancy is not the difference.
Peak live device memory of a PFPL Auto compress, measured with cuda_mem_probe against d511ebc: HACC/vx 4721 -> 2671 MB, NYX 2255 -> 1192 MB, CESM-2D/CLDHGH 108.9 -> 59.4 MB (input, archive, outliers, ~1/CR concat copy).