Skip to content

Commit 30579ed

Browse files
dfa1claude
dfa1
and
claude
committed
feat(writer): pco E3+E4 — histogram DP + tANS encoder
E3: Histogram + bin optimization DP (port of pco/src/bin_optimization.rs): - PcoHistBin: equal-count histogram buckets from unsigned-sorted latents - PcoBinOptimizer: DP minimizing Σ (bin_meta_cost + (ans_cost + offset_cost)×count) with single-bin and trivial-offset shortcuts; log2_approx port for fast float log E4: tANS weight quantization + encode table (port of pco/src/ans/encoding.rs): - PcoWeightQuantizer: proportional scaling to 2^sizeLog with rounding correction; reduces table when all weights share a common power-of-2 factor - PcoAnsEncoder: builds next-state table from spread_state_symbols; per-symbol renorm bounds; encode() returns (newState, bits, numBits) for LIFO collection Page encoding: LIFO encoding with 4-way interleaved tANS — process elements n-1→0, store bits at posInBatch index (implicit reversal), write forward per batch with offset bits interleaved. Works because posInBatch=i%BATCH_N maps each element's bits to the right position for the forward decode read. Key bug fixed: chunk meta must write quantized ANS weights, not raw bin counts. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent 9aa9803 commit 30579ed

5 files changed

Lines changed: 605 additions & 85 deletions

File tree

‎writer/src/main/java/io/github/dfa1/vortex/writer/encode/PcoAnsEncoder.java‎

Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
1+
package io.github.dfa1.vortex.writer.encode;
2+
3+
/// tANS encode table — port of {@code pco/src/ans/encoding.rs} and {@code spec.rs}.
4+
///
5+
/// <p>State is kept in Rust convention: {@code state ∈ [tableSize, 2*tableSize)}.
6+
/// The page header stores {@code state - tableSize} as the decoder initial state index.
7+
final class PcoAnsEncoder {
8+
9+
/// Result of one ANS encode step.
10+
///
11+
/// @param newState next ANS state (Rust convention: {@code ∈ [tableSize, 2*tableSize)})
12+
/// @param bits low {@code numBits} bits of the old state, to be written to the stream
13+
/// @param numBits number of bits to write (renormalization bits)
14+
record Step(int newState, int bits, int numBits) {
15+
}
16+
17+
private final int sizeLog;
18+
private final int tableSize;
19+
private final int[] minRenormBits;
20+
private final int[] renormBitCutoff;
21+
private final int[][] nextStates; // nextStates[sym][k] = Rust state for k-th occurrence
22+
23+
private PcoAnsEncoder(int sizeLog, int tableSize,
24+
int[] minRenormBits, int[] renormBitCutoff, int[][] nextStates) {
25+
this.sizeLog = sizeLog;
26+
this.tableSize = tableSize;
27+
this.minRenormBits = minRenormBits;
28+
this.renormBitCutoff = renormBitCutoff;
29+
this.nextStates = nextStates;
30+
}
31+
32+
/// Build an encoder from quantized ANS weights (sum must equal {@code 2^sizeLog}).
33+
///
34+
/// @param sizeLog log₂ of the ANS table size
35+
/// @param weights quantized weight per symbol (sum == {@code 1 << sizeLog})
36+
/// @return a ready-to-use encoder
37+
static PcoAnsEncoder build(int sizeLog, int[] weights) {
38+
int tableSize = 1 << sizeLog;
39+
int nSymbols = weights.length;
40+
41+
int[] minR = new int[nSymbols];
42+
int[] cutoff = new int[nSymbols];
43+
int[][] nextSt = new int[nSymbols][];
44+
45+
for (int sym = 0; sym < nSymbols; sym++) {
46+
int w = weights[sym];
47+
int maxXS = 2 * w - 1;
48+
int mr = sizeLog - (31 - Integer.numberOfLeadingZeros(maxXS));
49+
minR[sym] = mr;
50+
cutoff[sym] = 2 * w * (1 << mr);
51+
nextSt[sym] = new int[w];
52+
}
53+
54+
int[] stateSymbols = spreadStateSymbols(weights, tableSize);
55+
int[] symbolXs = weights.clone();
56+
for (int stateIdx = 0; stateIdx < tableSize; stateIdx++) {
57+
int sym = stateSymbols[stateIdx];
58+
int w = weights[sym];
59+
int k = symbolXs[sym] - w; // 0-based position into next_states[sym]
60+
nextSt[sym][k] = tableSize + stateIdx;
61+
symbolXs[sym]++;
62+
}
63+
64+
return new PcoAnsEncoder(sizeLog, tableSize, minR, cutoff, nextSt);
65+
}
66+
67+
/// Encode one symbol. Caller writes the low {@link Step#numBits()} bits of
68+
/// the old state to the bit stream (in LIFO order; caller reverses per batch).
69+
///
70+
/// @param state current ANS state {@code ∈ [tableSize, 2*tableSize)}
71+
/// @param symbol bin symbol index
72+
/// @return encode step
73+
Step encode(int state, int symbol) {
74+
int w = nextStates[symbol].length;
75+
int renormBits = state >= renormBitCutoff[symbol]
76+
? minRenormBits[symbol] + 1
77+
: minRenormBits[symbol];
78+
int bitsVal = state & ((1 << renormBits) - 1);
79+
int xS = state >>> renormBits; // in [w, 2w)
80+
int newState = nextStates[symbol][xS - w];
81+
return new Step(newState, bitsVal, renormBits);
82+
}
83+
84+
/// Starting state for all 4 interleaved streams.
85+
///
86+
/// @return {@code tableSize} (Rust convention initial state)
87+
int defaultState() {
88+
return tableSize;
89+
}
90+
91+
/// Convert a final encoder state to the decoder initial state index written in the page header.
92+
///
93+
/// @param state final Rust-convention state after encoding
94+
/// @return 0-based state index (= {@code state - tableSize})
95+
int toStateIdx(int state) {
96+
return state - tableSize;
97+
}
98+
99+
// Port of PcoTansDecoder.spreadStateSymbols (duplicated: writer cannot depend on reader).
100+
private static int[] spreadStateSymbols(int[] weights, int tableSize) {
101+
int[] stateSymbols = new int[tableSize];
102+
int stride = (3 * tableSize) / 5;
103+
if (stride % 2 == 0) {
104+
stride++;
105+
}
106+
int modMask = tableSize - 1;
107+
int step = 0;
108+
for (int sym = 0; sym < weights.length; sym++) {
109+
for (int k = 0; k < weights[sym]; k++) {
110+
stateSymbols[(stride * step) & modMask] = sym;
111+
step++;
112+
}
113+
}
114+
return stateSymbols;
115+
}
116+
}

