RapidJSON 程式碼剖析（二）：使用 SSE4.2 優化字串掃描

Milo Yip發表於2015-05-20
APIJSON優化字串

RapidJSON 程式碼剖析（二）：使用 SSE4.2 優化字串掃描

現在的 CPU 都提供了單指令流多資料流（single instruction multiple data, SIMD）指令集。最常見的是用於大量的浮點數計算，但其實也可以用在文書處理方面。

其中，SSE4.2 包含了一些專為字串而設的指令。我們通過使用這些指令，可以大幅提升某些 JSON 解析的效能。

（配圖為 2008 年發售的 Intel Core i7 晶片，它採用的 Nehalem 是第一個支援 SSE4.2 的微架構。）

跳過空白字元

我們知道，有一些 JSON 含有縮排（indentation），這些 JSON 有大量的空白字元（whitespace）。在解析 JSON 的時候，需要跳過這些空白字元。這個操作在 RapidJSON 下是這樣的（reader.h，為配合版面稍改排版）：

template<typename InputStream>
void SkipWhitespace(InputStream& is) {
    internal::StreamLocalCopy<InputStream> copy(is);
    InputStream& s(copy.s);

    while (s.Peek() == ' '  ||
           s.Peek() == '\n' ||
           s.Peek() == '\r' ||
           s.Peek() == '\t')
    {
        s.Take();
    }
}

我們先不關注 StreamLocalCopy 等東西。這段程式碼很簡單，就是凡在輸入流中遇到4種空白字元，都提取出來跳過，直至流裡的字元為非空白字元。

但這種程式碼會帶來很多分支（branching），而且我們每次只能處理一個字元。

SSE4.2

在 Intel 的 SSE4.2 指令集中，有一個 pcmpistrm 指令，它可以一次對一組16個字元與另一組字元作比較，也就是說一個指令可以作最多16×16=256次比較。

對於上面跳過空白字元的需求，我們只需要對16個輸入流裡的字元與4個空白字元比較，即16×4=64次比較。雖然這樣未用盡所有計算能力，但一個指令能代替64個比較以及「或」運算，還是很划算的。

我們可以使用 VC/gcc/clang 都支援的 instrinsic 函式去使用這個指令。這個指令的函式命名為 _mm_cmpistrm()，在nmmintrin.h中定義。

SkipWhitespace 的 SSE4.2 版本只能跳過字串的輸入流，其部分程式碼如下：

inline const char *SkipWhitespace_SIMD(const char* p) {
    // ... 非對齊處理

    static const char whitespace[16] = " \n\r\t";
    const __m128i w = _mm_load_si128((const __m128i *)&whitespace[0]);

    for (;; p += 16) {
        const __m128i s = _mm_load_si128((const __m128i *)p);
        const unsigned r = _mm_cvtsi128_si32(_mm_cmpistrm(w, s, 
            _SIDD_UBYTE_OPS | _SIDD_CMP_EQUAL_ANY |
            _SIDD_BIT_MASK | _SIDD_NEGATIVE_POLARITY));

        if (r != 0) {   // some of characters is non-whitespace
#ifdef _MSC_VER         // Find the index of first non-whitespace
            unsigned long offset;
            _BitScanForward(&offset, r);
            return p + offset;
#else
            return p + __builtin_ffs(r) - 1;
#endif
}

解析一下這裡 _mm_cmpistrm() 用上了的選項：

  • _SIDD_UBYTE_OPS: 操作單位是無號位元組，即16個 unsigned char。
  • _SIDD_CMP_EQUAL_ANY: 每次比較 s 裡的字元，是否和 w 中的任意字元相等。
  • _SIDD_BIT_MASK: 以位元方式返回結果。
  • _SIDD_NEGATIVE_POLARITY: 把結果反轉。這裡指返回值的1代表非空白字元。

然後，我們用_mm_cvtsi128_si32()指令，把返回的最低位32位元組儲存成普通的32位整數。如果含有非空白字元，就使用_BitScanForward()或__builtin_ffs()計算出最早出現的非空白字元，並把指標跳到那裡返回。

對齊問題

通過 SSE 讀寫記憶體，每次可以讀寫128位（16位元組）資料。理想地是使用 128位對齊的地址來讀寫，這樣會最大化讀寫速度。

最初我使用了 _mm_loadu_si128() 從非對齊的來源字串讀取16個字元。當時我覺得最多就是損失一些時間吧，問題似乎不大。但實際上還是出現了問題：

If rapidjson::SkipWhitespace_SIMD(char const*) is called at close to the end of string buffer which has less than 16 bytes of allocated space, the function will read beyond the memory it owns.

In our use case, we parse around 50 million JSON files/buffers per day and
we got hit by the bug around 100 times per day on average before the
workaround.

後來，我估計是因為用非對齊讀取，有可能在邊界會讀到未分配的記憶體分頁，做成很低機率的崩潰。因此，修正方法是先用普通程式碼處理未對齊的地址，然後才使用 SIMD 進行讀取。

inline const char *SkipWhitespace_SIMD(const char* p) {
    // ...

    // 16-byte align to the next boundary
    const char* nextAligned = reinterpret_cast<const char*>(
        (reinterpret_cast<size_t>(p) + 15) & ~15);

    while (p != nextAligned)
        if (*p == ' ' || *p == '\n' || *p == '\r' || *p == '\t')
            ++p;
        else
            return p;

    // The rest of string using SIMD
    // ...
}

快速返回

優化其實還要看實際情況。我們發現，有比較多的情況是，第一個字元已是非空白字元。尤其是已去除空白字元的JSON，上面程式碼的初始時間還是比較大。因此，我們把第一個字元的檢測獨立出來。

inline const char *SkipWhitespace_SIMD(const char* p) {
    // Fast return for single non-whitespace
    if (*p == ' ' || *p == '\n' || *p == '\r' || *p == '\t')
        ++p;
    else
        return p;

    // ...
}

效能測試

測試環境

  • iMac 2.7 GHz Intel Core i5
  • Apple LLVM version 6.1.0 (clang-602.0.49) (based on LLVM 3.6.0svn)

測試用例 1

跳過1M個空白字元1000次。

  • 基本實現: 675 ms
  • SSE4.2: 86 ms
  • strspn: 897 ms

測試用例 2

使用 SAX API 去原位解析（in situ parse）一個含縮排的 671KB sample.json，不處理事件（null handler）。

  • 基本實現: 934 ms
  • SSE4.2: 650 ms

結語

RapidJSON 中使用 SSE4.2 指令集跳過空白字元，可以在一個迭代中進行 64 次字元比較，而且每次讀取 128 位資料應該對記憶體頻寬友好。為了相容更舊的 x86 系 CPU，RapidJSON 也提供了一個 SSE2 的版本，但每個迭代需要執行更多指令，讀取可參考原始碼。

此優化只對含縮排的 JSON 有利，但我們通過「快速返回」使非縮排 JSON 也不會減慢，算是一種權衡之策。在後續的 v1.1 版本中，我希望嘗試利用 SIMD 指令去快速掃瞄需處理轉義（escaping）的字元，不需轉義的部分能使用到 128 位複製至目標緩衝。由於轉義符在 JSON 的出現率較低，此舉應該能進一步提升整體效能。

最後，關於 x86/x64 系的 SIMD 指令，我推薦 Intel Instrinsic Guide 及 Agner Fog 的5本優化手冊。

這兩期都是比較低階的東西，下期將會談一些比較高層一點的，敬請關注。

相關文章