FSE: fuse symbol spreading with table building on the fast path

lib/common/fse.h

@@ -622,4 +622,23 @@ MEM_STATIC unsigned FSE_endOfDState(const FSE_DState_t* DStatePtr)
 
 #define FSE_TABLESTEP(tableSize) (((tableSize)>>1) + ((tableSize)>>3) + 3)
 
+/* FSE_invTableStep() :
+ * Multiplicative inverse of FSE_TABLESTEP(tableSize) modulo tableSize.
+ * FSE_TABLESTEP() is odd for any tableSize >= 16, hence for any valid tableLog
+ * (>= FSE_MIN_TABLELOG), so the inverse exists.
+ * Symbol spreading places the i-th symbol (in symbol order) at position (i*step) & (tableSize-1),
+ * hence position u holds the symbol at index (u*inv) & (tableSize-1).
+ * This allows walking the table in position order while reading symbols
+ * from their sorted layout, instead of scattering them and reading them back. */
+MEM_STATIC U32 FSE_invTableStep(U32 tableSize)
+{
+    U32 const step = FSE_TABLESTEP(tableSize);
+    U32 inv = step;          /* odd*odd == 1 (mod 8) : 3 correct low bits */
+    inv *= 2 - step * inv;   /* each Newton iteration doubles the nb of correct bits : 6 */
+    inv *= 2 - step * inv;   /* 12 */
+    inv *= 2 - step * inv;   /* 24 >= FSE_TABLELOG_ABSOLUTE_MAX */
+    assert(((step * inv) & (tableSize - 1)) == 1);
+    return inv & (tableSize - 1);
+}
+
 #endif /* FSE_STATIC_LINKING_ONLY */

lib/compress/fse_compress.c

@@ -133,23 +133,21 @@ size_t FSE_buildCTable_wksp(FSE_CTable* ct,
                 pos += (size_t)n;
             }
         }
-        /* Spread symbols across the table. Lack of lowprob symbols means that
-         * we don't need variable sized inner loop, so we can unroll the loop and
-         * reduce branch misses.
+        /* Build the table directly, walking positions in order.
+         * Position u holds spread[(u * invStep) & tableMask], where invStep
+         * is the inverse of step modulo tableSize (see FSE_invTableStep()).
+         * This fuses symbol spreading with table building into a single pass,
+         * so symbols are never scattered into tableSymbol[] and read back.
          */
-        {   size_t position = 0;
-            size_t s;
-            size_t const unroll = 2; /* Experimentally determined optimal unroll */
-            assert(tableSize % unroll == 0); /* FSE_MIN_TABLELOG is 5 */
-            for (s = 0; s < (size_t)tableSize; s += unroll) {
-                size_t u;
-                for (u = 0; u < unroll; ++u) {
-                    size_t const uPosition = (position + (u * step)) & tableMask;
-                    tableSymbol[uPosition] = spread[s + u];
-                }
-                position = (position + (unroll * step)) & tableMask;
+        {   size_t const invStep = FSE_invTableStep(tableSize);
+            size_t s = 0;
+            U32 u;
+            for (u=0; u<tableSize; u++) {
+                FSE_FUNCTION_TYPE const symbol = spread[s];
+                tableU16[cumul[symbol]++] = (U16) (tableSize+u);   /* TableU16 : sorted by symbol order; gives next state value */
+                s = (s + invStep) & tableMask;
             }
-            assert(position == 0);   /* Must have initialized all positions */
+            assert(s == 0);   /* Must have visited all positions */
         }
     } else {
         U32 position = 0;
@@ -164,13 +162,13 @@ size_t FSE_buildCTable_wksp(FSE_CTable* ct,
                     position = (position + step) & tableMask;   /* Low proba area */
         }   }
         assert(position==0);  /* Must have initialized all positions */
-    }
 