‎writer/src/main/java/io/github/dfa1/vortex/writer/encode/PcoBinOptimizer.java‎

Lines changed: 156 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,156 @@
1+
package io.github.dfa1.vortex.writer.encode;
2+
3+
import java.util.ArrayList;
4+
import java.util.Collections;
5+
import java.util.List;
6+
7+
/// Bin optimization DP — port of {@code pco/src/bin_optimization.rs}.
8+
///
9+
/// <p>Input: equal-count histogram bins (sorted by value). Output: merged bins that
10+
/// minimize total bit cost = Σ bin_meta_cost + (ans_cost + offset_cost) × count.
11+
final class PcoBinOptimizer {
12+
13+
private PcoBinOptimizer() {
14+
}
15+
16+
private static final float SINGLE_BIN_SPEEDUP = 0.1f;
17+
private static final float TRIVIAL_OFFSET_SPEEDUP = 0.1f;
18+
19+
/// One optimized ANS bin.
20+
///
21+
/// @param lowerSortKey lower bound in sort-key space (latent XOR Long.MIN_VALUE)
22+
/// @param upperSortKey upper bound in sort-key space
23+
/// @param weight raw element count (quantized to ANS weight separately)
24+
record Bin(long lowerSortKey, long upperSortKey, int weight) {
25+
26+
int offsetBits() {
27+
long range = upperSortKey - lowerSortKey;
28+
return range == 0 ? 0 : 64 - Long.numberOfLeadingZeros(range);
29+
}
30+
31+
long lowerLatent() {
32+
return lowerSortKey ^ Long.MIN_VALUE;
33+
}
34+
}
35+
36+
static List<Bin> optimize(List<PcoHistBin> histBins, int ansSizeLog, int dtypeSize) {
37+
if (histBins.isEmpty()) {
38+
return List.of();
39+
}
40+
int n = histBins.size();
41+
float binMetaCost = ansSizeLog + dtypeSize + bitsToEncodeOffsetBits(dtypeSize);
42+
43+
long total = 0;
44+
for (PcoHistBin b : histBins) {
45+
total += b.count();
46+
}
47+
float totalLog2 = log2Approx((float) total);
48+
49+
long[] cCounts = new long[n + 1];
50+
long[] lowers = new long[n];
51+
long[] uppers = new long[n];
52+
for (int i = 0; i < n; i++) {
53+
cCounts[i + 1] = cCounts[i] + histBins.get(i).count();
54+
lowers[i] = histBins.get(i).lowerSortKey();
55+
uppers[i] = histBins.get(i).upperSortKey();
56+
}
57+
58+
float[] bestCosts = new float[n + 1];
59+
int[] bestJs = new int[n];
60+
61+
for (int i = 0; i < n; i++) {
62+
float best = Float.MAX_VALUE;
63+
int bestJ = 0;
64+
for (int j = i; j >= 0; j--) {
65+
long count = cCounts[i + 1] - cCounts[j];
66+
float cost = bestCosts[j] + binCost(binMetaCost, lowers[j], uppers[i], count, totalLog2);
67+
if (cost < best) {
68+
best = cost;
69+
bestJ = j;
70+
}
71+
}
72+
bestCosts[i + 1] = best;
73+
bestJs[i] = bestJ;
74+
}
75+
76+
float singleCost = binCost(binMetaCost, lowers[0], uppers[n - 1], total, totalLog2);
77+
if (singleCost < bestCosts[n] + SINGLE_BIN_SPEEDUP * total) {
78+
return List.of(new Bin(lowers[0], uppers[n - 1], (int) total));
79+
}
80+
81+
boolean allTrivial = true;
82+
for (PcoHistBin b : histBins) {
83+
if (b.lowerSortKey() != b.upperSortKey()) {
84+
allTrivial = false;
85+
break;
86+
}
87+
}
88+
if (allTrivial) {
89+
float trivialCost = 0f;
90+
for (PcoHistBin b : histBins) {
91+
trivialCost += binCost(binMetaCost, b.lowerSortKey(), b.upperSortKey(), b.count(), totalLog2);
92+
}
93+
if (trivialCost < bestCosts[n] + TRIVIAL_OFFSET_SPEEDUP * total) {
94+
List<Bin> trivial = new ArrayList<>(n);
95+
for (PcoHistBin b : histBins) {
96+
trivial.add(new Bin(b.lowerSortKey(), b.upperSortKey(), (int) b.count()));
97+
}
98+
return trivial;
99+
}
100+
}
101+
102+
List<int[]> partitions = new ArrayList<>();
103+
int idx = n - 1;
104+
while (true) {
105+
int j = bestJs[idx];
106+
partitions.add(new int[]{j, idx});
107+
if (j == 0) {
108+
break;
109+
}
110+
idx = j - 1;
111+
}
112+
Collections.reverse(partitions);
113+
114+
List<Bin> result = new ArrayList<>(partitions.size());
115+
for (int[] part : partitions) {
116+
long count = cCounts[part[1] + 1] - cCounts[part[0]];
117+
result.add(new Bin(lowers[part[0]], uppers[part[1]], (int) count));
118+
}
119+
return result;
120+
}
121+
122+
private static float binCost(float metaCost, long lower, long upper, long count, float totalLog2) {
123+
float ansLoss = totalLog2 - log2Approx((float) count);
124+
long range = upper - lower;
125+
float offsetCost = range == 0 ? 0f : 64 - Long.numberOfLeadingZeros(range);
126+
return metaCost + (ansLoss + offsetCost) * count;
127+
}
128+
129+
// Port of Rust's log2_approx (pco/src/bin_optimization.rs).
130+
static float log2Approx(float x) {
131+
final float z = 0.674f;
132+
final int signifMask = 0x7FFFFF;
133+
final int zSignif = Float.floatToRawIntBits(z) & signifMask;
134+
final float b = 2.0f / z;
135+
final float c = -b / (6.0f * z);
136+
final float a = -b - c;
137+
138+
int rawBits = Float.floatToRawIntBits(x);
139+
int exp = rawBits >>> 23;
140+
int signif = rawBits & signifMask;
141+
142+
int highBit = signif > zSignif ? 1 : 0;
143+
int logInt = exp + highBit - 127;
144+
float normalized = Float.intBitsToFloat(((0x7F ^ highBit) << 23) | signif);
145+
return logInt + a + normalized * (b + c * normalized);
146+
}
147+
148+
private static int bitsToEncodeOffsetBits(int dtypeSize) {
149+
return switch (dtypeSize) {
150+
case 64 -> 7;
151+
case 32 -> 6;
152+
case 16 -> 5;
153+
default -> throw new IllegalArgumentException("invalid dtypeSize: " + dtypeSize);
154+
};
155+
}
156+
}

0 commit comments

Comments
 (0)