Every form of a graph builder

A tour of what you can write inside @eagle.Graph.builder, one form at a time, each with a line saying what it does. The first-hour version is in the graph model; how it all works is in under the hood. Everything here runs on the CPU; CuPy arrays make the same functions CUDA graphs.

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

import warnings

import numpy as np
import eagle
import hawk
from hawk import Mutable, Param, Scalar, UInt32
from hawk import math as hm

@hawk.kernel
def inc(x: Scalar) -> Scalar:
    return x + 1.0

@hawk.kernel
def dbl(x: Scalar) -> Scalar:
    return 2.0 * x

@hawk.kernel
def add(a: Scalar, b: Scalar) -> Scalar:
    return a + b

@hawk.kernel
def square(v: Scalar) -> Scalar:
    return v * v

@hawk.kernel
def affine(x: Scalar, a: Param, b: Param) -> Scalar:
    return a * x + b

@hawk.kernel
def below(s: Scalar, lim: Param, flag: Mutable[UInt32]):       # writes a one-cell guard: 1 while any sample is below lim
    flag = hm.uint32(hawk.sum(s < lim))

@hawk.kernel
def over(s: Scalar, lim: Param, flag: Mutable[UInt32]):        # writes a one-cell flag: 1 when any sample is above lim
    flag = hm.uint32(hawk.sum(s > lim))

@hawk.kernel
def accumulate(s: Scalar, y: Mutable[Scalar]):                 # writes the caller's array y in place
    y = s + 1.0

def cell(v):                                                   # a one-cell uint32 array: a device condition
    return np.array([v], dtype=np.uint32)

x = np.arange(4.0)

A builder, a block, a block on a hand-made graph

@eagle.Graph.builder
def twice(g: eagle.Graph, v):                  # a BLOCK: a parameter annotated eagle.Graph; it appends to that graph
    return dbl(dbl(v))

@eagle.Graph.builder
def first(x):                                  # a BUILDER: no Graph parameter; it makes and returns one
    g = eagle.Graph()                          # the graph under construction
    g.result = twice(g, inc(x))                # kernel calls become nodes; the block appends two more
    return g                                   # a builder returns its graph

g = first(x)                                   # build: records the graph, runs nothing
print(g())                                     # run: replays it and returns `result`

from eagle.graph import opaque
src, dst = np.arange(4.0), np.zeros(4)
h = eagle.Graph()                              # a hand-made graph ...
h.add_node(opaque(lambda: np.add(src, 1.0, out=dst), reads={"x": src}, writes={"y": dst}), name="shift")
z = twice(h, dst)                              # ... that a block extends, after `shift`
h.build().launch()
print(z)

[ 4.  8. 12. 16.]
[ 4.  8. 12. 16.]

Forks, branches, wait

@eagle.Graph.builder
def forks(x):
    g = eagle.Graph()
    a = dbl(x)
    with g.parallel() as p:                    # a fork: its branches may run at the same time
        with p.branch() as b:                  # branch b: two nodes in order
            l = inc(a)
            l2 = dbl(l)
            with b.parallel() as q:            # a fork inside a branch: open it from the branch handle
                with q.branch() as b1:
                    u = inc(l2)
                with q.branch() as b2:
                    v = dbl(l2)
        with p.branch() as w:                  # branch w ...
            w.wait(l)                          # ... starts after the node that wrote `l`
            r = add(l, a)
    g.result = (u, v, r)                       # leaving a `with` joins: what follows waits for all branches
    return g

print(forks(x)())

(array([ 3.,  7., 11., 15.]), array([ 4., 12., 20., 28.]), array([ 1.,  5.,  9., 13.]))

if, a Python if, a Python for

@eagle.Graph.builder
def choose(x, flag, which, debug):
    g = eagle.Graph()
    if flag:                                   # a one-cell uint32 array: a device condition, read at every replay
        x = dbl(x)
    elif which:                                # elif on another cell
        x = inc(x)
    else:                                      # else
        x = add(x, x)
    if debug:                                  # a Python value: decided once, at build
        x = inc(x)
    for k in (inc, dbl):                       # a Python loop: unrolled at build
        x = k(x)
    g.result = x
    return g

c = cell(1)
g = choose(x, c, cell(1), True)
print(g())
c[0] = 0
print(g())                                     # same graph, the other arm: the cell is read per replay

[ 4.  8. 12. 16.]
[ 6.  8. 10. 12.]

Loops on the device

@eagle.Graph.builder
def capped(x, flag):
    g = eagle.Graph()
    for _ in range(50):                        # capped: at most 50 turns ...
        if not flag: break                     # ... leaves as soon as the cell reads 0
        x = dbl(x)
        below(x, 100.0, flag)                  # the body rewrites the cell
    g.result = x
    return g

