SIMD: vector lanes for codecs and constant-time code

SIMD (single instruction, multiple data) means one CPU instruction working on several values at once. A 128-bit SIMD register can hold sixteen bytes, eight 16-bit integers or four 32-bit integers, and one instruction adds all of them pairwise. Each slot is called a lane, and a type such as u8x16 names the lane type and count: sixteen u8 lanes. Video codecs, image filters and ciphers spend most of their time in loops that SIMD speeds up several times. x86 calls its instruction sets SSE and AVX2 (256-bit registers); ARM calls its NEON.

kola has two integer SIMD families that look the same at the call site and follow opposite rules. The ct lanes give up almost every operation so that they can be used inside constant-time code. The media lanes have the data-dependent operations a codec needs and are therefore banned from constant-time code. The checker keeps them from trading places. This chapter covers both families, the parts of their semantics that are pinned down exactly, a small float family, and how the backends lower them. Constant-time code is code whose running time and memory accesses do not depend on secret values, so that timing it reveals nothing; a ct fn is a function the compiler checks for that property (Constant time (ct)).

Two families, one shape

A lane type is a value. It is copied like a fixed-size array and is never a handle. Every lane type is built and read the same four ways: T::splat(x) fills every lane with one scalar, T::of(a, b, ...) takes one scalar per lane, and T::load(slice, off) and T::store(slice, off, v) move a vector in and out of a slice of its element type (&[u8] for u8x16, &[u32] for u32x4, and so on). None of this is new grammar. A lane type is a type and T::op(...) is a call, which is how the checker can hold the two families to different rules without a second syntax.

The ct lanes are u32x4 and its 256-bit, 8-lane twin u32x8. They offer bitwise & | ^, wrapping +% -% *%, and <</>> by a literal shift count, and nothing else: no runtime shift amount, no ==, no trapping arithmetic. Nothing in that set can branch or address memory on a lane's value, so the set is constant-time by construction, which is why these two types are allowed inside a ct fn.

The media lanes are i16x8 (video and DCT work), u8x16 (pixels) and i32x4 (transform accumulators), with 256-bit twins i16x16, u8x32 and i32x8. They provide saturating add and subtract (a result that would overflow sticks at the largest or smallest value instead of wrapping around, which is what you want for pixel brightness), min and max, compares that return a mask (a vector whose lanes are all ones where the comparison held and all zeros elsewhere), select, a byte shuffle, packs, unpacks, widens, the pairwise multiply-add madd, and a runtime uniform shift. Most of these are data-dependent by design (a byte shuffle is a table lookup), so the whole family is non-ct. The media lanes carry no bitwise or shift operators; those belong to the ct lanes.

Here are both shapes side by side: a saturating pixel add on the media lanes and a literal shift on the ct lanes.

fn media_construct(a: &[u8], b: &[u8], out: &mut [u8]) !{} {
    let va = u8x16::splat(a[0]);
    let vb = u8x16::load(b, 0);
    u8x16::store(out, 0, u8x16::sat_add(va, vb));
}

fn ct_construct(a: &[u32], out: &mut [u32]) !{} {
    let v = u32x4::of(a[0], a[1], a[2], a[3]);
    u32x4::store(out, 0, v << 1u32);
}

test both_families_construct_and_store {
    let a = [250u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8,
             0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8];
    let b = [10u8, 0u8, 30u8, 40u8, 0u8, 0u8, 0u8, 0u8,
             0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8];
    let mut mo = [0u8; 16];
    media_construct(&a, &b, &mut mo);
    assert mo[0] == 255u8;   // 250 + 10 saturates at 255
    assert mo[1] == 250u8;   // 250 + 0

    let au = [1u32, 2u32, 3u32, 4u32];
    let mut co = [0u32; 4];
    ct_construct(&au, &mut co);
    assert co[0] == 2u32 && co[3] == 8u32;
}

The narrow ct op set is part of the type everywhere, not only inside a ct fn. A runtime shift amount on u32x4 is refused in an ordinary function too:

