Across lanes

Across lanes#

A lane is one work item, and a kernel body is one lane’s code. x is this lane’s value, and nothing in the body reaches another lane unless you say so. This page shows the three ways to say so: reduce over lanes, scatter by key, and read another lane’s value. (A kernel with no Grid has one lane per sample.)

Colab or Kaggle: run this first cell (it installs the packages and downloads the helper files this notebook uses). Elsewhere it only checks the imports.

Show code cell source

Hide code cell source

# Colab setup, generated by python -m raptor_docs.nb colab_setup (do not edit).
import os

try:
    import google.colab  # noqa: F401
    HOSTED = True  # a Colab runtime
except ImportError:
    HOSTED = "KAGGLE_KERNEL_RUN_TYPE" in os.environ  # a Kaggle kernel

if HOSTED:
    import importlib.metadata, pathlib, re, shutil, subprocess, sys, urllib.request

    gpu = bool(shutil.which("nvidia-smi")) and subprocess.run(["nvidia-smi", "-L"], capture_output=True).returncode == 0
    extra = ""
    if gpu:
        # the CUDA major the runtime already has: a CuPy that is installed, else torch's, else the driver's
        major = None
        for m in ("13", "12"):
            try:
                importlib.metadata.version(f"cupy-cuda{m}x")
                major = m
                break
            except importlib.metadata.PackageNotFoundError:
                pass
        if major is None:
            try:
                import torch
                major = "13" if str(torch.version.cuda).startswith("13") else None
            except ImportError:
                pass
        if major is None:
            smi = re.search(r"CUDA Version: (\d+)", subprocess.run(["nvidia-smi"], capture_output=True, text=True).stdout)
            major = "13" if smi and int(smi.group(1)) >= 13 else "12"
        extra = f"[cuda{major}]"
        print(f"GPU found: installing the {extra} extras to match this runtime.")
    pre = ["--pre"] if os.environ.get("RAPTOR_PRE") else []  # optional: test a release candidate
    ref = os.environ.get("RAPTOR_REF", "main")  # optional: fetch the helper files from this git ref
    if pre or "RAPTOR_REF" in os.environ:
        print(f"Pre-release test: installing with --pre, helper files from {ref}.")
    # our packages with --upgrade, so a runtime that kept an older release gets the current one
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "--upgrade", *pre, f"raptor-eagle{extra}", f"raptor-hawk{extra}"])
    if not gpu:
        print("Running on the CPU (no GPU in this runtime). For a GPU: " + ("Settings → Accelerator → GPU T4 x2." if "KAGGLE_KERNEL_RUN_TYPE" in os.environ else "Runtime → Change runtime type → T4 GPU."))

try:
    import eagle, hawk  # noqa: F401
except ImportError as err:
    raise ImportError(
        "eagle, hawk not found. Install with `pip install raptor-eagle raptor-hawk`. For a GPU, add [cuda12] (or [cuda13]) to raptor-eagle and raptor-hawk."
    ) from err

Reduce: `hawk.sum`, `max`, `min`, `mean`#

hawk.sum(x) adds x over every lane. Store the result to a Mutable output and that output is a cell: one value for the whole run. where= keeps only the lanes where a mask is true.

import warnings
warnings.filterwarnings("ignore")

import numpy as np

import eagle
import hawk
import hawk.math as m
from hawk import Grid, Index, Matrix, Mutable, Param, Scalar, Table, Vector


@hawk.kernel
def stats(x: Scalar, total: Mutable[Scalar], mean: Mutable[Scalar], positive: Mutable[Scalar]):
    total = hawk.sum(x)
    mean = hawk.mean(x)
    positive = hawk.sum(x, where=x > 0.0)


@hawk.kernel
def peak(x: Scalar, hi: Mutable[Scalar]):
    hi = hawk.max(x)


stats

<Kernel stats slots=8>

eagle.deploy builds a kernel for the host and the GPU and hands back every cell finished, as a one-element array:

