How eagle records and replays a graph#
Build the step once, run it as many times as you like. This page follows one small graph through the four stages: record, instantiate, replay, retune.
Time: ~5 min · Runs on: CPU (the same cells capture a CUDA graph when the arrays are cupy arrays) · You need: eagle and NumPy; read first: Build a step as a graph: eagle.Graph
A CUDA graph is a recorded list of kernel launches. Launching the whole list costs one call instead of one call per kernel, which is why eagle records a step once and replays it. The C++ classes are StreamCapturer, Graph and Launcher; eagle.Graph is their Python face.
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.
import numpy as np
import eagle
from eagle.graph import opaque
1. Record#
A node is a callable plus the arrays it reads and writes. Two nodes that only read the same array, l and r below, can run side by side; join waits for both. Nothing runs yet, and print shows what eagle recorded.
x, left, right, out = np.arange(4.0), np.zeros(4), np.zeros(4), np.zeros(4)
calls = {"l": 0, "r": 0, "join": 0}
def run_l():
calls["l"] += 1
left[:] = 2 * x
def run_r():
calls["r"] += 1
right[:] = x + 1
def run_join():
calls["join"] += 1
out[:] = left + right
g = eagle.Graph()
g.add_group([opaque(run_l, reads={"x": x}, writes={"a": left}),
opaque(run_r, reads={"x": x}, writes={"b": right})], names=["l", "r"])
g.add_node(opaque(run_join, reads={"a": left, "b": right}, writes={"out": out}), name="join")
print(g)
eagle.Graph 3 nodes · 1 group · backend: (not built)
stage 0 ┌─────┐ ┌─────┐
│ <l> │ │ <r> │
└──┬──┘ └──┬──┘
a b
└────┬────┘
stage 1 ┌────┴────┐
│ <join> │
└─────────┘
2. Instantiate#
build() turns the recording into something launchable. On a GPU it issues the nodes on a stream between StreamCapturer::begin() and end(), so the launches go into a cudaGraph_t instead of running, and Graph::launcher() instantiates it once. On the CPU it fixes the order and the backend. The picture now names the backend and, for each stage, how the work is shared.
// A stream to record on, and a Graph that will own the captured work.
eagle::cuda::Stream stream;
eagle::cuda::Graph graph;
graph.stream(stream.cuda());
// Everything launched on the stream between begin() and end() is recorded
// into a cudaGraph_t instead of executing eagerly.
eagle::cuda::StreamCapturer capturer(stream.cuda());
capturer.begin();
addOne<<<1, N, 0, stream.cuda()>>>(dBuf, N);
graph.addNode(eagle::cuda::CapturedGraph{ capturer.end() });
g.build()
print(g)
eagle.Graph 3 nodes · 1 group · backend: cpu
stage 0 work_share ┌─────┐ ┌─────┐
│ <l> │ │ <r> │
└──┬──┘ └──┬──┘
a b
└────┬────┘
stage 1 work_share ┌────┴────┐
│ <join> │
└─────────┘
3. Replay#
launch(n) runs the instantiated graph n times. On a GPU each run is one cudaGraphLaunch, with no per-kernel setup; this is the C++ Launcher::launch().
// Instantiate the graph once, then replay it REPLAYS times. Each launch()
// is a single cudaGraphLaunch — no kernel-config marshalling per call.
eagle::cuda::Launcher launcher = graph.launcher();
for (int r = 0; r < REPLAYS; ++r)
launcher.launch();
launcher.synchronize();
The nodes ran once per replay, and the arrays hold the result.
g.launch(3)
print(calls)
print(out)
{'l': 3, 'r': 3, 'join': 3}
[ 1. 4. 7. 10.]
4. Retune#
A recorded graph is not frozen. Launcher::setLogicalSize (Python: set_logical_size) changes how many samples the next replay processes: eagle patches the new launch dimensions into the instantiated graph, with no re-capture. CUDA graphs has the details, and Finished samples: compaction, reorder, and when not to shows compaction shrinking the live set as samples finish. Launch geometry and the choice between host and device are the next page’s subject.