Add exact encoded-size oracle for plain/outlier AdaptiveBitpack

Plain/outlier AdaptiveBitpack stages declare their local coder policy;
finalize() binds compatible semantics to directly connected
AdaptiveLorenzo stages independently of fusion policy. The staged
selector now evaluates shared exact size quotes from full/rest/element-0
block statistics instead of a private plain-byte formula.

Author: skyler-ruiter

modules/coders/adaptive_bitpack/adaptive_bitpack_oracle.cuh

@@ -0,0 +1,80 @@
+#pragma once
+
+/**
+ * @file adaptive_bitpack_oracle.cuh
+ * @brief Exact, device-inline encoding quotes shared by staged and fused selectors.
+ */
+
+#include <cstdint>
+
+namespace fz::adaptive_bitpack_oracle {
+
+struct PlainFixedRateDecision {
+    uint8_t rate = 0;
+};
+
+struct PlainFixedRateQuote {
+    uint32_t payload_bytes = 0;
+    PlainFixedRateDecision decision{};
+};
+
+struct AdaptiveFixedRateDecision {
+    uint8_t rate = 0;
+    uint8_t selector = 0; ///< bit0=outlier; bits1-2=outlier bytes minus one.
+};
+
+struct AdaptiveFixedRateQuote {
+    uint32_t payload_bytes = 0;
+    AdaptiveFixedRateDecision decision{};
+};
+
+/// Exact payload emitted by plain AdaptiveBitpack for one coder unit: an all-zero
+/// unit emits nothing; otherwise it emits one sign word and `rate` plane words.
+__device__ __forceinline__ PlainFixedRateQuote quotePlainFixedRate(
+    uint32_t max_magnitude, uint32_t word_bytes)
+{
+    const uint8_t rate = max_magnitude
+        ? static_cast<uint8_t>(32 - __clz(max_magnitude))
+        : uint8_t{0};
+    PlainFixedRateQuote q;
+    q.decision.rate = rate;
+    q.payload_bytes = rate > 0
+        ? word_bytes * (static_cast<uint32_t>(rate) + 1u)
+        : 0u;
+    return q;
+}
+
+/// Exact plain-versus-element0-outlier choice used by AdaptiveBitpack's outlier
+/// mode. Ties select plain, matching the staged and warp-register encoders.
+__device__ __forceinline__ AdaptiveFixedRateQuote quoteAdaptiveFixedRate(
+    uint32_t max_magnitude, uint32_t max_rest_magnitude,
+    uint32_t first_magnitude, uint32_t word_bytes)
+{
+    const PlainFixedRateQuote plain = quotePlainFixedRate(max_magnitude, word_bytes);
+    const uint8_t rest_rate = max_rest_magnitude
+        ? static_cast<uint8_t>(32 - __clz(max_rest_magnitude))
+        : uint8_t{0};
+    const uint32_t outlier_bytes = first_magnitude
+        ? static_cast<uint32_t>((32 - __clz(first_magnitude) + 7) / 8)
+        : 0u;
+    const uint32_t outlier_payload = outlier_bytes +
+        (rest_rate > 0
+            ? word_bytes * (static_cast<uint32_t>(rest_rate) + 1u)
+            : word_bytes);
+
+    AdaptiveFixedRateQuote q;
+    if (plain.payload_bytes <= outlier_payload) {
+        q.payload_bytes = plain.payload_bytes;
+        q.decision.rate = plain.decision.rate;
+        q.decision.selector = 0;
+    } else {
+        q.payload_bytes = outlier_payload;
+        q.decision.rate = rest_rate;
+        // outlier_payload can beat plain only when outlier_bytes is nonzero.
+        q.decision.selector = static_cast<uint8_t>(
+            1u | ((outlier_bytes - 1u) << 1u));
+    }
+    return q;
+}
+
+} // namespace fz::adaptive_bitpack_oracle

modules/coders/adaptive_bitpack/adaptive_bitpack_stage.h

@@ -123,6 +123,25 @@ class AdaptiveBitpackStage : public Stage {
         return d;
     }
 
