The workload and the ways of running it

The workload and the ways of running it#

A large class of scientific code integrates many independent problems at once: one trajectory, one particle or one parameter sample per GPU thread, each with its own state, its own control flow and its own stopping point. Written for the GPU, this class tends to lose efficiency in four places:

  • launch overhead, when every step is a separate kernel launch from the host;

  • host round trips, when the host has to read a device value every step to decide whether to continue;

  • finished samples, when samples that have stopped keep occupying the batch and keep being computed;

  • idle lanes, when a launch is too small to fill the device.

The benchmark is built to exercise all four. It integrates a batch of damped harmonic oscillators, each with its own frequency, damping ratio and initial state, using a fixed-step classical Runge-Kutta (RK4) scheme. Each sample stops after its own number of steps. In the spread configurations those step counts are drawn log-uniformly over two decades, so most of the batch finishes early; in the uniform configuration every sample runs the full count, a dense, all-active batch. The batch size N runs from 10³ to 10⁶.

The step is written once, as a hawk kernel, one sample at a time:

@hawk.kernel
def oscillator_step(omega: Scalar, zeta: Scalar, nstop: UInt32, dt: Param,
                    terminated: Terminated, x: Mutable[Scalar],
                    v: Mutable[Scalar], k: Mutable[UInt32]):
    # One classical RK4 step of x'' + 2 zeta omega x' + omega^2 x = 0.
    x0 = x
    v0 = v
    w2 = omega * omega
    c = 2.0 * zeta * omega
    h2 = 0.5 * dt
    h6 = dt * 0.16666666666666666
    a1 = -w2 * x0 - c * v0
    x2 = x0 + h2 * v0
    v2 = v0 + h2 * a1
    a2 = -w2 * x2 - c * v2
    x3 = x0 + h2 * v2
    v3 = v0 + h2 * a2
    a3 = -w2 * x3 - c * v3
    x4 = x0 + dt * v3
    v4 = v0 + dt * a3
    a4 = -w2 * x4 - c * v4
    x = x0 + h6 * (v0 + 2.0 * v2 + 2.0 * v3 + v4)
    v = v0 + h6 * (a1 + 2.0 * a2 + 2.0 * a3 + a4)
    k1 = k + 1
    k = k1
    terminated = k1 >= nstop  # the sample stops after its own nstop steps

terminated: Terminated is the per-sample stop mask: a finished sample keeps its state. x, v and k are read as plain variables holding last launch’s value and reassigned to the next one, which is what updates the state in place, launch after launch.

The ways of running it#

All arms do the same arithmetic in the same precision (float64) on the same samples, on the same device; the card’s arm table lists each one’s loop structure, and the card records every arm’s code.

hawk + eagle. The kernel is built with eagle.deploy() and run until every sample has finished with eagle.until_done(): the step, which finishes its own samples, is captured once as the body of eagle.repeat_while(), one CUDA graph whose WHILE node repeats it until the finished count reaches N. One graph launch runs the whole integration; the host does not read anything until it is done.

# eagle_graph_steps: the same one call, on the K-steps kernel; forced
# off the fast path (it is a fixed-K kernel, not "auto", so it would
# never be fast-routed, but the row's device-loop labels/byte model
# assume the band -- forced explicitly so a future kernel change
# cannot silently drift this arm onto the fast path unnoticed)
s = _GpuArm(n, counters="uint32")
s.runner = eagle.until_done(eagle.deploy(oscillator_step), max_steps=max_steps,
                            dt=DT, _fast_mode=False, **s.planes())

The plain @hawk.kernel decorator builds the automatic kernel. By default eagle.until_done() auto-routes it past this device-loop/”band” shape for most batch sizes: a batch of at most a quarter of what one fused wave holds takes one fused launch of the whole run, and one of at least eight times the threads the persistent (work-stealing) launch keeps resident takes the persistent launch – both no-graph, no-map, one host-device sync total; in between, the runner times both entries and keeps the faster. The band shown above (each launch running as many steps per sample as the policy picks, 8 to 64, with the state in registers) is still reachable, and is what this row forces explicitly so its loop-structure description, byte model and kernel count stay meaningful; it is also what any artifact falls back to when a reorder is requested, or when the fast entries are unavailable. The other hawk + eagle arms vary one thing at a time:

  • compaction: the single step traced under a kind whose guard reads an active-set index map (hawk Guard(active_set=True)); every 16 steps, when a sample finished since the last time, an eagle.ActiveSet recomputes the map on the device, inside the same graph, so the launch covers only the samples still running;

  • compaction + reorder: the same loop, plus an occasional physical reorder of every per-sample plane once the live samples spread thin over their warps (eagle.ActiveSet(mask, reorder=theta)), restored to sample order at the end;

  • K fused steps per launch: @hawk.kernel(steps=16), with and without compaction (every 64 steps);

  • eagle auto (eagle picks the launch mode): @hawk.kernel(steps="auto") on the active-set kind – this row is the one left on eagle.until_done’s own default pick, so it takes the fast, no-graph launch described above (fused or persistent, by the size rule) and the band (compacting between launches) only when those entries are unavailable or a reorder is requested; the card records which one each run actually took;

  • hawk + eagle persistent: the plain automatic kernel forced through hawk’s persist entry – one launch, no graph, no map, no policy kernel, a grid sized from the SM count in which each lane takes the next unfinished sample off a counter whenever its own finishes;

  • eagle.simulate: the same automatic kernel run by eagle.simulate() (its bound form, eagle.simulation, so the repetitions reuse one build), given its state and parameters by name – it auto-routes the same way;

  • eager launch loop: the single step (.step, one step per launch) launched from a Python loop with a host read of the finished count after every step: the most direct loop, every intermediate state one host read away.

# eagle_simulate: the model, its state and its parameters by name; it
# runs until every sample has finished
m = _GpuArm(n, counters="uint32")
m.runner = eagle.simulation(oscillator_step, omega=m.omega, zeta=m.zeta,
                            nstop=m.nstop, dt=DT, x=m.x, v=m.v, k=m.k,
                            max_steps=max_steps)

CuPy and PyTorch, masked. The same arithmetic as array expressions over the whole batch, with where keeping finished samples unchanged and a host check that some sample is still running after every step; PyTorch also runs with a block of 16 steps captured as a CUDA graph and replayed. This is how the computation is commonly written in an array framework: no kernel to write, and elementwise kernels that stream memory efficiently. Finished samples are still computed and then discarded.

JAX and Warp. JAX compiles a per-sample lax.while_loop, batched with vmap, into one executable; Warp runs an explicit per-thread kernel, one thread per sample, each stopping at its own step.

CPU OpenMP. The same hawk kernel: its planes on the host select the host side of the same eagle.deploy, run by hawk + eagle’s OpenMP host team on at most 8 threads. It needs no device and no transfers, and it is the correctness reference the other arms are checked against.