finishes rest of ADM stage details on local machine

Author: skyler-ruiter

CHANGELOG.md

@@ -9,13 +9,24 @@ Version numbers follow [Semantic Versioning](https://semver.org/).
 
 ## [Unreleased] — 2.0.0
 
+### Fixed
+- `adm_map_decoupled_u16`/`adm_map_thrust_u16`/`adm_map_decoupled_u32`/`adm_map_thrust_u32`: added `#ifndef NDEBUG` overflow sentinel — kernels call `atomicOr(d_overflow_flag, 1)` when a thread's `bit_offset` exceeds `kChunk × kMaxSignalBytes × 8`; host checks the flag after `cudaStreamSynchronize` and throws a `std::runtime_error` to catch inputs that violate the algorithm's bounded-diff assumption
+- `adm_map_decoupled_u16`/`adm_map_decoupled_u32`: `__shared__ excl_sum` was uninitialized for the first warp block (warp=0), causing non-deterministic writes to `d_concat_signals`; initialize to 0 at kernel entry
+- `tests/stages/test_adm.cpp` AD2 (`U32RoundTrip`): `make_u32_data` amplitude (±12000) exceeded the algorithm's per-thread `local_bits` capacity (64 bytes, supports max diff ≤ 4032); reduced amplitude to ±250
+- `tests/stages/test_adm.cpp` AD7 (`SerializeDeserialize`): used `adm_encode<uint16_t>` with `dtype=U32`, causing the U32 kernel to mis-interpret uint16_t bytes as uint32_t values with huge diffs; changed to `adm_encode<uint32_t>` with matching `make_u32_data`
+
 ### Added
 - CLI `-v`/`-vv`/`-vvv` and `--verbose[=N]` flags: route library log output (INFO/DEBUG/TRACE) to stderr via `fz::Logger::enableStderr()`
 - CLI `--profile`: now prints the full per-stage GPU timing table (`PipelinePerfResult::print()`) in compress, decompress, and benchmark modes; benchmark captures both compress and decompress stage breakdowns from the last timed run
 - CLI `--print-pipeline`: calls `pipeline->printPipeline()` after finalize to display stage topology and connections
 - CLI `--bounds-check`: enables `pipeline->enableBoundsCheck(true)` for runtime buffer overrun detection
 - CLI `--report` now includes peak device memory usage (`pipeline->getPeakMemoryUsage()`) for compress and benchmark modes
 - CLI: TOML config path now respects `--warmup`, `--profile`, `--bounds-check`, and `--print-pipeline` flags (previously only the dynamic builder path applied these settings)
+- `ADMStage`: Adaptive Data Mapping stage adapted from the MANS project (Huang et al., BSD-3-Clause); remaps `uint16_t[]`/`uint32_t[]` streams into a compact 8-bit symbol domain before entropy coding; `isGraphCompatible()=false`; 12-byte FZM header stores dtype + `num_elements`; two encode paths: decoupled look-back prefix sum for gsize ≤ 1024 and Thrust `exclusive_scan` fallback for larger arrays; 9 persistent scratch device buffers managed via `MemoryPool::allocatePersistentDevice`
+- `ADMStage` kernel files: `modules/transforms/adm/mapping_uint16.cu` and `mapping_uint32.cu` — TU-private `__global__` kernels with exported host wrappers (`compress_u16`/`decompress_u16`/`get_max_u16_payload_bytes` and u32 equivalents); all per-call `cudaMalloc`/`cudaFree` from the MANS reference replaced with `AdmScratch` pool pointers
+- `ADMStage` registered in `stage_factory.h` (`case StageType::ADM:`) and `config.cpp` (`addADMStage`/`saveADMStage`/`kStageRegistry` entry with TOML type `"ADM"` and `dtype` key)
+- `ADMStage` added to `fzgpumodules.h` public include
+- `tests/stages/test_adm.cpp`: 12 tests (AD1–AD12) covering u16/u32 round-trip, small input (< one warp block), large input (Thrust fallback path), zero input, compression ratio, header serialization, save/restore state, graph-compatibility, Pipeline integration, file round-trip, and LorenzoQuant→ADM→ANS end-to-end pipeline
 - `modules/coders/ans/dietgpu/`: vendored dietGPU rANS headers (Meta Platforms, MIT license); namespace adapted to `fz::ans`, histogram functions stripped in favor of shared `fz::module::GPU_histogram_generic`
 - `ANSStage`: full class definition in `ans_stage.h` — `ANSConfig` (12-byte FZM header with prob_bits + original_bytes_), 7 persistent scratch pointer fields, `isGraphCompatible()=false`, `estimateOutputSizes()`, `serializeHeader()`/`deserializeHeader()`, `saveState()`/`restoreState()`, `getRequiredInputAlignment()=4`, `estimateDeviceFootprintBytes()`, `estimateScratchBytes()`, `onFinalize()` pre-allocation path
 - `ANSStage::execute()` in `ans_stage.cu`: forward path (histogram → `ansCalcWeights` → `ansEncodeBatch<10,4096>` → `batchExclusivePrefixSum` → `ansEncodeCoalesceBatch<64>` → D2H header readback); inverse path (D2H header peek → `ansDecodeTable<256>` → occupancy-based `ansDecodeKernel<128,10,4096>`); `initScratch()`/`onFinalize()`/`estimateDeviceFootprintBytes()`/`estimateScratchBytes()` implementations

CMakeLists.txt

@@ -257,6 +257,9 @@ add_library(fzgmod_modules
   modules/coders/huffman/huffman_stage.cu
   modules/coders/ans/ans_stage.cu
   modules/coders/ans/dietgpu/GpuANS.cu
+  modules/transforms/adm/adm_stage.cu
+  modules/transforms/adm/mapping_uint16.cu
+  modules/transforms/adm/mapping_uint32.cu
 )
 
 target_link_libraries(fzgmod_modules

THIRD_PARTY.md

@@ -180,10 +180,18 @@ URL: https://github.com/facebookresearch/dietgpu
 **Used by:** `ADMStage` (`modules/transforms/adm/`), `MANSStage` (`modules/fused/mans/`)
 
 **Relationship:**
-- The Adaptive Data Mapping (ADM) kernels in `ADMStage` and the fused
-  ADM+ANS logic in `MANSStage` are algorithm-faithful reimplementations
-  adapted from the MANS project by Huang et al.  The GPU rANS component
-  used inside these stages is covered by the dietGPU entry above.
+- `ADMStage` (`modules/transforms/adm/mapping_uint16.cu`,
+  `mapping_uint32.cu`) — GPU kernels are a direct port of
+  `nv/adm/mapping_uint16.cu` and `nv/adm/mapping_uint32.cu` from the MANS
+  repository.  Kernel logic is unchanged.  Changes from the original: unused
+  `MansParams` parameter removed; per-call `cudaMalloc`/`cudaFree` replaced
+  by pool-allocated `AdmScratch`; `check_cuda()` replaced by `FZ_CUDA_CHECK`;
+  namespace changed from `mans::nv::adm` to `fz::adm`; kernels renamed with
+  `_u16`/`_u32` suffix to avoid TU-level naming conflicts; inline Chinese
+  comments translated to English.
+- `MANSStage` (`modules/fused/mans/`) — to be added in a future release;
+  will follow the fused ADM+rANS design from the MANS repository.  The GPU
+  rANS component is covered by the dietGPU entry above.
 
 **License:**
 

docs/stages/adm.md

@@ -0,0 +1,216 @@
+# ADMStage {#stage_adm}
+
+**Header:** `transforms/adm/adm_stage.h`  
+**Class:** `fz::ADMStage`  
+**Category:** Transform (lossless)
+
+**Common instantiation:**
+```cpp
+auto* adm = p.addStage<fz::ADMStage>();
+adm->setDtype(fz::ADMDtype::U16);  // match upstream output type
+```
+
+---
+
+## What it does
+
+Remaps a `uint16_t[]` or `uint32_t[]` integer stream into a compact 8-bit symbol
+domain, dramatically improving the compression ratio of a downstream entropy coder
+(typically `ANSStage`).
+
+- **Forward:** `uint16_t[]` or `uint32_t[]` → opaque ADM payload (`uint8_t[]`)
+- **Inverse:** ADM payload → original integer array (exact reconstruction)
+
+**Algorithm:**  ADM partitions the input into 512-element warp blocks and computes
+a per-block center (mean).  Each element is then encoded as a *(code, unary-signal)*
+pair relative to the center.  The resulting byte codes have a highly skewed,
+low-entropy distribution ideal for GPU rANS (`ANSStage`) or Huffman coding.
+
+The output is an opaque ADM payload — `getOutputDataType()` returns `UNKNOWN` to
+opt out of downstream type checking.
+
+---
+
+## Stage settings
+
+| Setting | Type | Default | Purpose |
+|---|---|---|---|
+| `setDtype(ADMDtype)` | `ADMDtype` | `U16` | Input element type: `U16` or `U32` |
+
+`ADMDtype::U16` expects `uint16_t` input; `ADMDtype::U32` expects `uint32_t` input.
+The dtype must match the upstream stage's output element type.
+
+---
+
+## Typical pipeline
+
+### Standalone (integer array input)
+
+```cpp
+Pipeline p(in_bytes, MemoryStrategy::PREALLOCATE);
+
+auto* adm = p.addStage<ADMStage>();
+adm->setDtype(ADMDtype::U16);
+p.finalize();
+
+p.compress(d_in, in_bytes, stream);
+```
+
+### cuSZ-style Lorenzo + ADM + ANS (recommended)
+
+```cpp
+Pipeline p(in_bytes, MemoryStrategy::PREALLOCATE);
+
+auto* lrz = p.addStage<LorenzoQuantStage<float, uint16_t>>();
+lrz->setErrorBound(1e-3f);
+lrz->setQuantRadius(512);
+lrz->setZigzagCodes(true);
+
+auto* adm = p.addStage<ADMStage>();
+adm->setDtype(ADMDtype::U16);
+p.connect(adm, lrz, "codes");  // connect to "codes" port
+
+auto* ans = p.addStage<ANSStage>();
+p.connect(ans, adm);
+
+p.finalize();
+```
+
+**Why ADM before ANS:**  `ANSStage` operates on 8-bit symbols (256-entry alphabet).
+Raw `uint16_t` quantization codes have a 65536-symbol alphabet, which ANS cannot
+encode directly.  ADM acts as a symbol-domain adapter — it remaps the wide integer
+stream into the 8-bit domain while preserving all information losslessly.
+
+---
+
+## TOML configuration
+
+```toml
+[[stage]]
+name   = "adm"
+type   = "ADM"
+dtype  = "uint16"   # "uint16" (default) or "uint32"
+inputs = [{from = "lrz", port = "codes"}]
+```
+
+---
+
+## Execution flow (CPU–GPU movement pattern) {#adm-execution}
+
+### Forward pass
+
+ADM uses one of two prefix-sum strategies depending on the number of warp blocks
+(`gsize = ⌈N / 512⌉`):
+
+**Decoupled look-back path** (`gsize ≤ 1024`):
+
+```
+GPU  ←input uint16_t[]                            output ADM payload uint8_t[]→
+  1. adm_map_decoupled_u16 — per-warp center + encode, decoupled prefix sum
+  2. adm_concat_u16 — pack (code, signals) into contiguous payload
+     └─ cudaMemcpyAsync D2H + cudaStreamSynchronize ◄── HOST BARRIER
+        (reads actual output_lengths to determine actual_output_size_)
+```
+
+**Thrust fallback path** (`gsize > 1024`, ~512 K+ elements for uint16_t):
+
+```
+GPU  ←input uint16_t[]                            output ADM payload uint8_t[]→
+  1. adm_map_thrust_u16 — per-warp center + encode, Thrust exclusive_scan
+  2. adm_concat_u16 — pack (code, signals) into contiguous payload
+     └─ cudaMemcpyAsync D2H + cudaStreamSynchronize ◄── HOST BARRIER
+```
+
+### Inverse pass
+
+```
+GPU  ←input ADM payload uint8_t[]            output uint16_t[]→
+  1. adm_decompress_u16 — decode (code, signals) → original values using stored centers
+```
+
+**Consequence:** one host barrier per compress call — the stage is not CUDA Graph
+compatible.  `isGraphCompatible()` returns `false`.
+
+---
+
+## Scratch buffers / device footprint
+
+`ADMStage` pre-allocates 10 persistent device buffers from the pipeline
+`MemoryPool` at `finalize()` time (PREALLOCATE mode) or on the first `execute()`
+call (MINIMAL mode).  All buffers are grow-only.
+
+Let `gsize = ⌈N / 512⌉` and `sig = kMaxSignalBytes` (2 for U16, 4 for U32):
+
+| Buffer | Size formula |
+|---|---|
+| `d_signal_length_` | `gsize × 4` B |
+| `d_output_lengths_` | `(gsize + 1) × 4` B |
+| `d_centers_` | `gsize × sizeof(T)` (2 or 4 B per element) |
+| `d_block_flags_` | `⌈gsize × 512 / 32⌉ × 4` B |
+| `d_codes_` | `N` B |
+| `d_concat_signals_` | `N × sig` B |
+| `d_bit_signals_` | `N × sig` B (Thrust fallback path) |
+| `d_loc_offset_` | `(gsize + 1) × 4` B |
+| `d_prefix_state_` | `(gsize + 1) × 4` B |
+| `d_overflow_flag_` | `4` B (debug overflow sentinel; always allocated) |
+
+For 4096 uint16_t elements: ~68 KiB total device footprint.
+For 1M uint16_t elements: ~10 MiB total device footprint.
+
+---
+
+## Serialized header
+
+The FZM stage header is 12 bytes:
+
+```
+[0]     dtype         (0 = U16, 1 = U32)
+[1..3]  reserved      (zero)
+[4..11] num_elements  (uint64_t LE; element count of the uncompressed input)
+```
+
+`num_elements` is stored so that `estimateOutputSizes()` can return the exact
+decompressed byte count for inverse-pass output buffer allocation.
+
+---
+
+## Limitations {#adm-limitations}
+
+**Not CUDA Graph compatible.** A device-to-host synchronization occurs in every
+forward call to read the actual payload size.  `isGraphCompatible()` returns `false`.
+
+**Bounded-diff constraint.** Each GPU thread encodes `kChunk = 16` elements into a
+fixed `kChunk × kMaxSignalBytes` byte buffer (`32` B for U16, `64` B for U32).  The
+number of signal bits required per element is `⌈diff / 126⌉`, where `diff` is the
+element's absolute deviation from the warp-block center.  If the total per-thread
+signal bits exceed the buffer size, the kernel writes out-of-bounds.  In debug
+builds (`NDEBUG` not defined) an `atomicOr` sentinel detects this condition and
+`compress_u16`/`compress_u32` throw `std::runtime_error` rather than silently
+corrupting data.  ADM is designed for bounded quantization codes (output of
+`LorenzoQuantStage` or `QuantizerStage`) — arbitrary integer arrays with large
+value ranges are not supported.
+
+**Opaque output.** The ADM payload is not a self-describing format — it cannot be
+decoded without the FZM header (which stores `num_elements`).  The payload must
+always be wrapped in a `Pipeline` with `writeToFile`/`decompressFromFile` or paired
+with the pipeline `decompress()` that restores the header.
+
+**Dtype must be set before `finalize()`.** `setDtype()` must be called before
+`pipeline.finalize()` so the correct element width is used for scratch sizing and
+type checking.
+
+---
+
+## Acknowledgements
+
+`ADMStage` is a direct port of the ADM encode/decode kernels from **MANS**
+(Wenjing Huang, Jinwu Yang, JingKai Huang, Haoquan Long, Dingwen Tao,
+Guangming Tan) from `nv/adm/` in the MANS repository.  Kernel logic is
+unchanged; changes from the original are documented at the top of
+`modules/transforms/adm/mapping_uint16.cu` and `mapping_uint32.cu`.
+
+> Wenjing Huang, Jinwu Yang, JingKai Huang, Haoquan Long, Dingwen Tao,
+> Guangming Tan. *MANS: Multidimensional Adaptive Numerical Compressor for
+> Scientific Data.* https://github.com/hpdps-group/MANS
+
+See `THIRD_PARTY.md` for the full BSD-3-Clause license text.

docs/stages/ans.md

@@ -176,9 +176,8 @@ one in every inverse call (to read the header before decoding).
 
 **Byte-level encoding only.** `ANSStage` operates on `uint8_t` symbols (256-entry
 alphabet).  For multi-byte integer streams (e.g., `uint16_t` quantization codes),
-pair it with `ADMStage` (coming in Phase 2) which remaps the wide symbol space into
-the 8-bit domain before ANS coding — or use `MANSStage` (Phase 3) for the fused
-path.
+pair it with `ADMStage` (\ref stage_adm) which remaps the wide symbol space into
+the 8-bit domain before ANS coding — or use `MANSStage` for the fused path.
 
 **Compression ratio depends on upstream symbol compactness.** ANS achieves its
 theoretical Shannon entropy bound only when the symbol distribution is known at

docs/stages/transforms.md

@@ -4,3 +4,4 @@
 |---|---|
 | \subpage stage_zigzag | Zigzag encode/decode (signed <-> unsigned) |
 | \subpage stage_negabinary | Negabinary encode/decode (signed <-> unsigned) |
+| \subpage stage_adm | Adaptive Data Mapping — remaps u16/u32 into compact 8-bit domain (MANS) |

include/fzgpumodules.h

@@ -31,3 +31,4 @@
 #include "shufflers/bitshuffle/bitshuffle_stage.h"
 #include "coders/huffman/huffman_stage.h"
 #include "coders/ans/ans_stage.h"
+#include "transforms/adm/adm_stage.h"

include/fzm_format.h

@@ -91,6 +91,7 @@ enum class StageType : uint16_t {
     BITSHUFFLE = 17, 
     RZE        = 18,  
     ANS        = 20,   ///< TODO: describe this stage
+    ADM        = 19,   ///< TODO: describe this stage
 };
 
 /**
@@ -313,6 +314,7 @@ inline std::string stageTypeToString(StageType type) {
         case StageType::RZE:         return "RZE";
         case StageType::LORENZO:     return "Lorenzo";
         case StageType::ANS:  return "ANS";
+        case StageType::ADM:  return "ADM";
         default:                     return "Unknown";
     }
 }

include/stage/stage_factory.h

@@ -18,6 +18,7 @@
 #include "coders/bitpack/bitpack_stage.h"
 #include "coders/huffman/huffman_stage.h"
 #include "coders/ans/ans_stage.h"
+#include "transforms/adm/adm_stage.h"
 
 #include <memory>
 #include <stdexcept>
@@ -247,6 +248,13 @@ inline Stage* createStage(StageType type, const uint8_t* config, size_t config_s
             break;
         }
 
+        case StageType::ADM: {
+            auto* s = new ADMStage();
+            s->deserializeHeader(config, config_size);
+            stage = s;
+            break;
+        }
+
         default:
             throw std::runtime_error("Unknown stage type: "
                 + std::to_string(static_cast<uint16_t>(type)));