Разгоняем JS-парсер с помощью WebAssembly (часть 1: базовые возможности)

В прошлой статье, посвященной выяснению победителя в состязании JS-парсеров строки buffers-атрибута узла плана PostgreSQL, мы дошли до факта, что самый эффективный вариант - реализовать примитивный конечный автомат и никогда не трогать регулярные выражения и любые операции над строками сложнее .charCodeAt, если из строки вида 'Buffers: shared hit=123 read=456, local hit=789' мы хотим как можно быстрее получить JSON формата:

{
  "shared-hit"  : 123
, "shared-read" : 456
, "local-hit"   : 789
}

Такой код на тестовом нормализованном наборе показывает время порядка 48ms на 6.3MB или около 130MB/s, что примерно в 11 раз быстрее наивного варианта со .split.

const buffersKeys = ['shared-hit', 'shared-read', 'shared-dirtied', 'shared-written', 'local-hit', 'local-read', 'local-dirtied', 'local-written', 'temp-hit', 'temp-read', 'temp-dirtied', 'temp-written'];

const parseBuffers = line => {
  let rv = {};

  let state;
  let key;

  _word: for (let off = 9 /*'Buffers: '*/, ln = line.length; off < ln; ) {
    switch (line.charCodeAt(off)) {
      case 0x73: // s = shared
        state = 0;
        off += 7;
        continue _word;
      case 0x6c: // l = local
        state = 4;
        off += 6;
        continue _word;
      case 0x74: // t = temp
        state = 8;
        off += 5;
        continue _word;

      case 0x68: // h = hit
        key = state;
        off += 4;
        break;
      case 0x72: // r = read
        key = state + 1;
        off += 5;
        break;
      case 0x64: // d = dirtied
        key = state + 2;
        off += 8;
        break;
      case 0x77: // w = written
        key = state + 3;
        off += 8;
        break;
    }

    let val = 0;
    _digit: for (; off < ln; off++) {
      let ch = line.charCodeAt(off);
      switch (ch) {
        case 0x20: // ' '
          off++;
          break _digit;
        case 0x2c: // ','
          off += 2;
          break _digit;
        default:
          val = val * 10 + ch - 0x30; // '0'
      }
    }

    rv[buffersKeys[key]] = val;
  }
  return rv;
};

Но всегда остается вопрос: "А еще быстрее - можно?"

Чтобы приблизиться к возможностям "железа", но по-прежнему остаться в инфраструктуре JavaScript, сегодня мы научимся решать эту задачу с использованием WebAssembly, постаравшись по пути споткнуться обо все подводные камни.

Готовим трамплин

Поскольку из кода HSM-примера уже понятно, что хардкор мы любим, поэтому продолжим в том же духе - и писать ассемблерный код будем вручную, без использования Emscripten.

Для этого, в минимальном варианте, нам понадобится компилятор WAT-файлов (WebAssembly Text Format):

npm install wabt

И, прямо по примерам, соберем из него свой собственный "компилятор" compile-test.js:

const { readFileSync, writeFileSync } = require('fs');

const fn = 'buffers-test';

require('wabt')().then(wabt => {
  const inputWat = `${fn}.wat`;
  const outputWasm = `${fn}.wasm`;

  const wasmModule = wabt.parseWat(inputWat, readFileSync(inputWat, 'utf8'));
  const { buffer } = wasmModule.toBinary({});

  writeFileSync(outputWasm, Buffer.from(buffer));
});

Заметим, что он генерирует buffers-test.wasm из buffers-test.wat, который пока выглядит у нас примитивной заглушкой:

(module
  (func (result i32)
    (i32.const 42)
  )
  (export "helloWorld" (func 0))
)

Для стартового понимания "что это вообще такое?" хорошо подойдут статья из MDN Описание текстового формата WebAssembly и полный перечень доступных WASM-инструкций.

А вызывать скомпилированное мы будем из buffers-test.js:

const { readFileSync } = require('fs');

const fn = 'buffers-test';

const run = async () => {
  const buffer = readFileSync(`${fn}.wasm`);
  const module = await WebAssembly.compile(buffer);
  
  const instance = await WebAssembly.instantiate(module);
  console.log(instance.exports.helloWorld());
};