+    /// Exact local encoded-size policy for an immediately upstream adaptive
+    /// selector. Unlike getFusedOp(), this semantic declaration is available in
+    /// plain mode: AdaptiveLorenzo needs the staged coder's exact byte formula
+    /// even when no fused implementation has been selected.
+    EncodingOracleDecl getEncodingOracle() const override {
+        if (is_inverse_ || block_size_ == 0) return {};
+        EncodingOracleDecl d;
+        d.kind = outlier_selection_
+            ? EncodingOracleKind::AdaptiveFixedRateBitpack
+            : EncodingOracleKind::PlainFixedRateBitpack;
+        d.op_name = outlier_selection_ ? "AdaptiveBitpackCoder" : "PlainBitpackCoder";
+        d.include_header = "fused/fused_block/warp_fusion.cuh";
+        d.input_data_type = static_cast<uint8_t>(getElementDataType());
+        d.unit_elems = block_size_;
+        d.exact = true;
+        d.additive = true;
+        return d;
+    }
+
     /// Choose the warp coder policy the fused path composes (default
     /// "AdaptiveBitpackCoder"). Any policy in warp_fusion.cuh that emits an
     /// AdaptiveBitpack-decodable archive works — e.g. "PlainBitpackCoder" for an
@@ -139,6 +158,9 @@ class AdaptiveBitpackStage : public Stage {
         num_elements_       = num_elements;
         actual_output_size_ = archive_bytes;
     }
+    void setFusedArchiveResult(size_t archive_bytes, size_t orig_bytes) override {
+        setFusedResult(orig_bytes / sizeof(T), archive_bytes);
+    }
 
     /// Enable cuSZp2 per-block plain/outlier selection: each block may instead
     /// store element 0 as a raw 1..sizeof(T)-byte outlier and pack only the rest,

modules/fused/adaptive_lorenzo/adaptive_lorenzo_stage.cu

@@ -23,6 +23,7 @@
  */
 
 #include "fused/adaptive_lorenzo/adaptive_lorenzo_stage.h"
+#include "coders/adaptive_bitpack/adaptive_bitpack_oracle.cuh"
 #include "stage/stage_registry.h"
 #include <cstring>
 #include <algorithm>
@@ -56,14 +57,33 @@ __device__ __forceinline__ typename std::make_unsigned<T>::type absU(T v) {
     return (v < 0) ? static_cast<U>(~uv + static_cast<U>(1)) : uv;
 }
 