x = np.random.default_rng(1).normal(size=100)
out = eagle.deploy(stats).run(x=x)
print({k: float(v[0]) for k, v in out.items()})
print("numpy:", x.sum(), x.mean(), x[x > 0].sum())
print("max:", float(eagle.deploy(peak).run(x=x)[0]), "numpy:", x.max())

{'mean': -0.0736121212729471, 'positive': 29.154644636394995, 'total': -7.361212127294709}
numpy: -7.361212127294709 -0.07361212127294708 29.154644636395
max: 2.1178387550510482 numpy: 2.1178387550510482

A kernel reduces with one operation: hawk.sum and hawk.mean together (as above), or hawk.max, or hawk.min. Put a hawk.max in its own kernel, as peak does.

Name the lane axes with `Grid`#

Declare the lane space with a Grid parameter and the axes get names. Labels go in numpy order, a bare label has its extent given at run time (the first axis only), and "sample":50 fixes it. The fastest axis is the last one declared: consecutive lanes along it are consecutive threads (or SIMD lanes).

over= names the axis to add up, and the other axes are kept. You bind a runtime extent by its label, channel=4, and the result comes back shaped by what is left.

@hawk.kernel
def per_channel(lane: Grid["channel", "sample":50], x: Scalar) -> Scalar:
    return hawk.sum(x * x, over="sample")


x = np.random.default_rng(1).normal(size=(4, 50))        # (channel, sample): sample is the fastest axis
out = eagle.deploy(per_channel).run(x=x, channel=4)
print("hawk :", np.asarray(out).round(4))
print("numpy:", (x * x).sum(axis=1).round(4))

hawk : [38.8846 34.1721 43.4679 55.4933]
numpy: [38.8846 34.1721 43.4679 55.4933]

over= also takes a position (over=-1 is the fastest axis) or a tuple of labels. A reduced output is shaped (*payload, *kept axes): the value’s own shape first, then the lane axes that remain.

Sum inside the lane, or over an axis#

Two sums, two jobs. v.sum(axis=) folds the value’s own axes (a vector’s components, a matrix’s rows). It stays inside the lane and reaches nobody. hawk.sum(v, over=) folds lanes, and it is the one that combines work from many lanes.

A least-squares calibration of several channels uses both. Evaluating the fitted response adds up a lane’s own S basis terms, so it is a sum inside the lane. The gradient of the misfit with respect to the coefficients adds one term per sample, so it is a sum over the sample axis.

NC, NT, B, S = 3, 4, 64, 2


@hawk.kernel
def response(basis: Vector[S], theta: Table["slot":S]) -> Scalar:
    w = m.vec(theta.at(slot=0), theta.at(slot=1))
    return (basis * w).sum()                                   # inside the lane: this lane's S terms


@hawk.kernel
def coef_grad(lane: Grid["channel", "term":NT, "sample":B],
               feat: Table["slot":S, "term":NT, "sample":B],
               resid: Table["channel":NC, "sample":B]) -> Vector[S]:
    g = resid.at(channel=lane.channel, sample=lane.sample)
    terms = m.vec(feat.at(slot=0, term=lane.term, sample=lane.sample),
                  feat.at(slot=1, term=lane.term, sample=lane.sample)) * g
    return hawk.sum(terms, over="sample")                      # over an axis: B lanes become one


rng = np.random.default_rng(3)
basis, theta = rng.random((S, 10)), rng.random(S)
print("forward :", np.allclose(eagle.deploy(response).run(basis=basis, theta=theta), theta @ basis))

feat, resid = rng.random((S, NT, B)), rng.random((NC, B))
grad = eagle.deploy(coef_grad).run(feat=feat, resid=resid, channel=NC)
print("gradient:", grad.shape, np.allclose(grad, np.einsum("kib,ob->koi", feat, resid)))

forward : True
gradient: (2, 3, 4) True

Rule of thumb: if the numbers you add all live in one lane, write v.sum(axis=). If each lane holds one of the terms, write hawk.sum(v, over=).

