A Warp kernel in the middle of an eagle graph

A Warp kernel in the middle of an eagle graph#

Keep the Warp kernel you already wrote. eagle runs it between its own compiled steps, in order, on the same arrays, and the gradient flows through it.

node = opaque(launch_my_warp_kernel, reads={"y": y}, writes={"u": u}, stream=wp.get_stream(), name="damp")
g = eagle.Graph().add(square.bind(x=x, y=y)).add_node(node).add(shift.bind(u=u, z=z)).build(record=True)

The graph below prints itself with your kernel as the dashed middle box. ~8 min · CPU, and the GPU when one is visible · needs eagle, hawk, warp-lang.

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}"])
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "matplotlib", "warp-lang"])
    base = os.environ.get("RAPTOR_RAW_BASE", "https://raw.githubusercontent.com")  # optional: a local mirror
    top = os.environ.get("RAPTOR_WORKDIR") or ("/kaggle/working" if "KAGGLE_KERNEL_RUN_TYPE" in os.environ else "/content")
    root = pathlib.Path(top) / "eagle"
    for rel in (
        "docs/content/_shared/nb_helpers.py",
    ):
        (root / rel).parent.mkdir(parents=True, exist_ok=True)
        urllib.request.urlretrieve(f"{base}/amasat01/eagle/{ref}/{rel}", root / rel)
    (root / "docs/content/userguide/tutorials").mkdir(parents=True, exist_ok=True)
    os.chdir(root / "docs/content/userguide/tutorials")
    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 matplotlib raptor-eagle raptor-hawk warp-lang`. For a GPU, add [cuda12] (or [cuda13]) to raptor-eagle and raptor-hawk."
    ) from err

Hide code cell source

 
import sys, pathlib
sys.path.insert(0, str(pathlib.Path.cwd().parent.parent / "_shared"))
from nb_helpers import to_numpy, DEVICE

import tempfile

import numpy as np
import warp as wp

import eagle
import hawk
from eagle.graph import opaque
from hawk import Mutable, Scalar

wp.config.quiet = True
wp.init()
on_gpu = DEVICE == "gpu"
device = "cuda:0" if on_gpu else "cpu"
if on_gpu:
    import cupy as xp
    device = f"cuda:{xp.cuda.Device().id}"
else:
    xp = np
print("running on", device)

running on cuda:0

Two eagle kernels and one of yours#

The eagle side is two one-line hawk kernels: y = x² before your kernel and z = u + 1 after it. Your kernel is plain Warp: u = sin(y) · exp(-y).

@hawk.kernel
def square(x: Scalar, y: Mutable[Scalar]):
    y = x * x


@hawk.kernel
def shift(u: Scalar, z: Mutable[Scalar]):
    z = u + 1.0


@wp.kernel
def damp(y: wp.array(dtype=wp.float64), u: wp.array(dtype=wp.float64)):
    i = wp.tid()
    u[i] = wp.sin(y[i]) * wp.exp(-y[i])


square_plan, shift_plan = eagle.deploy([square, shift], targets=("cuda",) if on_gpu else ("host",),
                                       cache_dir=tempfile.mkdtemp())

Share the arrays, then declare them#

eagle owns x, y, u, z. Warp gets views of y and u over the same memory (no copy): wp.from_dlpack on the GPU, a copy=False array on the CPU. The node is a function with no arguments that launches your kernel, plus a declaration of what it reads and writes. The declaration is how eagle orders it against its neighbours.

On the GPU, stream= tells eagle which stream your Warp work runs on. eagle then makes that stream wait for the step before and the step after wait for it, so nothing blocks the host.

n = 8
x = xp.linspace(-1.0, 1.0, n)
y, u, z = xp.zeros(n), xp.zeros(n), xp.zeros(n)
view = (lambda a: wp.from_dlpack(a)) if on_gpu else (lambda a: wp.array(a, dtype=wp.float64, device="cpu", copy=False))
wy, wu = view(y), view(u)


def launch_damp():
    wp.launch(damp, dim=n, inputs=[wy], outputs=[wu], device=device)


node = opaque(launch_damp, reads={"y": y}, writes={"u": u},
              stream=wp.get_stream(device) if on_gpu else None, name="damp")
g = (eagle.Graph().add(square_plan.bind(x=x, y=y), name="square").add_node(node)
     .add(shift_plan.bind(u=u, z=z), name="shift").build(record=True))
print(g)

eagle.Graph  3 nodes · 0 groups · backend: gpu · recorded

  stage 0 capture         ┌─────────┐
                          │ square  │
                          └────┬────┘
                               y
                               │
  stage 1 eager_stream ┌╌╌╌╌╌╌╌┴╌╌╌╌╌╌╌┐
                       ╎ damp [stream] ╎
                       └╌╌╌╌╌╌╌┬╌╌╌╌╌╌╌┘
                               u
                               │
  stage 2 capture          ┌───┴───┐
                           │ shift │
                           └───────┘

Run it#

launch() runs all three in order. The check is plain NumPy.

g.launch()
xh = to_numpy(x)
want = np.sin(xh ** 2) * np.exp(-xh ** 2) + 1.0
print("largest difference from NumPy:", float(np.abs(to_numpy(z) - want).max()))
print("how each step ran:", g.schedule)

largest difference from NumPy: 0.0
how each step ran: ('capture', 'eager_stream', 'capture')

On the GPU the middle step is eager_stream: your Warp launch, ordered by streams. The neighbours stay captured in eagle’s CUDA graph. On the CPU every step simply runs in order.

The gradient through your kernel#

eagle can differentiate its own kernels. For yours it asks for one function, vjp, which maps the gradient arriving at u to the gradient for y. Warp’s wp.Tape writes it for you.

def vjp(primals, cots):
    yy = wp.array(to_numpy(primals["y"]), dtype=wp.float64, device=device, requires_grad=True)
    uu = wp.zeros(n, dtype=wp.float64, device=device, requires_grad=True)
    tape = wp.Tape()
    with tape:
        wp.launch(damp, dim=n, inputs=[yy], outputs=[uu], device=device)
    tape.backward(grads={uu: wp.array(to_numpy(cots["u"]), dtype=wp.float64, device=device)})
    return {"y": xp.asarray(yy.grad.numpy())}


grads = eagle.diff.reverse(g.record, {"z": xp.ones(n)}, vjps={"damp": vjp})
yv = xh ** 2
exact = (np.cos(yv) - np.sin(yv)) * np.exp(-yv) * 2 * xh
print("dz/dx, largest difference from the closed form:", float(np.abs(to_numpy(grads["x"]) - exact).max()))

dz/dx, largest difference from the closed form: 5.551115123125783e-17

What just happened#

  • opaque(fn, reads=, writes=, stream=, name=) placed your Warp launch between two eagle steps; the declared arrays were the only contract.

  • str(g) shows the tier of every step, and g.schedule returns it.

  • reverse(..., vjps={"damp": vjp}) handed your kernel’s share of the gradient to a function you wrote with wp.Tape.

Next#