-__device__ __forceinline__ int bitWidth32(uint32_t x) {
-    return x ? (32 - __clz(x)) : 0;
-}
-
-// AdaptiveBitpackStage's per-block encoded size: nothing when every residual is
-// zero, else a sign bitmap plus `r` bit-planes, one 32-bit word each.
-__device__ __forceinline__ uint32_t blockCost(int r) {
-    return (r > 0) ? 4u * (static_cast<uint32_t>(r) + 1u) : 0u;
+struct CoderStats {
+    uint32_t all;
+    uint32_t rest;
+    uint32_t first;
+};
+
+// Exact policy dispatch shared by staged selection and the future generated
+// tile harness. The quote carries the coder decision even though the staged
+// predictor needs only its payload length; the downstream staged coder will
+// recompute the same decision from the selected residuals.
+__device__ __forceinline__ uint32_t blockCost(
+    CoderStats stats, EncodingOracleKind oracle_kind)
+{
+    switch (oracle_kind) {
+        case EncodingOracleKind::AdaptiveFixedRateBitpack:
+            return adaptive_bitpack_oracle::quoteAdaptiveFixedRate(
+                stats.all, stats.rest, stats.first,
+                /*word_bytes=*/4u).payload_bytes;
+        case EncodingOracleKind::PlainFixedRateBitpack:
+            return adaptive_bitpack_oracle::quotePlainFixedRate(
+                stats.all, /*word_bytes=*/4u).payload_bytes;
+        default:
+            // Standalone/backward-compatible AdaptiveLorenzo uses the legacy
+            // plain policy. Unsupported downstream policies are not bound.
+            return adaptive_bitpack_oracle::quotePlainFixedRate(
+                stats.all, /*word_bytes=*/4u).payload_bytes;
+    }
 }
 
 // Modes are packed 2 bits per tile (4 tiles per byte): at one byte per tile the
@@ -131,14 +151,19 @@ __global__ void adaptive_lorenzo_forward_kernel(
     size_t n,
     uint32_t tile_size,
     bool enable_order2,
-    bool enable_centering)
+    bool enable_centering,
+    EncodingOracleKind oracle_kind)
 {
     // All shared state is now per-WARP, not per-element: nwarps <= 32 because a
     // tile is at most 32 coder blocks of 32. The tile-sized `s` staging buffer and
     // the tile-sized `red` reduction buffer are both gone, so this kernel needs no
     // dynamic shared memory at all (see the launch site).
-    __shared__ uint32_t  acc1[kMaxBlocksPerTile];    // per coder block, LZ1
-    __shared__ uint32_t  acc2[kMaxBlocksPerTile];    // per coder block, LZ2
+    __shared__ uint32_t  acc1[kMaxBlocksPerTile];    // all magnitudes, LZ1
+    __shared__ uint32_t  acc2[kMaxBlocksPerTile];    // all magnitudes, LZ2
+    __shared__ uint32_t  rest1[kMaxBlocksPerTile];   // lanes 1..31, LZ1
+    __shared__ uint32_t  rest2[kMaxBlocksPerTile];   // lanes 1..31, LZ2
+    __shared__ uint32_t  first1[kMaxBlocksPerTile];  // lane 0 magnitude, LZ1
+    __shared__ uint32_t  first2[kMaxBlocksPerTile];  // lane 0 magnitude, LZ2
     __shared__ long long red[kMaxBlocksPerTile];     // per-warp partial sums (mean)
     __shared__ T         sb_last[kMaxBlocksPerTile]; // v at lane 31 of each warp
     __shared__ T         sb_prev[kMaxBlocksPerTile]; // v at lane 30 of each warp
@@ -223,29 +248,42 @@ __global__ void adaptive_lorenzo_forward_kernel(
     const T q0 = s_q0;
 
     // ---- Per-coder-block magnitudes, uncentered ----
-    const uint32_t o1 = warpOr(live ? static_cast<uint32_t>(absU<T>(d1)) : 0u);
-    const uint32_t o2 = warpOr(live ? static_cast<uint32_t>(absU<T>(d2)) : 0u);
-    if (lane == 0) { acc1[warp] = o1; acc2[warp] = o2; }
+    const uint32_t m1 = live ? static_cast<uint32_t>(absU<T>(d1)) : 0u;
+    const uint32_t m2 = live ? static_cast<uint32_t>(absU<T>(d2)) : 0u;
+    const uint32_t o1 = warpOr(m1);
+    const uint32_t o2 = warpOr(m2);
+    const uint32_t r1rest = warpOr(lane > 0u ? m1 : 0u);
+    const uint32_t r2rest = warpOr(lane > 0u ? m2 : 0u);
+    if (lane == 0) {
+        acc1[warp] = o1;  rest1[warp] = r1rest; first1[warp] = m1;
+        acc2[warp] = o2;  rest2[warp] = r2rest; first2[warp] = m2;
+    }
 
     // ---- Centered variants: only coder block 0 can differ ----
-    uint32_t acc1c0 = 0u, acc2c0 = 0u;
+    CoderStats c1stats{0u, 0u, 0u}, c2stats{0u, 0u, 0u};
     if (enable_centering && warp == 0u) {
         const T c0 = static_cast<T>(q0 - mu);
-        const T r1 = (tid == 0u) ? c0 : d1;
-        T       r2 = d2;
-        if      (tid == 0u) r2 = c0;
-        else if (tid == 1u) r2 = static_cast<T>(d1 - c0);
-        acc1c0 = warpOr(live ? static_cast<uint32_t>(absU<T>(r1)) : 0u);
-        acc2c0 = warpOr(live ? static_cast<uint32_t>(absU<T>(r2)) : 0u);
+        const T cr1 = (tid == 0u) ? c0 : d1;
+        T       cr2 = d2;
+        if      (tid == 0u) cr2 = c0;
+        else if (tid == 1u) cr2 = static_cast<T>(d1 - c0);
+        const uint32_t cm1 = live ? static_cast<uint32_t>(absU<T>(cr1)) : 0u;
+        const uint32_t cm2 = live ? static_cast<uint32_t>(absU<T>(cr2)) : 0u;
+        c1stats.all   = warpOr(cm1);
+        c1stats.rest  = warpOr(lane > 0u ? cm1 : 0u);
+        c1stats.first = fz::backend::shfl(cm1, 0, 32);
+        c2stats.all   = warpOr(cm2);
+        c2stats.rest  = warpOr(lane > 0u ? cm2 : 0u);
+        c2stats.first = fz::backend::shfl(cm2, 0, 32);
     }
     __syncthreads();
 
     // ---- Cost each variant, pick the cheapest ----
     if (tid == 0) {
         uint32_t c_lz1 = 0, c_lz2 = 0;
         for (unsigned w = 0; w < nwarps; ++w) {
-            c_lz1 += blockCost(bitWidth32(acc1[w]));
-            c_lz2 += blockCost(bitWidth32(acc2[w]));
+            c_lz1 += blockCost(CoderStats{acc1[w], rest1[w], first1[w]}, oracle_kind);
+            c_lz2 += blockCost(CoderStats{acc2[w], rest2[w], first2[w]}, oracle_kind);
         }
         uint32_t costs[4];
         costs[0] = c_lz1;
@@ -259,11 +297,13 @@ __global__ void adaptive_lorenzo_forward_kernel(
             // genuinely saves those bytes. (The 2-bit mode is charged to nobody
             // because every tile pays it regardless of what it picks.)
             const uint32_t mean_cost = static_cast<uint32_t>(sizeof(T));
-            costs[2] = c_lz1 - blockCost(bitWidth32(acc1[0]))
-                             + blockCost(bitWidth32(acc1c0)) + mean_cost;
+            costs[2] = c_lz1
+                - blockCost(CoderStats{acc1[0], rest1[0], first1[0]}, oracle_kind)
+                + blockCost(c1stats, oracle_kind) + mean_cost;
             if (enable_order2)
-                costs[3] = c_lz2 - blockCost(bitWidth32(acc2[0]))
-                                 + blockCost(bitWidth32(acc2c0)) + mean_cost;
+                costs[3] = c_lz2
+                    - blockCost(CoderStats{acc2[0], rest2[0], first2[0]}, oracle_kind)
+                    + blockCost(c2stats, oracle_kind) + mean_cost;
         }
         uint32_t best = 0;
         for (uint32_t i = 1; i < 4; ++i)
@@ -384,15 +424,16 @@ template<typename T>
 void launchAdaptiveLorenzoForward(
     const T* d_input, T* d_residuals, uint8_t* d_modes_dense, T* d_means_dense,
     uint32_t* d_flags, size_t n, uint32_t tile_size,
-    bool enable_order2, bool enable_centering, cudaStream_t stream)
+    bool enable_order2, bool enable_centering, EncodingOracleKind oracle_kind,
+    cudaStream_t stream)
 {
     if (n == 0) return;
     const int grid = static_cast<int>((n + tile_size - 1) / tile_size);
     // No dynamic shared memory: the kernel's shared state is now per-warp
     // (<= 32 entries each) and lives in static __shared__ arrays.
     adaptive_lorenzo_forward_kernel<T><<<grid, tile_size, 0, stream>>>(
         d_input, d_residuals, d_modes_dense, d_means_dense, d_flags, n, tile_size,
-        enable_order2, enable_centering);
+        enable_order2, enable_centering, oracle_kind);
     FZ_CUDA_CHECK(cudaGetLastError());
 }
 
@@ -526,7 +567,8 @@ void AdaptiveLorenzoStage<T>::execute(
         launchAdaptiveLorenzoForward<T>(
             static_cast<const T*>(inputs[0]), static_cast<T*>(outputs[0]),
             d_modes_dense_, d_means_dense_, d_flags_,
-            n, tile, config_.enable_order2, config_.enable_centering, stream);
+            n, tile, config_.enable_order2, config_.enable_centering,
+            getBoundEncodingOracleKind(), stream);
 
         // flags[tiles] = 0 so offsets[tiles] lands on the total centered count.
         FZ_CUDA_CHECK(cudaMemsetAsync(d_flags_ + tiles, 0, sizeof(uint32_t), stream));
@@ -600,9 +642,11 @@ template class AdaptiveLorenzoStage<int16_t>;
 template class AdaptiveLorenzoStage<int32_t>;
 
 template void launchAdaptiveLorenzoForward<int16_t>(
-    const int16_t*, int16_t*, uint8_t*, int16_t*, uint32_t*, size_t, uint32_t, bool, bool, cudaStream_t);
+    const int16_t*, int16_t*, uint8_t*, int16_t*, uint32_t*, size_t, uint32_t,
+    bool, bool, EncodingOracleKind, cudaStream_t);
 template void launchAdaptiveLorenzoForward<int32_t>(
-    const int32_t*, int32_t*, uint8_t*, int32_t*, uint32_t*, size_t, uint32_t, bool, bool, cudaStream_t);
+    const int32_t*, int32_t*, uint8_t*, int32_t*, uint32_t*, size_t, uint32_t,
+    bool, bool, EncodingOracleKind, cudaStream_t);
 
 template void launchAdaptiveLorenzoCompact<int16_t>(
     const uint8_t*, const int16_t*, const uint32_t*, uint8_t*, int16_t*, size_t, cudaStream_t);

