31: Recording steps as CUDA graphs

H3101

A step that can’t be recorded as one CUDA graph. A step, a function that changes a model or optimizer state (a mut parameter of either) and does device work itself (value_and_grad, update or changing a field), is recorded on CUDA as one CUDA graph for each shape it’s called with, and replayed. This one, or a function it calls, does more than device work: it loops over a dataset, prints, saves, loads, writes to a log sink, uses Python or calls a Python extern fn, waits for the device with if on a tensor, or makes a size that depends on the data that isn’t kept padded. So it runs operation by operation: correct, with the same results, but slower. It’s a hint, not a warning.

model M:
    param w: [4 3] = normal(std=0.02)

fn loss(m: M, x: [b 4]) -> [] = sum(x @ m.w)

fn step(mut m: M, mut opt: AdamW[M], x: [b 4]):
    l, g = value_and_grad(loss, m, x)
    print(l)
    update(mut m, mut opt, g)

hint[H3101]: `step` isn't recorded as one step: it prints

Return what the step computes, and print or log it in the function that calls the step, as the README’s run writes the logs train_step returns. A loop that does device work itself is moved into a step of its own. See Recorded steps.