Lanes and the machine

This example shows ppy.simd, a few scalars operated on at once, and ppy.cpu, the machine as a facade with no instruction named. Both have a reference implementation under CPython and a lowering to a dialect of the IR.

Run it

python  lanes.ppy
ppy run lanes.ppy

What it prints

python lanes.ppy, ppy run lanes.ppy

10.0
18446744073709551615 0
-6.0 2.0 2.0
28 1.5
# vector width for float: 4 lanes
# avx2 here: True

The simd vocabulary

def dot(a: native.const_ptr[float], b: native.const_ptr[float]) -> float:
    return simd.reduce_add(simd.load[float, 4](a) * simd.load[float, 4](b))


def mix(p: native.ptr[float]) -> float:
    v = simd.load[float, 4](p)
    w = simd.shuffle(v, v, (3, 2, 1, 0))
    chosen = simd.select(v > w, v, w)
    simd.store(chosen / simd.splat[float, 4](2.0), p)
    return simd.extract(chosen, 0) + simd.reduce_add(-w)

• splat, load, store, insert, extract, shuffle, select, and the three reductions are the whole surface.
• + - * /, & | ^, and the comparisons work lane by lane. A comparison gives a Vector[bool, N] for select.
• A floating-point reduce_add folds in lane order, first to last, so the sum is one number everywhere.

bytes_sum masks eight u8 lanes with & 15 and reduces.

Integer lanes wrap

ramp starts three below the largest int and adds a lane-wise v + v. Vector integer arithmetic carries wrap semantics, since there is no unbounded integer in a register. The reference implementation wraps the same way.

The machine as constants

cpu.prefetch(p, locality=2) and cpu.pause() are hints and change no value. cpu.vector_width[float]() and "avx2" in cpu.features() are facts about the machine compiling, folded to constants natively.

The two lines that print them start with #, which is how an example marks output allowed to differ between machines.

Read on: Lanes and the machine · The IR: the simd and cpu dialects

lanes.ppy is hand-written; there is no .py source and no conversion step.

33_simd_and_cpu/lanes.ppy

import ppy
from ppy import cpu, native, simd


def dot(a: native.const_ptr[float], b: native.const_ptr[float]) -> float:
    return simd.reduce_add(simd.load[float, 4](a) * simd.load[float, 4](b))


def ramp(p: native.ptr[int], start: int) -> int:
    v = simd.splat[int, 4](start)
    v = simd.insert(v, 1, start + 1)
    v = simd.insert(v, 2, start + 2)
    v = simd.insert(v, 3, start + 3)
    simd.store(v + v, p)
    return simd.reduce_max(v) - simd.reduce_min(v)


def mix(p: native.ptr[float]) -> float:
    v = simd.load[float, 4](p)
    w = simd.shuffle(v, v, (3, 2, 1, 0))
    chosen = simd.select(v > w, v, w)
    simd.store(chosen / simd.splat[float, 4](2.0), p)
    return simd.extract(chosen, 0) + simd.reduce_add(-w)


def bytes_sum(p: native.const_ptr[ppy.u8]) -> int:
    v = simd.load[ppy.u8, 8](p)
    return simd.reduce_add(v & simd.splat[ppy.u8, 8](15))


def warm(p: native.const_ptr[float], n: int) -> float:
    for i in range(n):
        cpu.prefetch(native.offset(p, i), locality=2)
    cpu.pause()
    return native.load(native.offset(p, n - 1))


def main() -> None:
    a = native.stack_alloc[float](4)
    b = native.stack_alloc[float](4)
    for i in range(4):
        native.store(native.offset(a, i), 1.0 + i)
        native.store(native.offset(b, i), 0.5 * i)
    print(dot(a, b))
    ints = native.stack_alloc[int](4)
    print(ramp(ints, 9223372036854775805), native.load(native.offset(ints, 3)))
    print(mix(a), native.load(native.offset(a, 0)), native.load(native.offset(a, 3)))
    raw = native.stack_alloc[ppy.u8](8)
    for i in range(8):
        native.store(native.offset(raw, i), 240 + i)
    print(bytes_sum(raw), warm(b, 4))
    print(f"# vector width for float: {cpu.vector_width[float]()} lanes")
    print(f"# avx2 here: {'avx2' in cpu.features()}")


main()

Source: examples/33_simd_and_cpu.