#include <adaptive_lorenzo_stage.h>

Inheritance diagram for fz::AdaptiveLorenzoStage< T >:

Public Member Functions
void	setInverse (bool inv) override

void	execute (fz::stream_t stream, MemoryPool pool, const std::vector< void > &inputs, const std::vector< void * > &outputs, const std::vector< size_t > &sizes) override

std::string	getName () const override

std::vector< std::string >	getOutputNames () const override

void	postStreamSync (fz::stream_t stream) override

bool	isGraphCompatible () const override

size_t	estimateScratchBytes (const std::vector< size_t > &input_sizes) const override
	Defined in the .cu — sizing the CUB scan temp needs a CUDA translation unit.

std::vector< size_t >	estimateOutputSizes (const std::vector< size_t > &input_sizes) const override

void	saveState () override

std::unordered_map< std::string, size_t >	getActualOutputSizesByName () const override

size_t	getActualOutputSize (int index) const override

uint16_t	getStageTypeId () const override

uint8_t	getOutputDataType (size_t output_index) const override

uint8_t	getInputDataType (size_t input_index) const override

size_t	serializeHeader (size_t, uint8_t *buf, size_t max_size) const override

void	deserializeHeader (const uint8_t *buf, size_t size) override

size_t	getMaxHeaderSize (size_t) const override

Public Member Functions inherited from fz::Stage
virtual size_t	getRequiredInputAlignment () const

int	getOutputIndex (const std::string &name) const

virtual std::vector< std::string >	getRunNotes () const

virtual void	setDims (const std::array< size_t, 3 > &dims)

virtual void	onFinalize (size_t, MemoryPool *)

virtual size_t	estimateDeviceFootprintBytes (size_t) const

virtual size_t	estimatePinnedFootprintBytes (size_t) const

Additional Inherited Members
Static Public Member Functions inherited from fz::Stage
static constexpr bool	isSupportedOnBackend ()

Detailed Description

template<typename T>
class fz::AdaptiveLorenzoStage< T >

Per-tile adaptive multi-order Lorenzo predictor (FSZ prediction stage). Lossless.

Splits the flattened input into tiles of coder_block_size * blocks_per_tile elements and, for each tile independently, picks whichever of four prediction variants encodes smallest:

variant	residual
LZ1	`q_i - q_{i-1}`
LZ2	`q_i - 2q_{i-1} + q_{i-2}`
LZ1 + centering	LZ1 on `q_i - mu`
LZ2 + centering	LZ2 on `q_i - mu`

The prediction chain runs the length of the whole tile, not the coder block — that is the cross-block prediction state of FSZ, obtained here by making the tile a multiple of the coder's block. Only the tile's first one or two elements lack predecessors instead of one per coder block.

Selection is by exact encoded size, not entropy. For a coder block with maximum residual magnitude of r bits, AdaptiveBitpackStage emits 0 bytes when r == 0 and word_bytes * (r + 1) otherwise (a sign bitmap plus r bit-planes). The stage sums that over the tile's blocks for each variant and takes the minimum, adding sizeof(T) for the mean when a centered variant is in the running. This makes the selection exact only for a downstream AdaptiveBitpackStage with a matching block_size — routing the residuals into an entropy coder instead leaves the choice merely reasonable, not optimal.

Single-pass evaluation. All four variants are derived from one read of the tile, because a constant offset cancels out of a k-th order difference for every element with k predecessors (delta^k(q - mu) == delta^k(q)). The centered variants therefore differ from the uncentered ones only in the tile's first one (LZ1) or two (LZ2) residuals, so their costs come from adjusting the first coder block's rate rather than recomputing anything.