@eagle.Graph.builder
def counted(x, n):
    g = eagle.Graph()
    for i in range(n):                         # a counted loop: the cell n is the trip count, read per replay
        x = inc(x)
    g.result = x
    return g

@eagle.Graph.builder
def early_exit(x, flag, done):
    g = eagle.Graph()
    while flag:                                # an uncapped loop: it warns (see below)
        x = inc(x)
        over(x, 2.5, done)
        if done: break                         # mid-body break: skips the rest of the body, leaves the loop
        x = dbl(x)
        below(x, 50.0, flag)
    g.result = x
    return g

@eagle.Graph.builder
def skipping(x, flag, skip):
    g = eagle.Graph()
    while flag:
        x = inc(x)
        below(x, 6.0, flag)
        over(x, 2.5, skip)
        if skip: continue                      # mid-body continue: skips the rest of this turn only
        x = dbl(x)
    g.result = x
    return g

print(capped(np.ones(4), cell(1))())
print(counted(np.zeros(4), cell(3))())
with warnings.catch_warnings(record=True) as w:
    warnings.simplefilter("always")
    print(early_exit(np.zeros(4), cell(1), cell(0))())
print(w[0].message)                            # the warning an uncapped `while` gives
print(skipping(np.zeros(4), cell(1), cell(0))())

[128. 128. 128. 128.]
[3. 3. 3. 3.]
[3. 3. 3. 3.]
uncapped device loop at early_exit:4: a body that never clears its guard never returns; cap it with for _ in range(N): if not cell: break
[6. 6. 6. 6.]

match

@eagle.Graph.builder
def chooser(x, sel):
    g = eagle.Graph()
    match sel:                                 # a one-cell array: one arm runs, read per replay
        case 0:                                # a literal case
            x = dbl(x)
        case 2:
            x = inc(x)
        case _:                                # everything else
            x = square(x)
    g.result = x
    return g

for s in (0, 2, 5):
    print(s, chooser(np.full(4, 3.0), cell(s))()[0])

0 6.0
2 4.0
5 9.0

Buffers and results

@eagle.Graph.builder
def keep(x):
    g = eagle.Graph()
    mine = np.zeros(4)
    a = inc(x, buffer=True)                    # eagle allocates and keeps `a`
    b = dbl(a, buffer=mine)                    # your array is used for b
    g.result = {"a": a, "pair": (b, [square(b)])}     # a result tree: dict, tuple, list, nested
    return g

print(keep(x)())

@eagle.Graph.builder
def void(x, y):                                # no result: the outputs are the Mutable arguments
    g = eagle.Graph()
    accumulate(x, y)                           # writes the caller's y in place
    return g

y = np.zeros(4)
print(void(x, y)(), y)                         # g() is None; y holds the output

{'a': array([1., 2., 3., 4.]), 'pair': (array([2., 4., 6., 8.]), [array([ 4., 16., 36., 64.])])}
None [1. 2. 3. 4.]

Running

g = first(x)
print(g())                                     # one replay; returns result
print(g(10))                                   # ten replays; the result after the last
g.bind(x=x + 10.0)                             # new array, same graph, in place
print(g())
g.launch(2)                                    # two replays and wait (what g(n) does, without the return)
g.enqueue(3)                                   # three replays, without waiting ...
g.synchronize()                                # ... until you ask
print(g.result)                                # the result arrays, read without a launch

[ 4.  8. 12. 16.]
[ 4.  8. 12. 16.]
[44. 48. 52. 56.]
[44. 48. 52. 56.]

Gradients

@eagle.Graph.builder
def model(x, y, record=False):
    g = eagle.Graph(record=record)             # record=True keeps what the reverse needs
    a = affine(x, 2.0, 1.0, buffer=True)       # a top-level buffer: a name a gradient can use
    accumulate(a, y)                           # y: an argument a kernel writes in place
    g.result = square(a)
    return g

y = np.zeros(4)
gr = model(x, y, record=True)
gr()                                           # a recorded graph runs with g() or launch(), not enqueue
ones = np.ones(4)
r = eagle.diff.reverse(gr.record, {"result": ones}, wrt=("x", "a"))   # seed `result`; wrt an argument and a buffer
print(r["x"], r["a"])
r = eagle.diff.reverse(gr.record, {"y": ones}, wrt=("x",))            # seed an argument a kernel writes
print(r["x"])
r = eagle.diff.reverse(gr.record, {"a": ones}, wrt=("x",))            # seed a buffer
print(r["x"])

[ 4. 12. 20. 28.] [ 2.  6. 10. 14.]
[2. 2. 2. 2.]
[2. 2. 2. 2.]

What it refuses