modules/fused/adaptive_lorenzo/adaptive_lorenzo_stage.h

@@ -122,6 +122,46 @@ class AdaptiveLorenzoStage : public Stage {
         return config_.coder_block_size * config_.blocks_per_tile;
     }
 
+    /// Bind an exact additive fixed-rate policy exposed by a directly connected
+    /// AdaptiveBitpack stage. Both plain and element-0-outlier selection operate
+    /// independently on the same 32-element coder units.
+    bool bindDownstreamEncodingOracle(const EncodingOracleDecl& decl) override {
+        if (!decl.valid() || !decl.additive ||
+            (decl.kind != EncodingOracleKind::PlainFixedRateBitpack &&
+             decl.kind != EncodingOracleKind::AdaptiveFixedRateBitpack) ||
+            decl.input_data_type != static_cast<uint8_t>(getElementDataType()) ||
+            decl.unit_elems != config_.coder_block_size) {
+            return false;
+        }
+        bound_oracle_ = decl;
+        has_bound_oracle_ = true;
+        return true;
+    }
+
+    bool hasBoundEncodingOracle() const { return has_bound_oracle_; }
+    EncodingOracleKind getBoundEncodingOracleKind() const {
+        return has_bound_oracle_ ? bound_oracle_.kind
+                                 : EncodingOracleKind::PlainFixedRateBitpack;
+    }
+
+    FusionSpec getFusionSpec() const override {
+        if (is_inverse_ || !has_bound_oracle_) return {};
+        return FusionSpec{FusionAccess::TileAdaptive, getTileSize(),
+                          config_.coder_block_size};
+    }
+
+    std::vector<FusedAuxOutputDecl> getFusedAuxOutputs() const override {
+        if (!getFusionSpec().fusable()) return {};
+        return {
+            FusedAuxOutputDecl{1, "modes", FusedAuxSizeKind::FixedBitsPerUnit,
+                               static_cast<uint8_t>(DataType::UINT8), getTileSize(),
+                               2u, 0u},
+            FusedAuxOutputDecl{2, "means", FusedAuxSizeKind::CompactedElements,
+                               static_cast<uint8_t>(getElementDataType()), getTileSize(),
+                               0u, 1u},
+        };
+    }
+
     void execute(
         fz::stream_t stream,
         MemoryPool* pool,
@@ -185,6 +225,12 @@ class AdaptiveLorenzoStage : public Stage {
             ? actual_output_sizes_[index] : 0;
     }
 
+    void setFusedSideOutput(int output_index, size_t bytes) override {
+        if (actual_output_sizes_.size() < 3) actual_output_sizes_.resize(3, 0);
+        if (output_index == 1 || output_index == 2)
+            actual_output_sizes_[static_cast<size_t>(output_index)] = bytes;
+    }
+
     uint16_t getStageTypeId() const override {
         return static_cast<uint16_t>(StageType::ADAPTIVE_LORENZO);
     }
@@ -232,6 +278,8 @@ class AdaptiveLorenzoStage : public Stage {
 
 private:
     Config config_;
+    EncodingOracleDecl bound_oracle_;
+    bool has_bound_oracle_ = false;
     bool   is_inverse_ = false;
     size_t num_elements_ = 0;
     std::vector<size_t> actual_output_sizes_{0, 0, 0};
@@ -283,8 +331,8 @@ extern template class AdaptiveLorenzoStage<int32_t>;
 template<typename T>
 void launchAdaptiveLorenzoForward(
     const T* d_input, T* d_residuals, uint8_t* d_modes, T* d_means,
-    size_t n, uint32_t tile_size, bool enable_order2, bool enable_centering,
-    fz::stream_t stream);
+    uint32_t* d_flags, size_t n, uint32_t tile_size, bool enable_order2,
+    bool enable_centering, EncodingOracleKind oracle_kind, fz::stream_t stream);
 
 /// Inverse: replay each tile's recorded variant.
 template<typename T>