fn bad_shift(v: &[u32], n: u32, out: &mut [u32]) !{} {
    let va = u32x4::load(v, 0);
    u32x4::store(out, 0, va << n);
}

error[card §9]: vector shift amount must be a literal (shift-by-immediate only)
  --> bs.kl:3:32
   |
 3 |     u32x4::store(out, 0, va << n);
   |                                ^
   = help: build a variable shift from fixed-amount shifts if needed

Going the other way, a media op inside a ct fn is refused under §6b: the checker reports that the ct fn calls a non-ct function (load, sat_add, store) and names each call.

A constant-time kernel and a codec kernel want opposite things from the same 128 bits of register. One needs an op set that attacker data can never steer; the other needs exactly the data-steered operations (shuffle a row by a motion vector, pick a lane by a comparison). One type for both jobs would force a choice between a ct type that is no longer ct and a codec type without its shuffle. With two types, an edit that reaches for a media op in what used to be constant-time code is stopped at that line.

The slice rule for load and store

T::load and T::store take a slice argument, and a fixed-size local array does not coerce to a slice at the intrinsic itself. Writing u8x16::load(&a, 0) on a local [u8; 16] is refused with u8x16::load needs a &[u8] slice. Pass the array through an ordinary function parameter typed &[u8] or &mut [u8], where the usual array-to-slice coercion applies. Every kernel in this chapter is written that way: a small function over slices, called from a test or main that owns the arrays.

The media operations

The 128-bit media lanes share construction, memory access and wrapping +% -% *%. Beyond that:

The runtime shift takes an unsigned count (usize, u64 and so on); passing an i32 is refused because a negative count would have no meaning. A count of at least the lane width gives 0 for shl and the sign bit splatted for shr. There is deliberately no runtime shift on u8x16 (x86 has no packed-byte shift) and no per-lane variable shift. u8x16::store_low32(slice, off, v) stores only the low four bytes.

Pinned semantics

kola runs programs two ways: an interpreter, which executes the source directly and serves as the reference, and compiled backends that produce machine code. The interpreter defines what a program means, and every backend is compared against it lane for lane. That forces decisions the hardware does not always make for you. Several media operations have a result layout that one instruction set defines one way and another could define differently, and faster. kola pins one answer, because an unspecified layout is permission for two backends to disagree, and they will as soon as a program reads the lane that differs.

The clearest case is u8x16::sad, sum of absolute differences, which is psadbw on x86. A video encoder saves space by describing a block of pixels as "like that block in the previous frame, shifted a little", and to find the best match it compares candidate blocks by summing the absolute differences of their pixels. That makes sad the most executed operation in an encoder. You might expect it to return one number. It returns an i32x4: the sum of the low eight byte differences in lane 0, the sum of the high eight in lane 2, and zero in lanes 1 and 3. kola keeps exactly that layout. Returning a vector rather than a scalar is deliberate too: an encoder accumulates many rows with plain vector adds and reduces to a scalar once at the end, instead of paying a vector-to-scalar move per row. The 256-bit u8x32::sad is two independent 128-bit halves, so group g always lands in lane 2g.

fn sad_row16(a: &[u8], b: &[u8], lanes: &mut [i32]) !{} {
    i32x4::store(lanes, 0, u8x16::sad(u8x16::load(a, 0), u8x16::load(b, 0)));
}

test sad_matches_hand_computed_value {
    let a = [10u8, 10u8, 10u8, 10u8, 10u8, 10u8, 10u8, 10u8,
             20u8, 20u8, 20u8, 20u8, 20u8, 20u8, 20u8, 20u8];
    let b = [7u8, 7u8, 7u8, 7u8, 7u8, 7u8, 7u8, 7u8,
             15u8, 15u8, 15u8, 15u8, 15u8, 15u8, 15u8, 15u8];
    let mut lanes = [0i32; 4];
    sad_row16(&a, &b, &mut lanes);
    assert lanes[0] == 24;   // 8 lanes * |10 - 7|
    assert lanes[1] == 0;    // always zero
    assert lanes[2] == 40;   // 8 lanes * |20 - 15|
    assert lanes[3] == 0;    // always zero
    sad_row16(&a, &a, &mut lanes);
    assert lanes[0] == 0 && lanes[2] == 0;
}

