velt/simd#

import { u8x16, f32x4, mask8x16 } from "velt/simd";

The 128-bit vector types (SIMD vectors) and their methods. A module imports the types it names; the methods work on every vector, also one whose type the module does not name (the mask eq returns). T is a vector type (i8x16 u8x16 i16x8 u16x8 i32x4 u32x4 i64x2 u64x2 f32x4 f64x2), L the type of its lanes, n their number and M its mask type (mask8x16, mask16x8, mask32x4, mask64x2: lanes as wide as the vector's). Each method is one instruction, or a few, on x86-64 (SSE2) and arm64 (NEON) and WebAssembly (simd128), in debug builds too. WebAssembly SIMD is part of Velt's WebAssembly baseline: a program that uses vectors needs an engine that has it (every current browser, Node.js, wasmtime and wasmer do).

A mask m:

The methods keep the types they declare: a lane of a u8x16 is a u8, first() a usize. A position or lane index (at of fits, load, store, loadPartial and storePartial, i of lane and withLane) given as a number is used as itself when it is a whole number from 0 up; any other number (negative, fractional, NaN, ±Infinity) is out of range: fits gives false and the others panic with their RangeError. A negative literal written there (v.lane(-1)) is a compile error instead, since the parameter is a usize. A lane value or shift count given as a number is converted as the standard library converts every number it takes as an integer: truncated, clamped to the type's range, NaN as 0 (u8x16.splat(300) is 255 in every lane).

SHA instructions#

u32x4 has the CPU's SHA-1 and SHA-256 instructions as methods, named after the instruction, with its operands in the instruction's order and lane 0 the low 32 bits: the ARMv8 ones as in ARM's ACLE (vsha256hq_u32(abcd, efgh, wk) is abcd.sha256h(efgh, wk)) and the x86 ones as in Intel's intrinsics (_mm_sha256rnds2_epu32(cdgh, abef, wk) is cdgh.sha256rnds2(abef, wk)), so published code ports line by line.

Every method works everywhere: it is the instruction where the CPU has it, and otherwise an exact emulation in Velt with the same result. The checks are false on another architecture, on WebAssembly, in debug builds (Cranelift has no SHA instructions) and when VELT_SIMD_DISABLE=sha is set (to test the emulation on a CPU that has them). In a loop of them the flag is read once, before the loop, and tested in every iteration (a predictable branch per method); std's hashes run their loops in functions compiled for the extension, where the check is known to be true and there is no branch. With the shuffles above (T.shuffle, reverseBytes) these are what the published SHA-1 / SHA-256 routines use:

import { u32x4 } from "velt/simd";

const abcd = u32x4.splat(0x6a09e667).withLane(1, 0xbb67ae85).withLane(2, 0x3c6ef372).withLane(3, 0xa54ff53a);
const efgh = u32x4.splat(0x510e527f).withLane(1, 0x9b05688c).withLane(2, 0x1f83d9ab).withLane(3, 0x5be0cd19);
const wk = u32x4.splat(0x428a2f98);
// Four SHA-256 rounds: the instruction on arm64 release builds, the emulation elsewhere.
const next = abcd.sha256h(efgh, wk);
console.log(next.lane(0).toString(16)); // ab76ba6d

Carry-less multiplication#

u64x2 has the CPU's carry-less multiplication: x86's PCLMULQDQ and ARM's PMULL, which multiply two 64-bit lanes as polynomials over GF(2) (XOR instead of addition, so no carries) into a 128-bit product, its low 64 bits in lane 0. It is what CRC folding (std's gzip CRC-32, velt/compress) and GHASH are built on.

As with the SHA instructions, every method works everywhere: the instruction where the CPU has it, an exact emulation (a shift and XOR per set bit) elsewhere, on WebAssembly, in debug builds and with VELT_SIMD_DISABLE=clmul. The emulation is far slower than the instruction, so code with a table-driven alternative checks u64x2.hasClmul() first, as velt/compress's CRC-32 does (about 0.45 instructions per byte with the instruction, 3.4 with its tables).

import { u64x2 } from "velt/simd";

