eagle inside your NumPy and CuPy pipeline

eagle inside your NumPy and CuPy pipeline#

You write array code once and run it on NumPy or, with a GPU, on CuPy. This page changes one line of such a pipeline: a step written in array operations becomes a compiled kernel. The arrays you pass decide where it runs, and what you get back is the same kind of array.

It runs on the CPU with NumPy; the CuPy cells run when a CUDA GPU is present.

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.

Hide code cell source

 
# Colab setup, generated by docs/_tools/nb_colab_setup.py (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

1. Your program as it is#

A pipeline over 200,000 ball-flight samples: 50 drag steps, then a summary of the speeds. It is written against whichever array library the data comes from.

import warnings

warnings.filterwarnings("ignore", category=UserWarning, message=r"(?s).*CuPy may not function correctly.*")
import numpy as np

try:
    import cupy as cp

    HAS_GPU = cp.cuda.runtime.getDeviceCount() > 0
except Exception:
    HAS_GPU = False

N, STEPS, DT, G = 200_000, 50, 0.05, 9.81
rng = np.random.default_rng(1)
v0 = 5.0 * rng.normal(size=(3, N))
cd = rng.uniform(0.0, 0.1, N)


def lib(a):
    return cp.get_array_module(a) if HAS_GPU else np


def numpy_step(v, cd):
    xp = lib(v)
    a = -cd * xp.linalg.norm(v, axis=0) * v
    a[2] -= G
    return v + DT * a


def pipeline(step, v, cd):
    for _ in range(STEPS):
        v = step(v, cd)
    speed = lib(v).linalg.norm(v, axis=0)
    return v, float(speed.mean()), float(speed.max())


v_np, mean_np, max_np = pipeline(numpy_step, v0, cd)
print(f"NumPy: mean speed {mean_np:.6f}, max {max_np:.6f}")
if HAS_GPU:
    v_cp, mean_cp, max_cp = pipeline(numpy_step, cp.asarray(v0), cp.asarray(cd))
    print(f"CuPy:  mean speed {mean_cp:.6f}, max {max_cp:.6f}")

NumPy: mean speed 14.302729, max 42.139946
CuPy:  mean speed 14.302729, max 42.139946

2. Drop the kernel in#

The same step as a hawk kernel. eagle.deploy builds it once; plan.run takes NumPy or CuPy arrays and returns the same kind. The one changed line replaces numpy_step.

import hawk
from hawk import Mutable, Param, Scalar, Vector
from hawk.math import norm, vec


@hawk.kernel
def drag_step(v: Vector[3], cd: Scalar, dt: Param, g: Param, v_next: Mutable[Vector[3]]):
    a = vec(0.0, 0.0, -g) - cd * norm(v) * v
    v_next = v + dt * a

import eagle

plan = eagle.deploy(drag_step)


def eagle_step(v, cd):  # the one changed line: the step is now the kernel
    return plan.run(v=v, cd=cd, dt=DT, g=G)


v_np2, mean_np2, max_np2 = pipeline(eagle_step, v0, cd)
print(f"NumPy + eagle: mean speed {mean_np2:.6f}, max {max_np2:.6f}")
np.testing.assert_allclose([mean_np2, max_np2], [mean_np, max_np], rtol=1e-12)
if HAS_GPU:
    v_cp2, mean_cp2, max_cp2 = pipeline(eagle_step, cp.asarray(v0), cp.asarray(cd))
    print(f"CuPy + eagle:  mean speed {mean_cp2:.6f}, max {max_cp2:.6f}  ({type(v_cp2).__module__}.{type(v_cp2).__name__})")
    np.testing.assert_allclose([mean_cp2, max_cp2], [mean_np, max_np], rtol=1e-10)
print("same numbers as the array-operation step")

NumPy + eagle: mean speed 14.302729, max 42.139946
CuPy + eagle:  mean speed 14.302729, max 42.139946  (cupy.ndarray)
same numbers as the array-operation step

3. What stays the same#

Whatever you do with the result, you still do. Slicing, masks and reductions run on the array you get back:

fast = v_np2[:, lib(v_np2).linalg.norm(v_np2, axis=0) > 12.0]
print("samples faster than 12:", fast.shape[1], "| median vertical speed:", float(np.median(fast[2])))
if HAS_GPU:
    fast_gpu = v_cp2[:, cp.linalg.norm(v_cp2, axis=0) > 12.0]
    print("on the GPU, same count:", fast_gpu.shape[1] == fast.shape[1])

samples faster than 12: 121918 | median vertical speed: -15.123036825617957
on the GPU, same count: True

CuPy streams, kernels and the memory pool work next to an eagle launch. Here a CuPy RawKernel runs before and after the eagle launch on one stream, with no synchronization in between. A CuPy launch of the bound kernel (bind) rides CuPy’s current stream:

if HAS_GPU:
    scale = cp.RawKernel(r'''
    extern "C" __global__ void scale(double* x, double k, long long m, long long cycles) {
        long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x;
        long long t0 = clock64();
        while (clock64() - t0 < cycles) {}   // keeps the stream busy, so an unordered launch would be caught
        if (i < m) x[i] *= k;
    }''', "scale")
    vg, cdg = cp.asarray(v0), cp.asarray(cd)
    out = cp.empty_like(vg)
    m = vg.size
    stream = cp.cuda.Stream(non_blocking=True)
    with stream:
        work = vg * 1.0
        scale((m // 256 + 1,), (256,), (work, 2.0, m, 2_000_000))                   # CuPy kernel, on the stream
        plan.bind(v=work, cd=cdg, dt=DT, g=G, v_next=out).launch()                   # eagle launch, same stream
        scale((m // 256 + 1,), (256,), (out, 0.5, m, 0))                             # CuPy kernel again
    stream.synchronize()
    expected = numpy_step(2.0 * v0, cd) * 0.5
    np.testing.assert_allclose(cp.asnumpy(out), expected, rtol=1e-10, atol=1e-12)
    print("RawKernel, eagle, RawKernel on one stream, no sync between: in order")
    pool = cp.get_default_memory_pool()
    before = pool.used_bytes()
    result = plan.run(v=vg, cd=cdg, dt=DT, g=G)
    print("the result came from CuPy's pool:", pool.used_bytes() - before, "bytes =", result.nbytes)
else:
    print("needs a CUDA GPU")

RawKernel, eagle, RawKernel on one stream, no sync between: in order
the result came from CuPy's pool: 4800000 bytes = 4800000

4. What crosses, and who owns it#

NumPy arrays are read where they lie and the call returns when the result is ready. CuPy arrays are read where they lie, on the device, and the call returns at once on CuPy’s current stream. A result you did not supply is a new array of your library. One you supply is written in place and handed back.

one = plan.run(v=v0, cd=cd, dt=DT, g=G)
print("new array:", type(one).__name__, not np.shares_memory(one, v0), "| input untouched:", bool(v0.std() > 0))
mine = np.empty_like(v0)
back = plan.run(v=v0, cd=cd, dt=DT, g=G, v_next=mine)
print("your own buffer, written in place:", back is mine)
if HAS_GPU:
    on_gpu = plan.run(v=cp.asarray(v0), cd=cp.asarray(cd), dt=DT, g=G)
    print("CuPy in, CuPy out, no trip to the host:", type(on_gpu).__module__)

new array: ndarray True | input untouched: True
your own buffer, written in place: True
CuPy in, CuPy out, no trip to the host: cupy

5. What is refused, and why#

When the data cannot be used as given, the call says so and names the argument.

def refused(label, call):
    try:
        call()
    except Exception as err:
        msg = str(err).splitlines()[0]
        print(f"{label}: {type(err).__name__}: {msg if len(msg) <= 150 else msg[:147] + '...'}")


frozen = np.empty_like(v0)
frozen.flags.writeable = False
refused("read-only output", lambda: plan.run(v=v0, cd=cd, dt=DT, g=G, v_next=frozen))
refused("wrong length", lambda: plan.run(v=v0, cd=cd[:5], dt=DT, g=G))
if HAS_GPU:
    refused("host and device mixed", lambda: plan.run(v=v0, cd=cp.asarray(cd), dt=DT, g=G))

read-only output: ValueError: eagle.plan.run: mutable 'v_next' was supplied read-only (array.flags.writeable is False); Plan.run writes the launch's result into every supplied o...
wrong length: ValueError: eagle.plan.run: per_sample 'cd' is declared 1 component(s) wide, so at n=200000 its plane is (200000,); got (5,)
host and device mixed: ValueError: eagle.plan.auto: 'cd' is on the device but 'v' is on the host; the data decides where the kernel runs, and it runs in one place: move one side with...

6. Your CuPy or NumPy code inside an eagle batch#

The other direction: your array code runs as a node of an eagle graph, between kernels, over the same arrays. Here two array operations feed a hawk kernel whose weights you change in place:

import tempfile
import warnings

warnings.filterwarnings("ignore", category=UserWarning, message=r"(?s).*CuPy may not function correctly.*")

import numpy as np

import eagle, eagle.diff, eagle.plan
import hawk
from eagle.graph import opaque
from eagle.interop import import_buffer
from hawk import Mutable, Scalar, Table

gpu = HAS_GPU
xp = cp if gpu else np
@hawk.kernel
def lin(w: Table[Scalar], z: Scalar, y: Mutable[Scalar]):
    y = w.at(0) * z + w.at(1)       # w is ONE buffer, read by every sample

(plan,) = eagle.plan.auto([lin], targets=("cuda" if gpu else "host",), cache_dir=tempfile.mkdtemp())
x, z, y = xp.linspace(-1, 1, 8), xp.zeros(8), xp.zeros(8)
w = xp.array([2.0, 0.5])


def model():                                      # allocation-free: out= keeps it capturable
    xp.multiply(x, 0.7, out=z)
    xp.tanh(z, out=z)

to_np = (lambda a: a.get()) if gpu else np.asarray
g = (eagle.Graph()
     .add_node(opaque(model, reads={"x": x}, writes={"z": z}), name="model")
     .add(plan.bind(w=w, z=z, y=y), name="lin").build(record=True))
w *= 3.0                                          # in place: the next launch reads the new weights
g.launch()
print("a launch after the build sees the new weights:", bool(np.allclose(to_np(y), 6 * np.tanh(0.7 * to_np(x)) + 1.5)))

a launch after the build sees the new weights: True

The planes decide where it runs, and the ladder is short (the three tiers has the picture):

Your code

Write it as

Runs

CuPy with out= (no allocation)

opaque(fn)

captured with the graph

CuPy on a stream of your own

opaque(fn, stream=cp.cuda.Stream())

stream-ordered, the host never waits

NumPy on a GPU graph, or anything that reads back

opaque(fn, eager=True)

host-synchronized

NumPy on a CPU graph

opaque(fn)

a plain call, in order

Weights are one buffer. The Table plane is bound to your array w, so every sample reads the one copy and an in-place change is seen by the next launch, as above.

Gradients. Neither library has autodiff, so you write the vjp: it gets the recorded inputs and the incoming cotangents and returns one gradient per read.

def vjp(primals, cotangents):
    t = xp.tanh(0.7 * xp.asarray(primals["x"]))
    return {"x": xp.asarray(cotangents["z"]) * 0.7 * (1 - t * t)}


grads = eagle.diff.reverse(g.record, {"lin.y": xp.ones(8)}, vjps={"model": vjp})
want = 6 * 0.7 * (1 - np.tanh(0.7 * to_np(x)) ** 2)
print("d y / d x through your model:", bool(np.allclose(to_np(grads["x"]), want)))

d y / d x through your model: True

7. Limits#

What does not work today, one line each:

  • No gradients through NumPy or CuPy, because neither has autodiff (the frameworks' design): write the vjp yourself and pass it as `vjps=`, or use the [torch](torch.ipynb) or [JAX](jax.ipynb) bridge.
  • NumPy calls block until the result is ready, because host arrays have no stream (NumPy's design); CuPy arrays run asynchronously.
  • `plan.run` does not follow a CuPy stream you opened (our scope): bind the arrays and call `bind(...).launch()` as above to stay ordered on it.
  • Inputs of another float type are converted to float64, which is a copy (our scope here): pass float64 arrays to avoid it.
  • One GPU per process until 0.7 (our scope): run one process per GPU, or choose it with `CUDA_VISIBLE_DEVICES`.
  • A NumPy function in a GPU graph is host-synchronized, because host arrays have no stream (NumPy's design): write it as `opaque(fn, eager=True)`, or use CuPy for work that stays on the device.

8. Where next#