A value’s own axes are named "components", "rows" and "columns". hawk.sum of a matrix keeps its shape, and mm.rows[i], mm.columns[j] and v.components[k] read one row, column or component by name. A vector or matrix cell comes back as one array per entry, gram_0 ... gram_3 (row-major).

@hawk.kernel
def frame(mm: Matrix[2, 2], gram: Mutable[Matrix[2, 2]], top: Mutable[Vector[2]], colsum: Mutable[Vector[2]]):
    gram = hawk.sum(mm)                                  # a 2x2 cell
    top = hawk.sum(mm.rows[0])                           # one row of it
    colsum = hawk.sum(mm.sum(axis="rows"))               # the value's own axis first, then the lanes


mm = np.random.default_rng(2).normal(size=(4, 50))       # a matrix plane is (rows*columns, lanes)
out = eagle.deploy(frame).run(mm=mm)
print("gram  :", [round(float(out[f"gram_{k}"][0]), 4) for k in range(4)], " numpy:", mm.sum(1).round(4))
print("top   :", [round(float(out[f"top_{k}"][0]), 4) for k in range(2)], " numpy:", mm[:2].sum(1).round(4))
print("colsum:", [round(float(out[f"colsum_{k}"][0]), 4) for k in range(2)],
      " numpy:", (mm[:2] + mm[2:]).sum(1).round(4))

gram  : [1.0977, -1.9397, 3.7114, -4.4244]  numpy: [ 1.0977 -1.9397  3.7114 -4.4244]
top   : [1.0977, -1.9397]  numpy: [ 1.0977 -1.9397]
colsum: [4.8091, -6.364]  numpy: [ 4.8091 -6.364 ]

Scatter: add or store by key#

hawk.scatter(sink)[key] += v adds each lane’s value into the slot the key names, and = v stores it. The sink is a Mutable declared with named axes, Mutable["bin":8]; a key is an index, a tuple of coordinates, or an Index plane.

Many lanes may share a key. The adds are combined (atomics or a segmented sum), never a race, and the order is not fixed, so a floating-point total can differ in its last digits. A sink you pass keeps its contents and the adds land on top; an omitted one starts at zero.

@hawk.kernel
def hist(x: Scalar, bins: Mutable["bin":8]):
    hawk.scatter(bins)[m.index(x * 8.0, 8)] += 1.0


@hawk.kernel
def put(key: Index, x: Scalar, buf: Mutable["slot":6]):
    hawk.scatter(buf)[key] = x


u = np.random.default_rng(2).uniform(size=1000)
print("omitted :", eagle.deploy(hist).run(x=u))
start = np.full(8, 100.0)
print("passed  :", eagle.deploy(hist).run(x=u, bins=start))
print("numpy   :", np.histogram(u, bins=8, range=(0, 1))[0])
print("store   :", eagle.deploy(put).run(key=np.arange(5, -1, -1, dtype=np.int32), x=np.arange(6.0)))

omitted : eagle.Array(shape=(8,), dtype=float64, device=cpu)
passed  : [211. 246. 220. 226. 224. 224. 213. 236.]
numpy   : [111 146 120 126 124 124 113 136]
store   : eagle.Array(shape=(6,), dtype=float64, device=cpu)

If the keys arrive in order, say so with Sorted[Index] (equal keys sit next to each other). It is a promise that lets hawk add each run of equal keys once instead of once per lane, and the result is the same. Set HAWK_CHECK_SORTED=1 to have a bind verify the promise.

from hawk import Sorted


@hawk.kernel
def by_segment(seg: Sorted[Index], x: Scalar, total: Mutable["segment":4]):
    hawk.scatter(total)[seg] += x


seg = np.sort(np.random.default_rng(1).integers(0, 4, 100)).astype(np.int32)
v = np.random.default_rng(2).normal(size=100)
print("hawk :", np.asarray(eagle.deploy(by_segment).run(seg=seg, x=v)).round(4))
print("numpy:", np.bincount(seg, v).round(4))

hawk : [-0.9941  2.0919  2.2402 -4.1799]
numpy: [-0.9941  2.0919  2.2402 -4.1799]