run();

Возврат результатов из WASM в JS

В данном примере использован самый простой вариант получения результата работы WASM-кода - возврат значения функции. Но он не единственный - и раз мы хотим достичь наилучшего результата, нам стоит их сравнить.

Результат функции

Изменим наш код так, чтобы при каждом новом вызове функции он возвращал следующее натуральное число миллион раз:

// ...  
  const hw = instance.exports.helloWorld;

  const hrb = process.hrtime();
  // -- 8< --
  for (let i = 1; i < 1000000; i++) {
    let v = hw(); // простое получение результата
  }
  // -- 8< --
  const hre = process.hrtime(hrb);
  const usec = hre[0] * 1e+9 + hre[1];
  console.log(usec);
// ...

(module
  ;; это глобальная переменная модуля, доступ к которой имеют все его функции
  ;; let val = 0 - то есть изменяемое, не-const, целое 32-bit число
  (global $val (mut i32)
    (i32.const 0)
  )

  (func (export "helloWorld") (result i32) ;; сразу экспортируем с нужным именем
    ;; val++ - да, вот настолько сложно
    (global.set $val
      (i32.add
        (global.get $val)
        (i32.const 1)
      )
    )
    ;; return val
    (global.get $val)
  )
)

Для NodeJS 15.9.0, на которой будем проводить все наши тесты, время выполнения этого кода составляет от 10мс - в зависимости от мгновенной загрузки CPU.

Глобальная переменная

Глобальная переменная позволяет как передавать значения из JS в WASM-код, так и получать их обратно как результат какой-то работы.

// ...
  // глобальная WASM-переменная с начальным значением 0
  let val = new WebAssembly.Global({value : 'i32', mutable : true}, 0);
  // этот объект мы используем для передачи информации "внутрь" WASM
  const importObject = {
    js : {
      val
    }
  };
  const instance = await WebAssembly.instantiate(module, importObject);
// ...
  for (let i = 1; i < 1000000; i++) {
    hw();
    let v = val.value; // .value получает реальное значение WebAssembly.Global
  }

Изменение же в ассемблерном коде - минимальное:

(module
  ;; это глобальная переменная, переданная из JS-кода как importObject.js.val
  (global $val (import "js" "val") (mut i32))

  (func (export "helloWorld") ;; функция теперь ничего не возвращает
    (global.set $val
      (i32.add
        (global.get $val)
        (i32.const 1)
      )
    )
  )
)

Этот код делает ровно то же самое... но уже за 48ms. Почти в 5 раз медленнее, однако!

Замечу, что если вместо val.value использовать val.valueOf(), время выполнения возрастает и вовсе до 58ms.

Разделяемая память

Предыдущий результат выглядит совсем грустно, ведь мы хотим и передавать и получать достаточно большое количество значений. Но отчаиваться рано - на помощь придет возможность создать сегмент разделяемой памяти, содержащей общее адресное пространство между JS и WASM:

// ...
  // сегмент разделяемой памяти с начальным размером в 1 страницу
  const memory = new WebAssembly.Memory({initial : 1});
  // проекция этой памяти в качестве массива 32-bit-чисел
  const prj32 = new Uint32Array(memory.buffer);

  const importObject = {
    js : {
      memory
    }
  };
// ...
  for (let i = 1; i < 1000000; i++) {
    hw();
    let v = prj32[0]; // получение значения из проекции
  }
// ...

(module
  ;; это сегмент разделяемой памяти размером в одну 64KB-страницу importObject.js.memory
  (import "js" "memory" (memory 1))

  (func (export "helloWorld")
    ;; memory[0] = memory[0] + 1
    (i32.store
      (i32.const 0) ;; offset
      (i32.add      ;; data
        (i32.load
          (i32.const 0) ;; offset
        )
        (i32.const 1)
      )
    )
  )
)

И результат уже вполне достойный - от 12ms. То есть совсем немного медленнее, чем возврат результата функции, зато вернуть можно сразу несколько значений - надо только согласовать смещения данных между JS и WASM-кодом.