Three more pins cover the rest of that core. abs of the lane minimum (i16 -32768 or i32 -2147483648) wraps to itself rather than trapping or saturating, matching pabsw/pabsd and NEON's vabsq. avg on u8x16 and u8x32 is the rounding average (a + b + 1) >> 1, computed wide so it cannot overflow: 255 and 255 give 255, not 0. That rounding is H.264's bi-prediction average, and a plain (a + b) / 2 rounds down and reconstructs a different picture. hadd on i16x8 and i32x4 sums adjacent pairs with wrapping adds, a's pairs in the low half and b's in the high.

fn abs16(v: &[i16], out: &mut [i16]) !{} {
    i16x8::store(out, 0, i16x8::abs(i16x8::load(v, 0)));
}

fn avg8(a: &[u8], b: &[u8], out: &mut [u8]) !{} {
    u8x16::store(out, 0, u8x16::avg(u8x16::load(a, 0), u8x16::load(b, 0)));
}

test abs_of_int_min_wraps {
    let mut src = [0i16; 8];
    src[0] = -32768i16;
    let mut out = [0i16; 8];
    abs16(&src, &mut out);
    assert out[0] == -32768i16;   // wraps; no trap, no saturation
}

test avg_rounds_up {
    let mut a = [0u8; 16];
    let mut b = [1u8; 16];
    a[15] = 255u8;
    b[15] = 255u8;
    let mut out = [0u8; 16];
    avg8(&a, &b, &mut out);
    assert out[0] == 1u8;      // (0 + 1 + 1) >> 1 rounds up
    assert out[15] == 255u8;   // 255 and 255 do not overflow
}

The 256-bit twins and the per-half rule

u8x32, i16x16 and i32x8 offer the same operation shapes over 32, 16 and 8 lanes, plus i16x16::madd, i32x8::mul_hi, the packs, and T::join(lo, hi), which builds a 256-bit vector from two 128-bit halves of the matching type.

One rule here surprises people who assume that 256-bit means "twice the lanes, same behaviour". unpack_lo, unpack_hi and the saturating packs work per 128-bit half: each half of the 32-byte vector is unpacked or packed independently, as vpunpck* and vpack* do on real hardware, not as one linear 32-byte interleave. The widens are the opposite and are linear: widen_lo_i16 takes source lanes 0 to 15 in order and widen_hi_i16 lanes 16 to 31, like vpmovzxbw. If you port an AVX2 kernel by hand, both rules matter.

fn unpack32(a: &[u8], b: &[u8], lo: &mut [u8], hi: &mut [u8]) !{} {
    let va = u8x32::load(a, 0);
    let vb = u8x32::load(b, 0);
    u8x32::store(lo, 0, u8x32::unpack_lo(va, vb));
    u8x32::store(hi, 0, u8x32::unpack_hi(va, vb));
}

fn widen32(a: &[u8], lo: &mut [i16], hi: &mut [i16]) !{} {
    let v = u8x32::load(a, 0);
    i16x16::store(lo, 0, u8x32::widen_lo_i16(v));
    i16x16::store(hi, 0, u8x32::widen_hi_i16(v));
}

test unpack_is_per_128_bit_half {
    let mut a = [0u8; 32];
    let mut b = [0u8; 32];
    let mut i: usize = 0;
    while i < 32 {
        a[i] = narrow[u8](i).unwrap();
        b[i] = narrow[u8](100 + i).unwrap();
        i = i + 1;
    }
    let mut lo = [0u8; 32];
    let mut hi = [0u8; 32];
    unpack32(&a, &b, &mut lo, &mut hi);
    assert lo[0] == 0u8 && lo[1] == 100u8;
    assert lo[14] == 7u8 && lo[15] == 107u8;
    // the second half restarts at a[16] and b[16]; a linear
    // interleave would have continued from a[8] and b[8]
    assert lo[16] == 16u8 && lo[17] == 116u8;
}