Each refusal says what to change. The limits that belong to recorded loops are listed with the gradients tutorial.

def refused(label, f):
    try:
        f()
    except Exception as e:
        print(f"{label}:\n  {type(e).__name__}: {str(e)[:330]}\n")

def returns_array():
    @eagle.Graph.builder
    def bad(x):
        g = eagle.Graph(); return inc(x)
    bad(x)

def top_level_wait():
    @eagle.Graph.builder
    def bad(x):
        g = eagle.Graph(); g.wait(x); return g
    bad(x)

def overlap():
    @eagle.Graph.builder
    def bad(x):
        g = eagle.Graph(); a = dbl(x)
        with g.parallel() as p:
            with p.branch() as b:
                u = inc(a)
            with p.branch() as c:
                v = add(u, a)                  # reads b's u without waiting
        return g
    bad(x)

def uncapped_under_record():
    @eagle.Graph.builder
    def bad(x, flag):
        g = eagle.Graph(record=True)
        while flag:
            x = dbl(x); below(x, 10.0, flag)
        return g
    bad(np.ones(4), cell(1))

def shape_change():
    g = first(x)
    g.bind(x=np.arange(5.0))

def no_wrt():
    g = model(x, np.zeros(4), record=True); g()
    eagle.diff.reverse(g.record, {"result": np.ones(4)})

def temporary_key():
    g = model(x, np.zeros(4), record=True); g()
    eagle.diff.reverse(g.record, {"result": np.ones(4)}, wrt=("x", "t"))

def enqueue_recorded():
    g = model(x, np.zeros(4), record=True); g.enqueue()

for label, f in (("a builder returns arrays", returns_array), ("wait outside a fork", top_level_wait),
                 ("branches touching one array", overlap), ("an uncapped while under record=True", uncapped_under_record),
                 ("bind with another shape", shape_change), ("reverse without wrt", no_wrt),
                 ("wrt a name that is not a key", temporary_key), ("enqueue on a recorded graph", enqueue_recorded)):
    refused(label, f)

a builder returns arrays:
  TypeError: bad: a builder returns its graph: `return g`; declare the arrays with `g.result = ...`

wait outside a fork:
  RuntimeError: eagle.Graph.wait: only inside a parallel region, on a branch handle