Note: Prior work: the cross-block prediction state, the four-variant adaptive selection, and the finite-difference cancellation that makes a single-pass evaluation possible are the design of FSZ (Jiajun Huang, "FSZ: Breaking the Prediction-Throughput Trade-off in GPU Lossy Compression", SC'26, arXiv:2607.15413). This is an independent reimplementation as a modular DAG stage — FSZ fuses prediction, quantization and encoding into one kernel, whereas this stage is the prediction step alone and pairs with QuantizerStage upstream and AdaptiveBitpackStage downstream.

Forward outputs:

[0] output — residuals for the selected variant (T, one per element)
[1] modes — one byte per tile: bit 0 = order 2, bit 1 = centering
[2] means — one T per tile (meaningful only where bit 1 is set)

Template Parameters

T	Signed integer element type: int8_t, int16_t, int32_t, int64_t.

Member Function Documentation

◆ setInverse()

template<typename T >

void fz::AdaptiveLorenzoStage< T >::setInverse ( bool inverse )

inlineoverridevirtual

Switch between forward (compression) and inverse (decompression) mode. Affects getNumInputs()/getNumOutputs() for stages with asymmetric port counts.

Reimplemented from fz::Stage.

◆ execute()

template<typename T >

void fz::AdaptiveLorenzoStage< T >::execute	(	fz::stream_t	stream,
		MemoryPool *	pool,
		const std::vector< void * > &	inputs,
		const std::vector< void * > &	outputs,
		const std::vector< size_t > &	sizes
	)

overridevirtual

Execute the stage. Inputs, outputs, and sizes are device pointers/bytes.

Stages may call cudaStreamSynchronize(stream) or issue blocking D2H copies when the algorithm requires it (e.g. Huffman histogram readback for codebook construction, ANS renormalization tables). Such stages must return false from isGraphCompatible() and must document the sync points.

Note: the DAG dispatches sibling nodes (same topological level) via a sequential CPU loop, each enqueuing to its own stream. A sync inside execute() blocks the CPU from dispatching subsequent siblings until the synced stream is idle — this delays parallel branches in wide DAGs. In a linear pipeline there are no siblings and no extra cost.

Implements fz::Stage.

◆ getName()

template<typename T >

std::string fz::AdaptiveLorenzoStage< T >::getName ( ) const

inlineoverridevirtual

Human-readable name used in error messages and debug output.

Implements fz::Stage.

◆ getOutputNames()

template<typename T >

std::vector< std::string > fz::AdaptiveLorenzoStage< T >::getOutputNames ( ) const

inlineoverridevirtual

Output port names in order. Default: single port named "output". Multi-output stages (e.g. Lorenzo: "codes", "outliers") override this.

Reimplemented from fz::Stage.

◆ postStreamSync()

template<typename T >

void fz::AdaptiveLorenzoStage< T >::postStreamSync ( fz::stream_t stream )

overridevirtual

Reads back the number of tiles that actually chose centering and trims the means port to it. Called once the stream is idle, so the D2H never stalls the pipeline mid-flight.

Reimplemented from fz::Stage.

◆ isGraphCompatible()

template<typename T >

bool fz::AdaptiveLorenzoStage< T >::isGraphCompatible ( ) const

inlineoverridevirtual

The means length is data-dependent and resolved by a D2H in postStreamSync(), so the forward pass cannot be captured in a graph.

Reimplemented from fz::Stage.

◆ estimateOutputSizes()

template<typename T >

std::vector< size_t > fz::AdaptiveLorenzoStage< T >::estimateOutputSizes ( const std::vector< size_t > & input_sizes ) const

inlineoverridevirtual

Estimate output buffer sizes given input sizes. Used for buffer allocation planning in PREALLOCATE mode — must be a safe upper bound; under-estimation causes buffer overruns.

Implements fz::Stage.

◆ saveState()

template<typename T >

void fz::AdaptiveLorenzoStage< T >::saveState ( )

inlineoverridevirtual

Save/restore config state around a decompression pass. deserializeHeader() overwrites the stage's forward-pass config; saveState() is called before and restoreState() after so the stage returns to its original configuration.

Reimplemented from fz::Stage.

◆ getActualOutputSizesByName()

template<typename T >

std::unordered_map< std::string, size_t > fz::AdaptiveLorenzoStage< T >::getActualOutputSizesByName ( ) const

inlineoverridevirtual

Actual output sizes after execute(), keyed by output port name.

Implements fz::Stage.

◆ getActualOutputSize()

template<typename T >

size_t fz::AdaptiveLorenzoStage< T >::getActualOutputSize ( int index ) const

inlineoverridevirtual

Actual size of a single output by index after execute(). Avoids constructing the map for the common single-output case. Default delegates to getActualOutputSizesByName(); override to return directly from an internal field.

Reimplemented from fz::Stage.

◆ getStageTypeId()

template<typename T >

uint16_t fz::AdaptiveLorenzoStage< T >::getStageTypeId ( ) const

inlineoverridevirtual

Stage type identifier written into the FZM file header.

Implements fz::Stage.

◆ getOutputDataType()

template<typename T >

uint8_t fz::AdaptiveLorenzoStage< T >::getOutputDataType ( size_t output_index ) const

inlineoverridevirtual

DataType enum of the given output port.

Implements fz::Stage.

◆ getInputDataType()

template<typename T >

uint8_t fz::AdaptiveLorenzoStage< T >::getInputDataType ( size_t ) const

inlineoverridevirtual

Expected DataType of the given input port.

Used by Pipeline::finalize() to detect type mismatches between connected stages before any execution. Return DataType::UNKNOWN to opt out of checking — byte-transparent stages (Bitshuffle, RZE, RRE) and mock stages must return UNKNOWN; finalize() skips any connection where either side is UNKNOWN.

Reimplemented from fz::Stage.

◆ serializeHeader()

template<typename T >

size_t fz::AdaptiveLorenzoStage< T >::serializeHeader	(	size_t	output_index,
		uint8_t *	header_buffer,
		size_t	max_size
	)		const

inlineoverridevirtual

Serialize stage config into header_buffer (max 128 bytes) for the FZM file. Return the number of bytes written, or 0 if the stage has no config.

Reimplemented from fz::Stage.

◆ deserializeHeader()

template<typename T >

void fz::AdaptiveLorenzoStage< T >::deserializeHeader	(	const uint8_t *	header_buffer,
		size_t	size
	)

inlineoverridevirtual

Restore stage config from header_buffer during decompression.

Reimplemented from fz::Stage.

◆ getMaxHeaderSize()

template<typename T >

size_t fz::AdaptiveLorenzoStage< T >::getMaxHeaderSize ( size_t output_index ) const

inlineoverridevirtual

Maximum bytes this stage writes into its per-output FZM header slot.

Reimplemented from fz::Stage.

Public Member Functions

Additional Inherited Members

Detailed Description

Member Function Documentation

◆ setInverse()

◆ execute()

◆ getName()

◆ getOutputNames()

◆ postStreamSync()

◆ isGraphCompatible()

◆ estimateOutputSizes()

◆ saveState()

◆ getActualOutputSizesByName()

◆ getActualOutputSize()

◆ getStageTypeId()

◆ getOutputDataType()

◆ getInputDataType()

◆ serializeHeader()

◆ deserializeHeader()

◆ getMaxHeaderSize()