-    /* Build table */
-    {   U32 u; for (u=0; u<tableSize; u++) {
-        FSE_FUNCTION_TYPE s = tableSymbol[u];   /* note : static analyzer may not understand tableSymbol is properly initialized */
-        tableU16[cumul[s]++] = (U16) (tableSize+u);   /* TableU16 : sorted by symbol order; gives next state value */
-    }   }
+        /* Build table */
+        {   U32 u; for (u=0; u<tableSize; u++) {
+            FSE_FUNCTION_TYPE s = tableSymbol[u];   /* note : static analyzer may not understand tableSymbol is properly initialized */
+            tableU16[cumul[s]++] = (U16) (tableSize+u);   /* TableU16 : sorted by symbol order; gives next state value */
+        }   }
+    }
 
     /* Build Symbol Transformation Table */
     {   unsigned total = 0;

lib/decompress/zstd_decompress_block.c

@@ -528,7 +528,6 @@ void ZSTD_buildFSETable_body(ZSTD_seqSymbol* dt,
      */
     if (highThreshold == tableSize - 1) {
         size_t const tableMask = tableSize-1;
-        size_t const step = FSE_TABLESTEP(tableSize);
         /* First lay down the symbols in order.
          * We use a uint64_t to lay down 8 bytes at a time. This reduces branch
          * misses since small blocks generally have small table logs, so nearly
@@ -551,26 +550,28 @@ void ZSTD_buildFSETable_body(ZSTD_seqSymbol* dt,
                 pos += (size_t)n;
             }
         }
-        /* Now we spread those positions across the table.
-         * The benefit of doing it in two stages is that we avoid the
-         * variable size inner loop, which caused lots of branch misses.
-         * Now we can run through all the positions without any branch misses.
-         * We unroll the loop twice, since that is what empirically worked best.
+        /* Now build the decoding table directly, walking positions in order.
+         * Position u holds spread[(u * invStep) & tableMask], where invStep
+         * is the inverse of step modulo tableSize (see FSE_invTableStep()).
+         * This fuses symbol spreading with table building into a single pass,
+         * with no scattered writes into the table and no read-back of symbols.
          */
         {
-            size_t position = 0;
-            size_t s;
-            size_t const unroll = 2;
-            assert(tableSize % unroll == 0); /* FSE_MIN_TABLELOG is 5 */
-            for (s = 0; s < (size_t)tableSize; s += unroll) {
-                size_t u;
-                for (u = 0; u < unroll; ++u) {
-                    size_t const uPosition = (position + (u * step)) & tableMask;
-                    tableDecode[uPosition].baseValue = spread[s + u];
-                }
-                position = (position + (unroll * step)) & tableMask;
+            size_t const invStep = FSE_invTableStep(tableSize);
+            size_t s = 0;
+            U32 u;
+            for (u=0; u<tableSize; u++) {
+                U32 const symbol = spread[s];
+                U32 const nextState = symbolNext[symbol]++;
+                BYTE const nbBits = (BYTE) (tableLog - ZSTD_highbit32(nextState) );
+                tableDecode[u].nbBits = nbBits;
+                tableDecode[u].nextState = (U16) ( (nextState << nbBits) - tableSize);
+                assert(nbAdditionalBits[symbol] < 255);
+                tableDecode[u].nbAdditionalBits = nbAdditionalBits[symbol];
+                tableDecode[u].baseValue = baseValue[symbol];
+                s = (s + invStep) & tableMask;
             }
-            assert(position == 0);
+            assert(s == 0);   /* Must have visited all positions */
         }
     } else {
         U32 const tableMask = tableSize-1;
@@ -585,19 +586,19 @@ void ZSTD_buildFSETable_body(ZSTD_seqSymbol* dt,
                 while (UNLIKELY(position > highThreshold)) position = (position + step) & tableMask;   /* lowprob area */
         }   }
         assert(position == 0); /* position must reach all cells once, otherwise normalizedCounter is incorrect */
-    }
 
-    /* Build Decoding table */
-    {
-        U32 u;
-        for (u=0; u<tableSize; u++) {
-            U32 const symbol = tableDecode[u].baseValue;
-            U32 const nextState = symbolNext[symbol]++;
-            tableDecode[u].nbBits = (BYTE) (tableLog - ZSTD_highbit32(nextState) );
-            tableDecode[u].nextState = (U16) ( (nextState << tableDecode[u].nbBits) - tableSize);
-            assert(nbAdditionalBits[symbol] < 255);
-            tableDecode[u].nbAdditionalBits = nbAdditionalBits[symbol];
-            tableDecode[u].baseValue = baseValue[symbol];
+        /* Build Decoding table */
+        {
+            U32 u;
+            for (u=0; u<tableSize; u++) {
+                U32 const symbol = tableDecode[u].baseValue;
+                U32 const nextState = symbolNext[symbol]++;
+                tableDecode[u].nbBits = (BYTE) (tableLog - ZSTD_highbit32(nextState) );
+                tableDecode[u].nextState = (U16) ( (nextState << tableDecode[u].nbBits) - tableSize);
+                assert(nbAdditionalBits[symbol] < 255);
+                tableDecode[u].nbAdditionalBits = nbAdditionalBits[symbol];
+                tableDecode[u].baseValue = baseValue[symbol];
+            }
         }
     }
 }

tests/fuzzer.c

@@ -46,6 +46,7 @@
 #include "threading.h"    /* ZSTD_pthread_create, ZSTD_pthread_join */
 #include "compress/hist.h" /* HIST_count_wksp */
 #include "compress/zstd_compress_internal.h" /* ZSTD_get1BlockSummary */
+#include "decompress/zstd_decompress_block.h" /* ZSTD_buildFSETable */
 
 
 /*-************************************
@@ -974,6 +975,104 @@ static unsigned test_get1BlockSummary(unsigned testNb)
     return testNb;
 }
 
+/*-*******************************************************
+*  FSE table construction reference
+*********************************************************/
+/* FSE_buildCTable_wksp() and ZSTD_buildFSETable() build their tables in a
+ * single pass when no symbol has a low probability (-1 count), using the
+ * modular inverse of FSE_TABLESTEP() to visit positions in order.
+ * The functions below implement the plain construction, spreading symbols
+ * with FSE_TABLESTEP() and reading them back in position order, as a reference.
+ */
+
+/* Fills norm[0..maxSV1-1] with a random distribution summing to 1<<tableLog.
+ * When lowProb is set, some symbols (never symbol 0) receive a -1 count. */
+static void FUZ_genNormalizedCount(short* norm, unsigned maxSV1, unsigned tableLog, int lowProb, U32* seed)
+{
+    unsigned remaining = 1u << tableLog;
+    unsigned s;
+    for (s = 0; s < maxSV1; s++) norm[s] = 0;
+    if (lowProb) {
+        for (s = 1; s < maxSV1 && remaining > 1; s++) {
+            if ((FUZ_rand(seed) & 3) == 0) { norm[s] = -1; remaining--; }
+    }   }
+    while (remaining > 0) {
+        s = FUZ_rand(seed) % maxSV1;
+        if (norm[s] >= 0) { norm[s]++; remaining--; }
+    }
+}
+
+/* Reference spreading : low probability symbols at the end of the table,
+ * then each symbol repeated norm[s] times at FSE_TABLESTEP() intervals. */
+static void FUZ_refSpreadSymbols(BYTE* tableSymbol, const short* norm, unsigned maxSV1, unsigned tableLog)
+{
+    U32 const tableSize = 1u << tableLog;
+    U32 const tableMask = tableSize - 1;
+    U32 const step = FSE_TABLESTEP(tableSize);
+    U32 highThreshold = tableSize - 1;
+    U32 s, position = 0;
+    for (s = 0; s < maxSV1; s++) {
+        if (norm[s] == -1) tableSymbol[highThreshold--] = (BYTE)s;
+    }
+    for (s = 0; s < maxSV1; s++) {
+        int i;
+        for (i = 0; i < norm[s]; i++) {
+            tableSymbol[position] = (BYTE)s;
+            position = (position + step) & tableMask;
+            while (position > highThreshold) position = (position + step) & tableMask;
+    }   }
+    assert(position == 0);
+}
+
+/* Reference for the header and state table of an FSE_CTable : the first 2 + tableSize U16.
+ * The symbol transformation table that follows does not depend on the spreading. */
+static void FUZ_refBuildCStateTable(U16* tableU16, const BYTE* tableSymbol, const short* norm, unsigned maxSV1, unsigned tableLog)
+{
+    U32 const tableSize = 1u << tableLog;
+    U16 cumul[MaxSeq + 2];
+    U32 u;
+    tableU16[0] = (U16)tableLog;
+    tableU16[1] = (U16)(maxSV1 - 1);
+    cumul[0] = 0;
+    for (u = 1; u <= maxSV1; u++) {
+        cumul[u] = (U16)(cumul[u-1] + (norm[u-1] == -1 ? 1 : norm[u-1]));
+    }
+    for (u = 0; u < tableSize; u++) {
+        tableU16[2 + cumul[tableSymbol[u]]++] = (U16)(tableSize + u);
+    }
+}
+
+/* Reference for a complete sequence decoding table, header included. */
+static void FUZ_refBuildFSETable(ZSTD_seqSymbol* dt, const BYTE* tableSymbol, const short* norm, unsigned maxSV1,
+                                 const U32* baseValue, const U8* nbAdditionalBits, unsigned tableLog)
+{
+    ZSTD_seqSymbol* const tableDecode = dt + 1;
+    U32 const tableSize = 1u << tableLog;
+    S16 const largeLimit = (S16)(1 << (tableLog - 1));
+    U16 symbolNext[MaxSeq + 1];
+    ZSTD_seqSymbol_header DTableH;
+    U32 s, u;
+    DTableH.tableLog = tableLog;
+    DTableH.fastMode = 1;
+    for (s = 0; s < maxSV1; s++) {
+        if (norm[s] == -1) {
+            symbolNext[s] = 1;
+        } else {
+            if (norm[s] >= largeLimit) DTableH.fastMode = 0;
+            symbolNext[s] = (U16)norm[s];
+    }   }
+    memcpy(dt, &DTableH, sizeof(DTableH));
+    for (u = 0; u < tableSize; u++) {
+        U32 const symbol = tableSymbol[u];
+        U32 const nextState = symbolNext[symbol]++;
+        BYTE const nbBits = (BYTE)(tableLog - ZSTD_highbit32(nextState));
+        tableDecode[u].nbBits = nbBits;
+        tableDecode[u].nextState = (U16)((nextState << nbBits) - tableSize);
+        tableDecode[u].nbAdditionalBits = nbAdditionalBits[symbol];
+        tableDecode[u].baseValue = baseValue[symbol];
+    }
+}
+
 /* ============================================================= */
 
 static int basicUnitTests(U32 const seed, double compressibility)
@@ -4771,6 +4870,50 @@ static int basicUnitTests(U32 const seed, double compressibility)
     }
     DISPLAYLEVEL(3, "OK \n");
 
+    DISPLAYLEVEL(3, "test%3i : FSE table construction matches reference spreading : ", testNb++);
+    {   /* Both builders take a fused single-pass path when no symbol has a -1 count,
+         * and the plain spread-then-read-back path otherwise.
+         * Check that each produces the reference table for every tableLog. */
+        U32 rSeed = seed;
+        short norm[MaxSeq + 1];
+        U32 baseValue[MaxSeq + 1];
+        U8 nbAdditionalBits[MaxSeq + 1];
+        BYTE tableSymbol[FSE_MAX_TABLESIZE];
+        FSE_CTable ct[FSE_CTABLE_SIZE_U32(FSE_MAX_TABLELOG, MaxSeq)];
+        U16 refCt[2 + FSE_MAX_TABLESIZE];
+        U32 cwksp[FSE_BUILD_CTABLE_WORKSPACE_SIZE_U32(MaxSeq, FSE_MAX_TABLELOG)];
+        ZSTD_seqSymbol dt[SEQSYMBOL_TABLE_SIZE(MaxFSELog)];
+        ZSTD_seqSymbol refDt[SEQSYMBOL_TABLE_SIZE(MaxFSELog)];
+        U32 dwksp[ZSTD_BUILD_FSE_TABLE_WKSP_SIZE_U32];
+        unsigned tableLog, s;
+        for (s = 0; s <= MaxSeq; s++) {
+            baseValue[s] = 1 + s * s;
+            nbAdditionalBits[s] = (U8)(s % 32);
+        }
+        for (tableLog = FSE_MIN_TABLELOG; tableLog <= FSE_MAX_TABLELOG; tableLog++) {
+            U32 const tableSize = 1u << tableLog;
+            int lowProb;
+            for (lowProb = 0; lowProb < 2; lowProb++) {
+                int n;
+                for (n = 0; n < 16; n++) {
+                    unsigned const maxSV1 = 2 + FUZ_rand(&rSeed) % MaxSeq;
+                    FUZ_genNormalizedCount(norm, maxSV1, tableLog, lowProb, &rSeed);
+                    FUZ_refSpreadSymbols(tableSymbol, norm, maxSV1, tableLog);
+
+                    CHECK_Z( FSE_buildCTable_wksp(ct, norm, maxSV1 - 1, tableLog, cwksp, sizeof(cwksp)) );
+                    FUZ_refBuildCStateTable(refCt, tableSymbol, norm, maxSV1, tableLog);
+                    if (memcmp(ct, refCt, (2 + tableSize) * sizeof(U16)) != 0) goto _output_error;
+
+                    if (tableLog <= MaxFSELog) {
+                        memset(dt, 0, sizeof(dt));
+                        memset(refDt, 0, sizeof(refDt));
+                        ZSTD_buildFSETable(dt, norm, maxSV1 - 1, baseValue, nbAdditionalBits, tableLog, dwksp, sizeof(dwksp), 0);
+                        FUZ_refBuildFSETable(refDt, tableSymbol, norm, maxSV1, baseValue, nbAdditionalBits, tableLog);
+                        if (memcmp(dt, refDt, (1 + tableSize) * sizeof(ZSTD_seqSymbol)) != 0) goto _output_error;
+        }   }   }   }
+    }
+    DISPLAYLEVEL(3, "OK \n");
+
     DISPLAYLEVEL(3, "test%3i : testing bitwise intrinsics PR#3045: ", testNb++);
     {
         U32 seed_copy = seed; /* need non-const seed to avoid compiler warning for FUZ_rand(&seed) */