branches touching one array:
  ValueError: eagle.Graph.build: two branches of one parallel region touch one plane and one writes it: put them in one branch, or order them with <branch>.wait(plane) [eagle.Graph.build: 'u' (node 1) and 'v' (node 2) are siblings of one fork/join group but overlap (RAW): 'u' plane 'out' and 'v' plane 'a' share memory; give each branch its ow

an uncapped while under record=True:
  ValueError: bad:4: an uncapped device loop cannot be recorded (record=True): cap it: for _ in range(N): if not cell: break

bind with another shape:
  ValueError: eagle.Graph.bind: 'x' changed shape/dtype: call the builder again

reverse without wrt:
  ValueError: eagle.diff.reverse: a builder graph has no default gradients: pass wrt=(...) naming the arguments it could take: x, y

wrt a name that is not a key:
  KeyError: "eagle.diff.reverse: 't' is not a builder argument, 'result' or a buffer= label: gradients flow through the arrays that enter and leave the graph. For a temporary, give it a dedicated buffer: `t = kernel(..., buffer=True)` (or pass an array), then use that name"

enqueue on a recorded graph:
  RuntimeError: eagle.Graph.enqueue: a record=True graph runs with launch(), which waits for the run its Record describes

One graph with everything

The pieces above, together: a fork with a wait, blocks, a capped device loop, a state seeded from an argument, an output written in place, a bind, a replay and a gradient. Each cell below builds on the one before it.

The kernels. Toy bodies; the roles are hawk’s.

import numpy as np
import hawk
from hawk import Mutable, Param, Scalar, UInt32

@hawk.kernel
def seed(x: Scalar) -> Scalar: return 0.5 * x
@hawk.kernel
def left(a: Scalar) -> Scalar: return a + 1.0
@hawk.kernel
def right(l: Scalar) -> Scalar: return l - 1.0
@hawk.kernel
def polish(v: Scalar) -> Scalar: return 2.0 * v
@hawk.kernel
def tweak(v: Scalar) -> Scalar: return v * v
@hawk.kernel
def merge(s: Scalar, t: Scalar) -> Scalar: return s - t
@hawk.kernel
def remaining(s: Scalar, flag: Mutable[UInt32]):                  # a reduction writes the guard cell
    flag = hawk.math.uint32(hawk.sum(s > 1.0))
@hawk.kernel
def accumulate(s: Scalar, r: Scalar, r2: Scalar, y: Mutable[Scalar]):   # writes the caller's y in place
    y = s + r + r2
@hawk.kernel
def scale(x: Scalar, a: Param, b: Param) -> Scalar: return a * x + b
@hawk.kernel
def square(v: Scalar) -> Scalar: return v * v

x = np.arange(8.0); x2 = x + 1.0; y = np.zeros(8); flag = np.array([1], dtype=np.uint32)

Blocks. A parameter annotated eagle.Graph means “append to the graph you are given”.

import eagle

@eagle.Graph.builder                               # a BLOCK: appends to the open graph it is given,
def twice(g: eagle.Graph, v):                      # returns what it computed, like any function
    return tweak(tweak(v))

@eagle.Graph.builder                               # a block calling a block, with a device loop
def settle(g: eagle.Graph, s, l, flag):
    for _ in range(50):                            # the capped device loop: at most 50 turns
        if not flag: break                         #   the guard: a 1-cell uint32 array
        s = merge(s, twice(g, l))                  #   one buffer, in place; twice: loop-local
        remaining(s, flag)                         #   the body writes the guard through its Mutable
    return s

Build. A builder makes a graph, records into it and returns it. Calling it on arrays builds that graph; nothing runs.

@eagle.Graph.builder                               # a BUILDER: no Graph parameter; makes and returns one
def whole(x, flag, y, record=False) -> eagle.Graph:
    g = eagle.Graph(record=record)                 # the graph under construction; keywords = build's
    a = seed(x)
    with g.parallel() as p:                        # fork; the `with` exit IS the join
        with p.branch() as b:                      # branch b: two nodes
            l = left(a)
            l2 = polish(l)                         #   b goes on
        with p.branch() as w:                      # branch w waits for the END of `left`: the plane is the mark
            w.wait(l)                              #   a sync point; w may now read l
            r = right(l)
        with p.branch() as c:                      # branch c: its own stream; a block inside a branch
            r2 = twice(c, a)
    s = settle(g, x, l2, flag)                     # a block appended after the join; x seeds the state
    accumulate(s, r, r2, y)                        # y: a builder argument written in place -> an output
    g.result = s                                   # the declared result: a hawk return tree
    return g

g = whole(x, flag, y)                              # BUILD + first BIND: records, lowers, instantiates

Bind. New arrays, same program, in place.

g.bind(x=x2)                                       # by the builder's parameter names; no copy, no rebuild;
                                                   # same shape/dtype/device/n (else: "call whole again")

<eagle.Graph 10 nodes, 5 stages, cpu>

Run. Call the graph. It returns the declared result, and the Mutable arguments hold theirs.

s = g()                                            # one replay -> the result tree; y holds its output
s = g(10)                                          # ten replays, the result after the last
g.launch(10); g.enqueue(5); g.synchronize()        # pipeline verbs, frozen (return the Graph)
latest = g.result                                  # the same arrays, read without a launch
print(latest)

[1. 2. 3. 4. 5. 6. 7. 8.]

Gradients. The keys are the names that enter and leave the graph: builder arguments and result.

gr = whole(x, flag, y, record=True)                # a recorded twin (record=True is a build fact)
s = gr()                                           # a recorded graph runs with launch()/g(), never enqueue
grads = eagle.diff.reverse(gr.record, {"result": np.ones(8), "y": np.ones(8)}, wrt=("x",))
print(grads["x"])

[ 2.5   2.75  4.5   9.25 18.5  33.75 56.5  88.25]

The result forms. A first-hour graph, a tree of results, and a graph with no result.

import numpy as np, eagle
@eagle.Graph.builder
def chain(x):
    g = eagle.Graph()
    g.result = square(scale(x, 2.0, 1.0))
    return g
g = chain(np.arange(8.0))                          # build
z = g()                                            # run: z == (2x+1)**2
print(z)

[  1.   9.  25.  49.  81. 121. 169. 225.]

@eagle.Graph.builder
def pair(x):                                       # a tree: exactly the shapes hawk's `return` allows
    g = eagle.Graph(); v = scale(x, 2.0, 1.0)
    g.result = {"v": v, "sq": (square(v), [v])}
    return g
@eagle.Graph.builder
def void(x, y):                                    # no result: the outputs are the Mutable arguments
    g = eagle.Graph(); accumulate(x, x, x, y)
    return g
pg, vg = pair(x), void(x, y)                       # pg() -> the dict; vg() -> None, y written
print(pg(), vg(), y)

{'v': array([ 1.,  3.,  5.,  7.,  9., 11., 13., 15.]), 'sq': (array([  1.,   9.,  25.,  49.,  81., 121., 169., 225.]), [array([ 1.,  3.,  5.,  7.,  9., 11., 13., 15.])])} None [ 0.  3.  6.  9. 12. 15. 18. 21.]