Extra: a graph with parallel branches

Extra: a graph with parallel branches#

The main tutorials run a chain of kernels, one after another. Real steps are often wider: several independent pieces of work read the same input, their results are combined, and the combination is reduced to one number. This extra page builds exactly that, as a graph with a fork, a join and a reduce, and runs it on the GPU and on the CPU. It is optional reading; the main path does not depend on it.

x --+-- branch_a --+
    +-- branch_b --+-- join --> s --> sum(s)
    +-- branch_c --+

Time: ~6 min · Runs on: CPU, GPU if visible (the GPU route is also run when a GPU is present) · You need: What a graph is (Plan.bind, GraphPipeline).

You will build: three branches over one input, a join kernel that combines them, a reduction, and a check that the GPU graph and the CPU run agree.

import sys, pathlib
sys.path.insert(0, str(pathlib.Path.cwd().parent.parent / "_shared"))
from nb_helpers import gpu_available
HAVE_GPU = gpu_available()
print("GPU visible:", HAVE_GPU)

GPU visible: True

import numpy as np

import hawk
from hawk import Mutable, Param, Scalar, Terminated

from eagle import GraphPipeline, deploy

Four small kernels#

Each is an ordinary per-sample function. The three branches read the same x and each writes its own output array, which is what makes them independent: no branch reads anything another branch writes. join reads all three and writes s. deploy compiles them together and returns one plan per kernel, in order.

def branch_a(x: Scalar, k: Param, terminated: Terminated, a: Mutable[Scalar]):
    a = x * x + k


def branch_b(x: Scalar, k: Param, terminated: Terminated, b: Mutable[Scalar]):
    b = 3.0 * x - k


def branch_c(x: Scalar, terminated: Terminated, c: Mutable[Scalar]):
    c = 0.5 * x + 1.0


def join(a: Scalar, b: Scalar, c: Scalar, terminated: Terminated, s: Mutable[Scalar]):
    s = a * b + c


kernels = [hawk.kernel(fn) for fn in (branch_a, branch_b, branch_c, join)]
plan_a, plan_b, plan_c, plan_join = deploy(kernels)

Fork, join, reduce#

Plan.bind packs each kernel’s arguments once, against the arrays it will use. The two routes then differ only in how the bound launches are scheduled:

  • GPU: add_concurrent() registers the three branches as one fork/join stage, and add appends the join after it. The join starts only when every branch has finished. build() captures the whole thing into one CUDA graph; launch() replays it. The device is free to run the branches side by side, but is never required to.

  • CPU: there is no graph to capture; the same bound launches run on the host team, in order.

The reduce is a plain array sum of s on either route, after the graph. A reduction allocates scratch memory, which a captured graph must not do, so it sits just outside the capture rather than inside it.

def bind_all(xp, n=100_000, k=0.5):
    """Allocate the ensemble's arrays and pack each kernel's arguments once."""
    x = xp.asarray(np.linspace(-1.0, 1.0, n))
    a, b, c, s = (xp.zeros(n) for _ in range(4))
    dead = xp.zeros(n, dtype=xp.bool_)  # the `terminated` mask: nobody is finished
    runs = [
        plan_a.bind(x=x, k=k, terminated=dead, a=a),
        plan_b.bind(x=x, k=k, terminated=dead, b=b),
        plan_c.bind(x=x, terminated=dead, c=c),
        plan_join.bind(a=a, b=b, c=c, terminated=dead, s=s),
    ]
    return runs, s


def run_gpu():
    import cupy as cp

    (ra, rb, rc, rj), s = bind_all(cp)
    pipe = GraphPipeline()
    pipe.add_concurrent([ra.launch, rb.launch, rc.launch], names=["a", "b", "c"])  # fork
    pipe.add(rj.launch, name="join")                                                # join
    pipe.build()
    pipe.launch()
    cp.cuda.Device().synchronize()
    return cp.asnumpy(s), float(cp.sum(s))                                          # reduce


def run_cpu():
    (ra, rb, rc, rj), s = bind_all(np)
    for run in (ra, rb, rc, rj):  # same launches, host team, in order
        run.launch()
    return s, float(np.sum(s))

Run both and compare#

s_cpu, total_cpu = run_cpu()
print(f"cpu total = {total_cpu:.9f}")

cpu total = 58332.999996667

if HAVE_GPU:
    s_gpu, total_gpu = run_gpu()
    print(f"gpu total = {total_gpu:.9f}")
    print(f"max |gpu - cpu| over {s_cpu.size} samples = {np.abs(s_gpu - s_cpu).max():.3e}")
else:
    print("no GPU visible; the CPU route above is the whole run")

gpu total = 58332.999996667
max |gpu - cpu| over 100000 samples = 0.000e+00

Both routes agree because they are the same four launches in the same dependency order. The graph only changes how the work is issued: one replay instead of four launches, with the branches free to overlap.

Try this#

  • Add a fourth branch kernel, include it in the add_concurrent list and in join’s arguments; the graph’s shape changes, nothing else does.

  • Add a second add_concurrent stage after the join to fork again. Stages are chained, each joining fully before the next begins.

  • Branches must not write the same array. If you are unsure, eagle.pipeline._concurrent_write_conflict_check can be run on the write targets before build() (it is an internal helper, so treat it as a debugging aid).

Next#