hawk.sum(terms, over="sample") and hawk.scatter(out)[lane.channel, lane.term] += terms are the same operation: one returns the result, the other adds it into a sink you hand over.

Read another lane: `.at`#

u.at(i) is u as lane i holds it, and on a Grid the axes are named, u.at(row=i - 1, col=j). It is the same .at a Table has. Only a plane that existed before the launch can be read at another lane, so u must be an input, not something this kernel computes. An index out of range is undefined, so guard the edge with an if.

N = 6


@hawk.kernel
def reverse(x: Scalar, i: Index, y: Mutable[Scalar]):
    y = x.at(i) * 2.0


@hawk.kernel
def laplacian(lane: Grid["row", "col":N], u: Scalar, h2inv: Param) -> Scalar:
    i, j = lane.row, lane.col
    if m.land(m.land(i > 0, i < N - 1), m.land(j > 0, j < N - 1)):
        out = (u.at(row=i - 1, col=j) + u.at(row=i + 1, col=j)
               + u.at(row=i, col=j - 1) + u.at(row=i, col=j + 1) - 4.0 * u) * h2inv
    else:
        out = 0.0
    return out


print("reverse  :", eagle.deploy(reverse).run(x=np.arange(6.0), i=np.arange(5, -1, -1, dtype=np.int32)))
u = np.random.default_rng(0).normal(size=(N, N))
lap = eagle.deploy(laplacian).run(u=u, h2inv=1.0, row=N)
want = np.zeros_like(u)
want[1:-1, 1:-1] = u[:-2, 1:-1] + u[2:, 1:-1] + u[1:-1, :-2] + u[1:-1, 2:] - 4.0 * u[1:-1, 1:-1]
print("laplacian:", np.allclose(lap, want))

reverse  : eagle.Array(shape=(6,), dtype=float64, device=cpu)
laplacian: True

The sum contract and `policy=`#

Every sum is within k ULP of the exactly rounded sum; k is on the per-target card. A plain sum’s bound grows with the number of terms added into a slot (at worst that many times eps * sum(|term|)); the card tests it at the run lengths it lists. A compensated sum is tested at k up to 1, however many terms. Both are measured as the error over eps * sum(|term|) against math.fsum. policy= on the operation picks the contract, with the same three words as eagle.reduce:

`policy=`	What you get
`"auto"` (default)	for a sum over lanes, the fast form (atomics or a segmented sum), bound growing with the number of terms; for a scatter, the compensated form
`"compensated"`	the tightest form, k up to 1, a little slower. A scatter keeps its correction in a companion plane `<sink>_c` that eagle makes and owns: you never pass it
`"plain"`	(scatter only) the atomic add alone, no companion, bound growing with the number of terms
`"ordered"`	the same bits every run, every thread count, CPU and GPU (a sum over every lane, over some axes, and `hawk.scatter` of a scalar or vector payload, also in a `for` and several per plane, where a masked, terminated or index-set lane adds nothing, and so does the key `-1` (any other key outside the target is an error, reported by name); a scatter adds each key’s lanes in fixed chunks of lane order, so the same lane count and target size give the same bits, and a different one differs from `np.add.at` only in the last bits)

max and min are exact, so they ignore the policy. hawk.run takes a cell as the output you bind: a (1,) array for a scalar cell, a (columns,) array for a vector cell. It adds the partial sums on the host, and the result is bitwise equal to eagle’s "ordered" policy. Every eagle door (eagle.deploy, a graph, the JAX bridge) hands back the same finished cell. To add up values yourself, use eagle.reduce.

import tempfile
import hawk.artifact

xs = np.random.default_rng(0).normal(size=1000)


@hawk.kernel
def energy(x: Scalar, total: Mutable[Scalar]):
    total = hawk.sum(x * x)


work = tempfile.mkdtemp()
hawk.artifact.build(energy, work, targets=("host",))
cell = np.zeros(1)
hawk.run(hawk.load(work, "energy"), x=xs, total=cell)