test widen_is_linear {
    let mut a = [0u8; 32];
    let mut i: usize = 0;
    while i < 32 {
        a[i] = narrow[u8](i).unwrap();
        i = i + 1;
    }
    let mut lo = [0i16; 16];
    let mut hi = [0i16; 16];
    widen32(&a, &mut lo, &mut hi);
    assert lo[0] == 0i16 && lo[15] == 15i16;
    assert hi[0] == 16i16 && hi[15] == 31i16;
}

Float lanes

f64x2 and f64x4 hold two and four IEEE binary64 lanes. They are outside ct like scalar f64. They have the same construction and memory intrinsics over &[f64], lane-wise + - * /, a truly fused T::fma(a, b, c) with one rounding, and T::hsum(v), which adds in a fixed order (v[0] + v[1] for two lanes, (v[0] + v[2]) + (v[1] + v[3]) for four) so the result is the same on every backend. As with scalar f64, an operation whose result is NaN traps, and loads and stores carry a NaN that is already held through unchanged.

fn square_plus_one(a: &[f64], out: &mut [f64]) !{} {
    let v = f64x2::load(a, 0);
    f64x2::store(out, 0, f64x2::fma(v, v, f64x2::splat(1.0)));
}

test fused_multiply_add {
    let a = [2.0, 3.0];
    let mut o = [0.0, 0.0];
    square_plus_one(&a, &mut o);
    assert o[0] == 5.0 && o[1] == 10.0;
    assert f64x2::hsum(f64x2::of(1.5, 2.5)) == 4.0;
}

Backends

The interpreter is the reference for every lane type. Every test above passes under kola test, which runs it on the interpreter and through the C backend and compares the two file by file; a disagreement is a compiler bug by definition. Each pinned rule in this chapter was written first as interpreter behaviour and then matched on x86-64 SSE and AArch64 NEON.

• The C backend (kola build --native, kola test --native) lowers every lane type. Each is a plain struct in kola.h and each operation a helper that uses SSE or NEON intrinsics where the target has them and a portable scalar path elsewhere, byte-exact to the interpreter. The 256-bit lanes become real AVX2 vectors only when you build with --cpu x86-64-v3; the resulting binary then requires that CPU level. At the default baseline they use the fallback path.

• The own backend, which needs no C compiler (kola test --own, and kola build --link, which tries it first), also lowers every lane type, running each 256-bit operation as two 128-bit halves on SSE2 or NEON. It does not yet cover every operation: avg, for one, is refused by name at build time, and kola test --native is the fallback.

Lane types may appear in function signatures, as parameters, return values, tuple elements and struct fields; the kola.h struct is the ABI.

Intrinsics and hardware legs

kola has no inline assembly and no FFI. Fast code comes from two places. The first is named intrinsics like the ones above, each with a precise definition the interpreter implements and every backend must match. When a codec kernel is slow, the question to ask first is whether an intrinsic is missing; a SAD loop written in scalar code cannot tell you whether you need assembly or just psadbw.

The second is a hardware leg: a hand-written native implementation of one function, used only on CPUs that support it. A function can carry a clause such as hw x86_64(aes, pclmul) = "symbol" beside its req and ens, naming a C or assembly symbol with the function's ABI that runs only where every listed CPU feature is present. The kola body stays the definition: it is what the interpreter runs, what the checker proves, and what kola test --native compares the leg against; --no-hw drops every leg. std uses this for AES-NI, PCLMUL, SHA-NI, CRC-32C and their AArch64 counterparts. The rules are in SPEC-CARD §8 and Under the hood.