c: vectorize table lookups

indutny · indutny · commit 25d43780a332 · 2019-07-08T11:32:12.000-07:00
Use SSE4.2 vector instructions to speed up matching. When node has a lot characters to match and loops to itself, the table lookups could be replaced with a call to `_mm_cmpestri`, skipping all matching characters in the input stream. PR-URL: #24
diff --git a/src/implementation/c/compilation.ts b/src/implementation/c/compilation.ts
@@ -19,6 +19,12 @@ const BLOB_GROUP_SIZE = 11;
 
 type WrappedNode = frontend.IWrap<frontend.node.Node>;
 
+interface IBlob {
+  readonly alignment: number | undefined;
+  readonly buffer: Buffer;
+  readonly name: string;
+}
+
 // TODO(indutny): deduplicate
 export interface ICompilationOptions {
   readonly debug?: string;
@@ -32,7 +38,7 @@ export interface ICompilationProperty {
 
 export class Compilation {
   private readonly stateMap: Map<string, ReadonlyArray<string>> = new Map();
-  private readonly blobs: Map<Buffer, string> = new Map();
+  private readonly blobs: Map<Buffer, IBlob> = new Map();
   private readonly codeMap: Map<string, Code<frontend.code.Code>> = new Map();
   private readonly matchSequence:
       Map<string, MatchSequence> = new Map();
@@ -64,14 +70,20 @@ export class Compilation {
       return;
     }
 
-    for (const [ blob, name ] of this.blobs) {
-      out.push(`static const unsigned char ${name}[] = {`);
+    for (const blob of this.blobs.values()) {
+      const buffer = blob.buffer;
+      let align = '';
+      if (blob.alignment) {
+        align = ` ALIGN(${blob.alignment})`;
+      }
 
-      for (let i = 0; i < blob.length; i += BLOB_GROUP_SIZE) {
-        const limit = Math.min(blob.length, i + BLOB_GROUP_SIZE);
+      out.push(`static const unsigned char${align} ${blob.name}[] = {`);
+
+      for (let i = 0; i < buffer.length; i += BLOB_GROUP_SIZE) {
+        const limit = Math.min(buffer.length, i + BLOB_GROUP_SIZE);
         const hex: string[] = [];
         for (let j = i; j < limit; j++) {
-          const value = blob[j] as number;
+          const value = buffer[j] as number;
 
           const ch = String.fromCharCode(value);
           // `'`, `\`
@@ -84,7 +96,7 @@ export class Compilation {
           }
         }
         let line = '  ' + hex.join(', ');
-        if (limit !== blob.length) {
+        if (limit !== buffer.length) {
           line += ',';
         }
         out.push(line);
@@ -304,12 +316,17 @@ export class Compilation {
     return JSON.stringify(value);
   }
 
-  public blob(value: Buffer): string {
+  public blob(value: Buffer, alignment?: number): string {
     if (this.blobs.has(value)) {
-      return this.blobs.get(value)!;
+      return this.blobs.get(value)!.name;
     }
+
     const res = BLOB_PREFIX + this.blobs.size;
-    this.blobs.set(value, res);
+    this.blobs.set(value, {
+      alignment,
+      buffer: value,
+      name: res,
+    });
     return res;
   }
 }
diff --git a/src/implementation/c/index.ts b/src/implementation/c/index.ts
@@ -35,6 +35,27 @@ export class CCompiler {
     out.push('#include <stdlib.h>');
     out.push('#include <stdint.h>');
     out.push('#include <string.h>');
+    out.push('');
+
+    // NOTE: Inspired by https://github.com/h2o/picohttpparser
+    // TODO(indutny): Windows support for SSE4.2.
+    // See: https://github.com/nodejs/llparse/pull/24#discussion_r299789676
+    // (There is no `__SSE4_2__` define for MSVC)
+    out.push('#ifdef __SSE4_2__');
+    out.push(' #ifdef _MSC_VER');
+    out.push('  #include <nmmintrin.h>');
+    out.push(' #else  /* !_MSC_VER */');
+    out.push('  #include <x86intrin.h>');
+    out.push(' #endif  /* _MSC_VER */');
+    out.push('#endif  /* __SSE4_2__ */');
+    out.push('');
+
+    out.push('#ifdef _MSC_VER');
+    out.push(' #define ALIGN(n) _declspec(align(n))');
+    out.push('#else  /* !_MSC_VER */');
+    out.push(' #define ALIGN(n) __attribute__((aligned(n)))');
+    out.push('#endif  /* _MSC_VER */');
+
     out.push('');
     out.push(`#include "${this.options.header || info.prefix}.h"`);
     out.push(``);
diff --git a/src/implementation/c/node/table-lookup.ts b/src/implementation/c/node/table-lookup.ts
@@ -7,6 +7,11 @@ import { Node } from './base';
 const MAX_CHAR = 0xff;
 const TABLE_GROUP = 16;
 
+// _mm_cmpestri takes 8 ranges
+const SSE_RANGES_LEN = 16;
+const MAX_SSE_CALLS = 2;
+const SSE_ALIGNMENT = 16;
+
 interface ITable {
   readonly name: string;
   readonly declaration: ReadonlyArray<string>;
@@ -24,8 +29,13 @@ export class TableLookup extends Node<frontend.node.TableLookup> {
     this.prologue(out);
 
     const transform = ctx.unwrapTransform(this.ref.transform!);
-    const current = transform.build(ctx, `*${ctx.posArg()}`);
 
+    // Try to vectorize nodes matching characters and looping to themselves
+    // NOTE: `switch` below triggers when there is not enough characters in the
+    // stream for vectorized processing.
+    this.buildSSE(out);
+
+    const current = transform.build(ctx, `*${ctx.posArg()}`);
     out.push(`switch (${table.name}[(uint8_t) ${current}]) {`);
 
     for (const [ index, edge ] of this.ref.edges.entries()) {
@@ -53,6 +63,102 @@ export class TableLookup extends Node<frontend.node.TableLookup> {
     out.push('}');
   }
 
+  private buildSSE(out: string[]): boolean {
+    const ctx = this.compilation;
+
+    // Transformation is not supported atm
+    if (this.ref.transform && this.ref.transform.ref.name !== 'id') {
+      return false;
+    }
+
+    if (this.ref.edges.length !== 1) {
+      return false;
+    }
+
+    const edge = this.ref.edges[0];
+    if (edge.node.ref !== this.ref) {
+      return false;
+    }
+
+    // NOTE: keys are sorted
+    let ranges: number[] = [];
+    let first: number | undefined;
+    let last: number | undefined;
+    for (const key of edge.keys) {
+      if (first === undefined) {
+        first = key;
+      }
+      if (last === undefined) {
+        last = key;
+      }
+
+      if (key - last > 1) {
+        ranges.push(first, last);
+        first = key;
+      }
+      last = key;
+    }
+    if (first !== undefined && last !== undefined) {
+      ranges.push(first, last);
+    }
+
+    if (ranges.length === 0) {
+      return false;
+    }
+
+    // Way too many calls would be required
+    if (ranges.length > MAX_SSE_CALLS * SSE_RANGES_LEN) {
+      return false;
+    }
+
+    out.push('#ifdef __SSE4_2__');
+    out.push(`if (${ctx.endPosArg()} - ${ctx.posArg()} >= 16) {`);
+    out.push('  __m128i ranges;');
+    out.push('  __m128i input;');
+    out.push('  int avail;');
+    out.push('  int match_len;');
+    out.push('');
+    out.push('  /* Load input */');
+    out.push(`  input = _mm_loadu_si128((__m128i const*) ${ctx.posArg()});`);
+    for (let off = 0; off < ranges.length; off += SSE_RANGES_LEN) {
+      const subRanges = ranges.slice(off, off + SSE_RANGES_LEN);
+
+      const blob = ctx.blob(Buffer.from(subRanges), SSE_ALIGNMENT);
+      out.push(`  ranges = _mm_loadu_si128((__m128i const*) ${blob});`);
+      out.push('');
+
+      out.push('  /* Find first character that does not match `ranges` */');
+      out.push(`  match_len = _mm_cmpestri(ranges, ${subRanges.length},`);
+      out.push('      input, 16,');
+      out.push('      _SIDD_UBYTE_OPS | _SIDD_CMP_RANGES |');
+      out.push('        _SIDD_NEGATIVE_POLARITY);');
+      out.push('');
+      out.push('  if (match_len != 0) {');
+      out.push(`    ${ctx.posArg()} += match_len;`);
+
+      const tmp: string[] = [];
+      assert.strictEqual(edge.noAdvance, false);
+      this.tailTo(tmp, {
+        noAdvance: true,
+        node: edge.node,
+      });
+      ctx.indent(out, tmp, '    ');
+
+      out.push('  }');
+    }
+
+    {
+      const tmp: string[] = [];
+      this.tailTo(tmp, this.ref.otherwise!);
+      ctx.indent(out, tmp, '  ');
+    }
+    out.push('}');
+
+    out.push('#endif  /* __SSE4_2__ */');
+
+    return true;
+  }
+
   // TODO(indutny): reduce copy-paste between `C` and `bitcode` implementations
   private buildTable(): ITable {
     const table: number[] = new Array(MAX_CHAR + 1).fill(0);
diff --git a/test/compiler-test.ts b/test/compiler-test.ts
@@ -200,6 +200,19 @@ describe('llparse/Compiler', () => {
       const binary = build(p, start, 'escape-char');
       await binary.check('\\\'', 'off=1\noff=2\n');
     });
+
+    it('should hit SSE4.2 optimization for table-lookup', async () => {
+      const start = p.node('start');
+
+      start
+        .match(ALPHA, start)
+        .skipTo(printOff(p, start));
+
+      // TODO(indutny): validate compilation result?
+      const binary = build(p, start, 'match-bit-check-sse');
+      await binary.check('abcdabcdabcdabcdabcdabcdabcd.abcd.',
+        'off=29\noff=34\n');
+    });
   });
 
   describe('`.peek()`', () => {
diff --git a/test/fixtures/index.ts b/test/fixtures/index.ts
@@ -9,6 +9,7 @@ export { ERROR_PAUSE } from 'llparse-test-fixture';
 const fixtures = new Fixture({
   buildDir: path.join(__dirname, '..', 'tmp'),
   extra: [
+    '-msse4.2',
     '-DLLPARSE__TEST_INIT=llparse__test_init',
     path.join(__dirname, 'extra.c'),
   ],