ordered = np.zeros(1)
eagle.Graph().add_node(eagle.reduce(xs * xs, ordered, policy="ordered"), name="sum").build().launch()
print("hawk.run cell :", cell[0], " numpy:", (xs * xs).sum())
print("bitwise equal to eagle.reduce(policy='ordered'):", cell[0] == ordered[0])

hawk.run cell : 956.3530648422026  numpy: 956.3530648422026
bitwise equal to eagle.reduce(policy='ordered'): True

Read a total in the same kernel#

Reading a reduced value later in the body is allowed. The total exists only after every lane has been added, so a read of it is a launch boundary: the kernel is built as stages, one launch each, with the total passed between them as a cell. hawk.run and the eagle doors (Plan.run, bind, the reverse, the torch, JAX and Warp bridges) run the stages in sequence. Each stage is compiled and cached on its own, so editing one stage recompiles only that one. A per-lane value a later stage needs is computed again there. A cell that only a later stage reads is marked internal in the sidecar’s cells and is not handed back to you.

import json, pathlib


@hawk.kernel
def share(x: Scalar, z: Mutable[Scalar]):
    z = x / hawk.sum(x)                    # reads the total: two stages


split_dir = tempfile.mkdtemp()
hawk.artifact.build_bundle([share], split_dir, targets=("host",))     # build() publishes one launch
sidecar = json.loads((pathlib.Path(split_dir) / "share.json").read_text())
print("stages:", [(s["kernel"], s["cells_in"], s["cells_out"]) for s in sidecar["stages"]])

xs = np.random.default_rng(0).uniform(1.0, 2.0, 10)
z = np.zeros(10)
hawk.run(hawk.load(split_dir, "share"), x=xs, z=z)
print("z sums to one:", np.isclose(z.sum(), 1.0), " matches numpy:", np.allclose(z, xs / xs.sum()))
try:
    hawk.run(hawk.load(split_dir, "share"), x=xs, z=z, count=5)
except Exception as exc:
    print("refused:", str(exc).split(": ", 1)[-1])

stages: [('share', [], ['cut0']), ('share_2', ['cut0'], [])]
z sums to one: True  matches numpy: True
refused: this kernel reads a value reduced over every sample, so it runs over every sample (no base= or count=)

build_bundle is the door for a split kernel; hawk.artifact.build publishes exactly one launch and says so. until_done and simulate refuse a split kernel, because they step it many times.

Forms that are gone#

hawk.all_samples and the other older spellings are refused, each with the line that replaces it:

import re
import textwrap


def refused(label, build):
    try:
        build()
    except Exception as exc:
        text = re.sub(r"^(line \d+|\S+:\d+): ", "", str(exc))
        print(f"{label}\n" + textwrap.indent(textwrap.fill(text, 96), "    ") + "\n")


def accum():
    hawk.Accum

def contribute():
    @hawk.kernel
    def k(x: Scalar, y: Mutable[Scalar]):
        y.contribute(x)

def add_at():
    @hawk.kernel
    def k(x: Scalar, key: Index, y: Mutable[Scalar]):
        y.add(x, at=key)

def own_sum():
    @hawk.kernel
    def k(x: Scalar, y: Mutable[Scalar]):
        y = x.sum()

def samples_axis():
    @hawk.kernel
    def k(v: Vector[3], y: Mutable[Scalar]):
        y = v.sum(axis="samples")

def masked_store():
    @hawk.kernel
    def k(x: Scalar, y: Mutable[Scalar]):
        y = x
        hawk.all_samples(y)[x > 0.0] = x

def in_a_branch():
    @hawk.kernel
    def k(x: Scalar, y: Mutable[Scalar]):
        if x > 0.0:
            y = hawk.sum(x)


refused("hawk.all_samples", lambda: hawk.all_samples)
refused("Accum", accum)
refused("y.contribute(x)", contribute)
refused("y.add(x, at=key)", add_at)
refused("x.sum() of one lane", own_sum)
refused('v.sum(axis="samples")', samples_axis)
refused("all_samples(y)[mask] = v", masked_store)
refused("hawk.sum inside an if", in_a_branch)

