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.
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())
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
Or hand eagle the graph you captured#
If you already capture your Warp work as a CUDA graph, give eagle the graph instead of a launch function. child
takes the wp.Graph, and one eagle launch then runs the three steps. This needs a GPU: on the CPU a child is
refused when the graph is built, so call your kernel in an opaque, as above.
from eagle.graph import child
if on_gpu:
with wp.ScopedCapture(device=device) as capture:
launch_damp()
z2 = xp.zeros(n)
node = child(capture.graph, reads={"y": y}, writes={"u": u}, name="damp")
g_child = (eagle.Graph().add(square_plan.bind(x=x, y=y)).add_node(node)
.add(shift_plan.bind(u=u, z=z2)).build())
g_child.launch()
print("same answer as the launch-function node:", bool(np.array_equal(to_numpy(z2), to_numpy(z))))
else:
print("a captured graph is GPU work: on the CPU, call your kernel in an opaque, as above")
same answer as the launch-function node: True
if on_gpu:
try:
with wp.ScopedCapture(device=device):
g_child.launch()
except RuntimeError as err:
print(err)
eagle: eagle.Graph.launch was called inside a stream capture (torch.cuda.graph, wp.capture_begin or a cupy graph): an eagle graph replays on its own streams and synchronizes, so the capture would record nothing. Go the other way: give your framework's graph to eagle with eagle.graph.child(graph, reads=..., writes=...) and launch the eagle graph once, outside any capture.
Go this way round. eagle does not launch inside your capture (a launch inside wp.ScopedCapture or torch.cuda.graph is refused), so hand the graph over as above. The same works for a torch.cuda.CUDAGraph(keep_graph=True);
the torch page builds one.
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, andg.schedulereturns it.reverse(..., vjps={"damp": vjp})handed your kernel’s share of the gradient to a function you wrote withwp.Tape.
Next#
Use eagle with NVIDIA Warp has the other direction: an eagle kernel as a
wp.launchreplacement inside your Warp program.The graph model: eagle.Graph explains the three ways your own code can sit in a graph.