lazy: speed up the row matchfinder (levels 5-12) by 2-8%, identical output

lib/compress/zstd_compress.c

@@ -2719,17 +2719,69 @@ static void ZSTD_reduceIndex (ZSTD_MatchState_t* ms, ZSTD_CCtx_params const* par
 
 /* See doc/zstd_compression_format.md for detailed format description */
 
+#if defined(ZSTD_ARCH_ARM_NEON)
+/* ZSTD_seqToCodes_8():
+ * Vectorized equivalent of the scalar loop in ZSTD_seqToCodes() for 8 sequences.
+ * Small lengths are coded by counting the code thresholds they reach, which matches
+ * the LL_Code[] / ML_Code[] tables of ZSTD_LLcode() / ZSTD_MLcode():
+ *   LL : ll < 16 ? ll : ll < 64  ? 16 + #{18,20,22,24,28,32,40,48 <= ll}         : highbit(ll) + 19
+ *   ML : ml < 32 ? ml : ml < 128 ? 32 + #{34,36,38,40,44,48,56,64,80,96 <= ml}   : highbit(ml) + 36
+ * Requires a little-endian 64-bit target (longOffsets can't happen there).
+ * NEON only: an SSE2 version measured slower than the scalar loop on x86 (no vector lzcnt, 4 lanes). */
+
+/* adds 1 to acc for each lane where x >= t */
+#define ZSTD_COUNT_GE(acc, x, t) acc = vsubq_u16(acc, vcgeq_u16(x, vdupq_n_u16(t)))
+
+FORCE_INLINE_TEMPLATE void
+ZSTD_seqToCodes_8(const SeqDef* seqs, BYTE* llCode, BYTE* ofCode, BYTE* mlCode)
+{
+    /* de-interleave: offBase low half, offBase high half, litLength, mlBase */
+    uint16x8x4_t const s = vld4q_u16((const U16*)(const void*)seqs);
+    uint16x8_t const ll = s.val[2];
+    uint16x8_t const ml = s.val[3];
+
+    {   /* offBase != 0, so if the high half is zero the low half is not */
+        uint16x8_t const hbLo = vsubq_u16(vdupq_n_u16(15), vclzq_u16(s.val[0]));
+        uint16x8_t const hbHi = vsubq_u16(vdupq_n_u16(31), vclzq_u16(s.val[1]));
+        vst1_u8(ofCode, vmovn_u16(vbslq_u16(vtstq_u16(s.val[1], s.val[1]), hbHi, hbLo)));
+    }
+    {   uint16x8_t mid = vdupq_n_u16(16);
+        uint16x8_t const big = vsubq_u16(vdupq_n_u16(15 + 19), vclzq_u16(ll));
+        ZSTD_COUNT_GE(mid, ll, 18); ZSTD_COUNT_GE(mid, ll, 20); ZSTD_COUNT_GE(mid, ll, 22); ZSTD_COUNT_GE(mid, ll, 24);
+        ZSTD_COUNT_GE(mid, ll, 28); ZSTD_COUNT_GE(mid, ll, 32); ZSTD_COUNT_GE(mid, ll, 40); ZSTD_COUNT_GE(mid, ll, 48);
+        vst1_u8(llCode, vmovn_u16(vbslq_u16(vcltq_u16(ll, vdupq_n_u16(16)), ll,
+                                  vbslq_u16(vcltq_u16(ll, vdupq_n_u16(64)), mid, big))));
+    }
+    {   uint16x8_t mid = vdupq_n_u16(32);
+        uint16x8_t const big = vsubq_u16(vdupq_n_u16(15 + 36), vclzq_u16(ml));
+        ZSTD_COUNT_GE(mid, ml, 34); ZSTD_COUNT_GE(mid, ml, 36); ZSTD_COUNT_GE(mid, ml, 38); ZSTD_COUNT_GE(mid, ml, 40);
+        ZSTD_COUNT_GE(mid, ml, 44); ZSTD_COUNT_GE(mid, ml, 48); ZSTD_COUNT_GE(mid, ml, 56); ZSTD_COUNT_GE(mid, ml, 64);
+        ZSTD_COUNT_GE(mid, ml, 80); ZSTD_COUNT_GE(mid, ml, 96);
+        vst1_u8(mlCode, vmovn_u16(vbslq_u16(vcltq_u16(ml, vdupq_n_u16(32)), ml,
+                                  vbslq_u16(vcltq_u16(ml, vdupq_n_u16(128)), mid, big))));
+    }
+}
+#undef ZSTD_COUNT_GE
+
+#endif /* ZSTD_ARCH_ARM_NEON */
+
 int ZSTD_seqToCodes(const SeqStore_t* seqStorePtr)
 {
     const SeqDef* const sequences = seqStorePtr->sequencesStart;
     BYTE* const llCodeTable = seqStorePtr->llCode;
     BYTE* const ofCodeTable = seqStorePtr->ofCode;
     BYTE* const mlCodeTable = seqStorePtr->mlCode;
     U32 const nbSeq = (U32)(seqStorePtr->sequences - seqStorePtr->sequencesStart);
-    U32 u;
+    U32 u = 0;
     int longOffsets = 0;
     assert(nbSeq <= seqStorePtr->maxNbSeq);
-    for (u=0; u<nbSeq; u++) {
+#if defined(ZSTD_ARCH_ARM_NEON)
+    if (MEM_64bits() && MEM_isLittleEndian()) {
+        for (; u + 8 <= nbSeq; u += 8)
+            ZSTD_seqToCodes_8(sequences + u, llCodeTable + u, ofCodeTable + u, mlCodeTable + u);
+    }
+#endif
+    for (; u<nbSeq; u++) {
         U32 const llv = sequences[u].litLength;
         U32 const ofCode = ZSTD_highbit32(sequences[u].offBase);
         U32 const mlv = sequences[u].mlBase;

lib/compress/zstd_compress_internal.h

@@ -897,6 +897,27 @@ MEM_STATIC size_t ZSTD_count(const BYTE* pIn, const BYTE* pMatch, const BYTE* co
     return (size_t)(pIn - pStart);
 }
 
+/** ZSTD_countShort() :
+ *  Same result as ZSTD_count(), for callers whose matches are mostly short: the first 16 bytes
+ *  are resolved with selects rather than data-dependent branches, which mispredict heavily when
+ *  match lengths are unpredictable. In the lazy row matchfinder ~85% of lengths are below 16.
+ *  Slower than ZSTD_count() when matches are typically long (e.g. the binary tree strategies). */
+MEM_STATIC size_t ZSTD_countShort(const BYTE* pIn, const BYTE* pMatch, const BYTE* const pInLimit)
+{
+    if (MEM_64bits() && (pInLimit - pIn) >= 16) {
+        U64 const d0 = MEM_read64(pMatch) ^ MEM_read64(pIn);
+        U64 const d1 = MEM_read64(pMatch + 8) ^ MEM_read64(pIn + 8);
+        if (d0 | d1) {
+            /* first differing word, and its byte offset; bit 63 keeps ctz defined without changing the result */
+            U64 const m = d0 ? d0 : d1;
+            size_t const base = d0 ? 0 : 8;
+            return base + (ZSTD_countTrailingZeros64(m | ((U64)1 << 63)) >> 3);
+        }
+        return 16 + ZSTD_count(pIn + 16, pMatch + 16, pInLimit);
+    }
+    return ZSTD_count(pIn, pMatch, pInLimit);
+}
+
 /** ZSTD_count_2segments() :
  *  can count match length with `ip` & `match` in 2 different segments.
  *  convention : on reaching mEnd, match count continue starting from iStart

lib/compress/zstd_lazy.c

@@ -791,17 +791,25 @@ MEM_STATIC U32 ZSTD_VecMask_next(ZSTD_VecMask val) {
     return ZSTD_countTrailingZeros64(val);
 }
 
-/* ZSTD_row_nextIndex():
- * Returns the next index to insert at within a tagTable row, and updates the "head"
- * value to reflect the update. Essentially cycles backwards from [1, {entries per row})
+/* ZSTD_row_nextIndexFromHead():
+ * Returns the next index to insert at within a tagTable row, given its current "head"
+ * value (*tagRow), and updates the head. Essentially cycles backwards from [1, {entries per row}).
+ * Indices are size_t so they can be used in addresses without zero-extension.
  */
-FORCE_INLINE_TEMPLATE U32 ZSTD_row_nextIndex(BYTE* const tagRow, U32 const rowMask) {
-    U32 next = (*tagRow-1) & rowMask;
+FORCE_INLINE_TEMPLATE size_t ZSTD_row_nextIndexFromHead(BYTE* const tagRow, size_t const head, U32 const rowMask) {
+    size_t next = (head - 1) & rowMask;
     next += (next == 0) ? rowMask : 0; /* skip first position */
     *tagRow = (BYTE)next;
     return next;
 }
 
+/* ZSTD_row_nextIndex():
+ * Same as ZSTD_row_nextIndexFromHead(), reading the head from the row.
+ */
+FORCE_INLINE_TEMPLATE size_t ZSTD_row_nextIndex(BYTE* const tagRow, U32 const rowMask) {
+    return ZSTD_row_nextIndexFromHead(tagRow, *tagRow, rowMask);
+}
+
 /* ZSTD_isAligned():
  * Checks that a pointer is aligned to "align" bytes which must be a power of 2.
  */
@@ -889,19 +897,22 @@ void ZSTD_row_update_internalImpl(ZSTD_MatchState_t* ms,
 {
     U32* const hashTable = ms->hashTable;
     BYTE* const tagTable = ms->tagTable;
+    U32* const hashCache = ms->hashCache;
     U32 const hashLog = ms->rowHashLog;
     const BYTE* const base = ms->window.base;
+    /* Read once: the BYTE stores into tagTable below may alias *ms, which would force a reload per position */
+    U64 const hashSalt = ms->hashSalt;
 
     DEBUGLOG(6, "ZSTD_row_update_internalImpl(): updateStartIdx=%u, updateEndIdx=%u", updateStartIdx, updateEndIdx);
     for (; updateStartIdx < updateEndIdx; ++updateStartIdx) {
-        U32 const hash = useCache ? ZSTD_row_nextCachedHash(ms->hashCache, hashTable, tagTable, base, updateStartIdx, hashLog, rowLog, mls, ms->hashSalt)
-                                  : (U32)ZSTD_hashPtrSalted(base + updateStartIdx, hashLog + ZSTD_ROW_HASH_TAG_BITS, mls, ms->hashSalt);
-        U32 const relRow = (hash >> ZSTD_ROW_HASH_TAG_BITS) << rowLog;
+        U32 const hash = useCache ? ZSTD_row_nextCachedHash(hashCache, hashTable, tagTable, base, updateStartIdx, hashLog, rowLog, mls, hashSalt)
+                                  : (U32)ZSTD_hashPtrSalted(base + updateStartIdx, hashLog + ZSTD_ROW_HASH_TAG_BITS, mls, hashSalt);
+        size_t const relRow = (size_t)(hash >> ZSTD_ROW_HASH_TAG_BITS) << rowLog;
         U32* const row = hashTable + relRow;
         BYTE* tagRow = tagTable + relRow;
-        U32 const pos = ZSTD_row_nextIndex(tagRow, rowMask);
+        size_t const pos = ZSTD_row_nextIndex(tagRow, rowMask);
 
-        assert(hash == ZSTD_hashPtrSalted(base + updateStartIdx, hashLog + ZSTD_ROW_HASH_TAG_BITS, mls, ms->hashSalt));
+        assert(hash == ZSTD_hashPtrSalted(base + updateStartIdx, hashLog + ZSTD_ROW_HASH_TAG_BITS, mls, hashSalt));
         tagRow[pos] = hash & ZSTD_ROW_HASH_TAG_MASK;
         row[pos] = updateStartIdx;
     }
@@ -1165,6 +1176,18 @@ ZSTD_row_getMatchMask(const BYTE* const tagRow, const BYTE tag, const U32 headGr
 #endif
 }
 
+/* ZSTD_row_maskOutHead():
+ * Slot 0 of a row holds the row's head index rather than an entry, so it is never a candidate.
+ * Clears its bit group from a mask returned by ZSTD_row_getMatchMask() (which is rotated by
+ * head * groupWidth), so the candidate loops don't need to test for it on every tag hit. */
+FORCE_INLINE_TEMPLATE ZSTD_VecMask
+ZSTD_row_maskOutHead(ZSTD_VecMask matches, U32 const head, U32 const rowEntries, U32 const groupWidth)
+{
+    U32 const headSlotBit = ((rowEntries - head) & (rowEntries - 1)) * groupWidth;
+    assert(headSlotBit + groupWidth <= sizeof(ZSTD_VecMask) * 8);
+    return matches & ~((((ZSTD_VecMask)1 << groupWidth) - 1) << headSlotBit);
+}
+
 /* The high-level approach of the SIMD row based match finder is as follows:
  * - Figure out where to insert the new entry:
  *      - Generate a hash for current input position and split it into a one byte of tag and `rowHashLog` bits of index.
@@ -1259,21 +1282,23 @@ size_t ZSTD_RowFindBestMatch(
     ms->hashSaltEntropy += hash; /* collect salt entropy */
 
     {   /* Get the hash for ip, compute the appropriate row */
-        U32 const relRow = (hash >> ZSTD_ROW_HASH_TAG_BITS) << rowLog;
+        size_t const relRow = (size_t)(hash >> ZSTD_ROW_HASH_TAG_BITS) << rowLog;
         U32 const tag = hash & ZSTD_ROW_HASH_TAG_MASK;
         U32* const row = hashTable + relRow;
         BYTE* tagRow = (BYTE*)(tagTable + relRow);
-        U32 const headGrouped = (*tagRow & rowMask) * groupWidth;
+        U32 const head = *tagRow & rowMask;
+        U32 const headGrouped = head * groupWidth;
         U32 matchBuffer[ZSTD_ROW_HASH_MAX_ENTRIES];
         size_t numMatches = 0;
         size_t currMatch = 0;
-        ZSTD_VecMask matches = ZSTD_row_getMatchMask(tagRow, (BYTE)tag, headGrouped, rowEntries);
+        ZSTD_VecMask matches = ZSTD_row_maskOutHead(
+                ZSTD_row_getMatchMask(tagRow, (BYTE)tag, headGrouped, rowEntries), head, rowEntries, groupWidth);
 
         /* Cycle through the matches and prefetch */
         for (; (matches > 0) && (nbAttempts > 0); matches &= (matches - 1)) {
             U32 const matchPos = ((headGrouped + ZSTD_VecMask_next(matches)) / groupWidth) & rowMask;
             U32 const matchIndex = row[matchPos];
-            if(matchPos == 0) continue;
+            assert(matchPos != 0);
             assert(numMatches < rowEntries);
             if (matchIndex < lowLimit)
                 break;
@@ -1289,7 +1314,8 @@ size_t ZSTD_RowFindBestMatch(
         /* Speed opt: insert current byte into hashtable too. This allows us to avoid one iteration of the loop
            in ZSTD_row_update_internal() at the next search. */
         {
-            U32 const pos = ZSTD_row_nextIndex(tagRow, rowMask);
+            /* nothing above wrote to tagRow, so the head is still the one the mask was built from */
+            size_t const pos = ZSTD_row_nextIndexFromHead(tagRow, head, rowMask);
             tagRow[pos] = (BYTE)tag;
             row[pos] = ms->nextToUpdate++;
         }
@@ -1306,7 +1332,7 @@ size_t ZSTD_RowFindBestMatch(
                 assert(matchIndex >= dictLimit);   /* ensures this is true if dictMode != ZSTD_extDict */
                 /* read 4B starting from (match + ml + 1 - sizeof(U32)) */
                 if (MEM_read32(match + ml - 3) == MEM_read32(ip + ml - 3))   /* potentially better */
-                    currentMl = ZSTD_count(ip, match, iLimit);
+                    currentMl = ZSTD_countShort(ip, match, iLimit);
             } else {
                 const BYTE* const match = dictBase + matchIndex;
                 assert(match+4 <= dictEnd);
@@ -1335,16 +1361,18 @@ size_t ZSTD_RowFindBestMatch(
         const U32 dmsSize              = (U32)(dmsEnd - dmsBase);
         const U32 dmsIndexDelta        = dictLimit - dmsSize;
 
-        {   U32 const headGrouped = (*dmsTagRow & rowMask) * groupWidth;
+        {   U32 const head = *dmsTagRow & rowMask;
+            U32 const headGrouped = head * groupWidth;
             U32 matchBuffer[ZSTD_ROW_HASH_MAX_ENTRIES];
             size_t numMatches = 0;
             size_t currMatch = 0;
-            ZSTD_VecMask matches = ZSTD_row_getMatchMask(dmsTagRow, (BYTE)dmsTag, headGrouped, rowEntries);
+            ZSTD_VecMask matches = ZSTD_row_maskOutHead(
+                    ZSTD_row_getMatchMask(dmsTagRow, (BYTE)dmsTag, headGrouped, rowEntries), head, rowEntries, groupWidth);
 
             for (; (matches > 0) && (nbAttempts > 0); matches &= (matches - 1)) {
                 U32 const matchPos = ((headGrouped + ZSTD_VecMask_next(matches)) / groupWidth) & rowMask;
                 U32 const matchIndex = dmsRow[matchPos];
-                if(matchPos == 0) continue;
+                assert(matchPos != 0);
                 if (matchIndex < dmsLowestIndex)
                     break;
                 PREFETCH_L1(dmsBase + matchIndex);
@@ -1426,8 +1454,21 @@ size_t ZSTD_RowFindBestMatch(
         return ZSTD_HcFindBestMatch(ms, ip, iLimit, offsetPtr, mls, ZSTD_##dictMode); \
     }                                                                                 \
 
+/* On x86-64, the per-call cost of the outlined search (callee-saved register spills, dispatch,
+ * re-deriving constants from ms) measured ~3% of level 6 compression time, so the hot noDict row
+ * search is inlined into the lazy parsers there. Measured neutral on aarch64, so it stays outlined
+ * on other targets, as do the other modes, to limit code size. */
+#if defined(__x86_64__) || defined(_M_X64)
+#define ZSTD_ROW_SEARCH_FN_ATTRS_noDict FORCE_INLINE_ATTR
+#else
+#define ZSTD_ROW_SEARCH_FN_ATTRS_noDict ZSTD_SEARCH_FN_ATTRS
+#endif
+#define ZSTD_ROW_SEARCH_FN_ATTRS_extDict ZSTD_SEARCH_FN_ATTRS
+#define ZSTD_ROW_SEARCH_FN_ATTRS_dictMatchState ZSTD_SEARCH_FN_ATTRS
+#define ZSTD_ROW_SEARCH_FN_ATTRS_dedicatedDictSearch ZSTD_SEARCH_FN_ATTRS
+
 #define GEN_ZSTD_ROW_SEARCH_FN(dictMode, mls, rowLog)                                          \
-    ZSTD_SEARCH_FN_ATTRS size_t ZSTD_ROW_SEARCH_FN(dictMode, mls, rowLog)(                     \
+    ZSTD_ROW_SEARCH_FN_ATTRS_##dictMode size_t ZSTD_ROW_SEARCH_FN(dictMode, mls, rowLog)(      \
             ZSTD_MatchState_t* ms,                                                             \
             const BYTE* ip, const BYTE* const iLimit,                                          \
             size_t* offsetPtr)                                                                 \