hawk.all_samples
    hawk.all_samples is gone. all_samples(x).sum() is hawk.sum(x) (max/min/mean the same;
    over="<axis>" keeps the other lane axes); all_samples(x)[mask].sum() is hawk.sum(x, where=mask);
    all_samples(m).rows[0].sum() is hawk.sum(m.rows[0]); all_samples(y)[key] += v is
    hawk.scatter(y)[key] += v; all_samples(y)[mask] = v is y = m.where(mask, v, y);
    all_samples(u)[i] is u.at(i)

Accum
    Accum is gone: declare Mutable["bin":n] and write hawk.scatter(bins)[key] += term in the body

y.contribute(x)
    'y': write y = hawk.sum(v) (or hawk.max/min/mean)

y.add(x, at=key)
    'y': write hawk.scatter(y)[key] += v (accumulate by key) or hawk.scatter(y)[key] = v (store by
    key)

x.sum() of one lane
    x.sum(): a rank-0 value is this lane's and has no axis to fold; the sum over the lanes is
    hawk.sum(x)

v.sum(axis="samples")
    x.sum(axis='samples') is gone: the sum over every lane is hawk.sum(x), over some lane axes
    hawk.sum(x, over=...)

all_samples(y)[mask] = v
    hawk.all_samples is gone: a store by key is hawk.scatter(y)[key] += v (or = v), and a masked
    store keeps the old value where the mask is false: y = m.where(mask, v, y)

hawk.sum inside an if
    hawk.sum(...) reduces over the lanes, a collective every lane reaches: it may not sit inside a
    branch or a `for`; reduce at the top level and fold the condition into where=

The rewrite table#

Retired	Write instead
`hawk.all_samples(x).sum()`, `.max()`, `.min()`, `.mean()`	`hawk.sum(x)`, `hawk.max(x)`, `hawk.min(x)`, `hawk.mean(x)`
`hawk.all_samples(x)[mask].sum()`	`hawk.sum(x, where=mask)`
`hawk.all_samples(x).sum(axis="samples")`, `x.sum(axis="samples")`	`hawk.sum(x)`, or `hawk.sum(x, over=...)` for some lane axes
`hawk.all_samples(m).rows[0].sum()` (also `.columns[j]`, `.components[k]`)	`hawk.sum(m.rows[0])`
`hawk.all_samples(u)[i]`	`u.at(i)`
`hawk.all_samples(y)[key] += v`	`hawk.scatter(y)[key] += v`
`hawk.all_samples(y)[key] = v`	`hawk.scatter(y)[key] = v`
`hawk.all_samples(y)[mask] = v`	`y = m.where(mask, v, y)`
`Accum[...]` plane, `y.add(v, at=key)`	`Mutable["bin":n]` and `hawk.scatter(y)[key] += v`
`Reduce[...]` plane, `y.contribute(v)`	`y = hawk.sum(v)` (or `max`, `min`, `mean`)
`y.write(v, at=key)`	`hawk.scatter(y)[key] = v`
`x.sum()` on one lane’s scalar	`hawk.sum(x)`

Limits#

A split kernel runs over every lane: base= and count= on hawk.run are refused. Doors that run single-stage kernels only (until_done, simulate, the step loops) refuse it by name: '<kernel>' runs in N stages; <door> runs single-stage kernels in this release. Use a plain kernel per step there.
A reduction cannot sit under an if or in a for: hoist it out and fold the condition into where=. A scatter cannot sit under an if either: add m.where(cond, v, 0.0) instead. A scatter in a for over range(...) is fine, each pass is its own scatter.
One kernel reduces with one operation (sum or mean, or max, or min). Split the kernel.
max and min over some of the lane axes, and policy="ordered" over some axes or on a scatter, come in the next release. Use hawk.max(x) (one value for the run) or policy="auto" for now.
A scatter adds with += or stores with =; there is no by-key max or min yet.
u.at(i) reads the input u as it was before the launch: not a value computed in this kernel, and not a plane the launch writes.