// (x + 1)(x + 1) = x^2 + 1 over GF(2): 3 times 3 is 5 without carries.
const p = u64x2.splat(3).clmulLow(u64x2.splat(3));
console.log(p.lane(0), p.lane(1)); // 5 0

Generic code given a vector (xs.includes(v), a Map key, xs.join()) compares and hashes all 128 bits (so NaN equals NaN and 0 differs from -0 there) and prints the type's name, [u8x16]; typeof v is "object". A vector cannot be the value of a module-level const yet.

A module that does not import a name from velt/simd can use it for a type of its own (as TypeScript code that declares its own u8x16 may); one that imports it and declares it too gets the usual "defined multiple times" error. import { u8x16 as Bytes } works as for any import.

import { f32x4, u8x16 } from "velt/simd";

// The dot product, sixteen floats per step in four sums.
function dot(a: f32[], b: f32[]): f32 {
  let s0 = f32x4.splat(0);
  let s1 = f32x4.splat(0);
  let s2 = f32x4.splat(0);
  let s3 = f32x4.splat(0);
  let i: usize = 0;
  while (f32x4.fits(a, i + 12) && f32x4.fits(b, i + 12)) {
    s0 = s0.add(f32x4.load(a, i).mul(f32x4.load(b, i)));
    s1 = s1.add(f32x4.load(a, i + 4).mul(f32x4.load(b, i + 4)));
    s2 = s2.add(f32x4.load(a, i + 8).mul(f32x4.load(b, i + 8)));
    s3 = s3.add(f32x4.load(a, i + 12).mul(f32x4.load(b, i + 12)));
    i += 16;
  }
  let s = s0.add(s1).add(s2.add(s3)).sum();
  while (i < a.length && i < b.length) {
    s += a[i] * b[i];
    i++;
  }
  return s;
}

// ASCII letters to lower case, in place.
function lower(d: u8[]) {
  let i: usize = 0;
  while (u8x16.fits(d, i)) {
    const v = u8x16.load(d, i);
    const upper = v.ge(u8x16.splat(65)).and(v.le(u8x16.splat(90)));
    u8x16.select(upper, v.add(u8x16.splat(32)), v).store(d, i);
    i += 16;
  }
  while (i < d.length) {
    if (d[i] >= 65 && d[i] <= 90) {
      d[i] = d[i] + 32;
    }
    i++;
  }
}

const a: f32[] = [];
const b: f32[] = [];
for (let k = 0; k < 20; k++) {
  a.push(k as f32);
  b.push(2);
}
console.log(dot(a, b)); // 380
const text: u8[] = [72, 101, 76, 76, 79, 32, 87, 79, 82, 76, 68, 33, 32, 65, 66, 67, 68];
lower(text);
console.log(text[0], text[16]); // 104 100

Write the loop as above: a step while a whole vector fits (T.fits(xs, i); with several vectors per step, the last one's position). fits is the check load and store make, so in such a loop the compiler removes theirs; a condition written by hand (i + 16 <= d.length) works too, with the check left in.

The last elements, fewer than a vector, can be done one at a time as in dot above, or as one more vector with loadPartial and storePartial, which is shorter: each reads or writes only the elements that exist, in two overlapping scalar accesses (three single bytes for fewer than four bytes). From about four elements left that is faster than a loop over them (on x86-64, a count and an XOR of a 15-byte array took 0.4 times the instructions); for one to three, the loop is (about half the instructions for one byte):

import { mask8x16, u8x16 } from "velt/simd";

// The number of bytes equal to `b`.
function count(d: u8[], b: u8): usize {
  const needle = u8x16.splat(b);
  let n: usize = 0;
  let i: usize = 0;
  while (u8x16.fits(d, i)) {
    n += u8x16.load(d, i).eq(needle).count();
    i += 16;
  }
  // The lanes past the end are 0, which `b` may be too: only the ones that exist count.
  const rest = mask8x16.prefix((d.length as usize) - i);
  return n + u8x16.loadPartial(d, i).eq(needle).and(rest).count();
}

const d: u8[] = Buffer.alloc(37);
d[3] = 9;
d[36] = 9;
console.log(count(d, 9), count(d, 0)); // 2 35