Стоит обратить внимание, что WASM оперирует смещениями памяти "в байтах", а JS - "в ячейках". То есть если нам нужно прочитать prj32[2], то писать i32.store придется по смещению 8 = 2 * 4 ("байтовый" размер элемента проекции можно узнать как prj32.BYTES_PER_ELEMENT).

Заметим, что нам ничто не мешает иметь одновременно несколько "разнотипных" проекций на один и тот же сегмент памяти.

Поскольку мы находимся все-таки не в браузере, а на серверсайде в Node.js-окружении, у нас есть больше возможностей - например, использовать чтение числовых значений непосредственно из Buffer-проекции того же сегмента памяти:

// ...
  const buf = Buffer.from(memory.buffer); // еще одна 1-байтовая проекция
// ...
  for (let i = 1; i < 1000000; i++) {
    hw();
    let v = buf.readUInt32LE(0); // тут смещение тоже побайтовое
  }
// ...

Преимущество этого подхода в возможности читать данные блоками нужной размерности, но по произвольному байтовому смещению, без выравнивания по BYTES_PER_ELEMENT. Увы, за все приходится платить - и результат получается чуть хуже - 16ms.

Вызов JS-функции из WASM

Тут возникает резонный вопрос: если мы имеем потери времени при обращении за значением из JS в WASM - так не эффективнее ли будет это значение передать прямо из WASM-кода в качестве аргумента нужной JS-функции?

// ...
  const func = v => {
    // эта функция будет вызываться из WASM
    // console.log(v)
  };

  const importObject = {
    js : {
      memory
    , func
    }
  };
// ...
  for (let i = 1; i < 1000000; i++) {
    hw();
  }
// ...

(module
  ;; объявляем внутреннее имя и сигнатуру JS-функции
  (import "js" "func" (func $js_func (param i32)))
  (import "js" "memory" (memory 1))

  (func (export "helloWorld")
    ;; memory[0] = memory[0] + 1
    (i32.store
      (i32.const 0)
      (i32.add
        (i32.load
          (i32.const 0)
        )
        (i32.const 1)
      )
    )
    ;; вызов js.func(memory[0])
    (call $js_func
      (i32.load
        (i32.const 0)
      )
    )
  )
)

Такой вариант показывает время еще чуть хуже - порядка 20ms, поскольку все-таки требует операций сохранения/восстановления стека при вызове функции.

Подведем промежуточные итоги:

BigInt vs умножение

Пока что мы протестировали варианты возврата 32bit-значений, но случается, что приходится оперировать и числами, выходящими за рамки этого диапазона. В этом случае у нас есть два варианта: использовать тип BigInt или воспользоваться умножением пары 32bit-значений.

В общем, выбор достаточно очевиден - по возможности стоит всегда пользоваться 32bit-проекцией.

Передача данных из JS в WASM

Воспользуемся возможностью иметь несколько проекций, чтобы сделать доступным в WASM-коде сразу весь файл, зарезервировав первую страницу памяти под обмен дополнительными значениями:

// ...
  const prj32 = new Uint32Array(memory.buffer); // 4-байтовая проекция
  const prj8 = new Uint8Array(memory.buffer);   // 1-байтовая проекция

  const offset = 1 << 16; // размер первой 64KB-страницы

  const data = readFileSync('buffers.txt');
  // _на_ столько страниц нам надо расширить разделяемую память
  memory.grow((data.byteLength >> 16) + 1);

  // заливаем все прочитанные данные, начиная со 2-й страницы
  data.copy(prj8, offset);
// ...

Так даже может какое-то время работать, но недолго и нестабильно, потому что memory.grow порождает уже новый сегмент памяти, поэтому все проекции надо создавать уже после наращивания памяти:

// ...
  memory.grow((data.byteLength >> 16) + 1);
  const prj32 = new Uint32Array(memory.buffer);
  const prj8 = new Uint8Array(memory.buffer);
// ...

В следующей части мы перейдем к написанию реального прикладного кода.

Материалы:

• WABT: The WebAssembly Binary Toolkit

• MDN: Описание текстового формата WebAssembly

• WebAssembly: Instructions

• MDN: WebAssembly.Global

• MDN: WebAssembly.Memory