GitHub

This is a tiny compiler (<180LOC) experiment for lowering simple kernel loops into explicit data-parallel code. It takes a small hand-written AST and prints a lowered form with lanes, masks, masked loads, and gathers.

python compiler.py
EXAMPLE: Scale audio volume
SOURCE
kernel scale_audio(samples, out, n, volume):
    for i in range(n):
        # Each sample can be adjusted without looking at its neighbors.
        # That independence lets several samples run side by side.
        out[i] = samples[i] * volume
LOWERED
kernel scale_audio(samples, out, n, volume):
  vector_for base in range(0, n, LANES):
    let i = (base + lane_id)
    let active = (i < n)
    masked_store(out, i, (masked_load(samples, i, active) * volume), active)
EXAMPLE: Move particles
SOURCE
kernel move_particles(position, velocity, out, n, dt):
    for i in range(n):
        # Particle i moves using its own position and velocity.
        # It does not depend on particle i - 1 or particle i + 1.
        out[i] = position[i] + velocity[i] * dt
LOWERED
kernel move_particles(position, velocity, out, n, dt):
  vector_for base in range(0, n, LANES):
    let i = (base + lane_id)
    let active = (i < n)
    masked_store(out, i, (masked_load(position, i, active) + (masked_load(velocity, i, active) * dt)), active)
EXAMPLE: Color by number
SOURCE
kernel color_by_number(color_number, colors, out, n):
    for i in range(n):
        # Each pixel stores a small color number, like a color-by-number page.
        number = color_number[i]
        # Neighboring pixels can name completely different colors.
        # The lowered code must let each lane read its own color entry.
        out[i] = colors[number]
LOWERED
kernel color_by_number(color_number, colors, out, n):
  vector_for base in range(0, n, LANES):
    let i = (base + lane_id)
    let active = (i < n)
    let number = masked_load(color_number, i, active)
    masked_store(out, i, gather(colors, number, active), active)

Read the original on github.com ↗