Batch the Work: One Tensor Op Beats a Python Loop

This function applies the same scale-and-bias to every row of a matrix, but it does it with a Python for-loop that touches one row at a time. Each tiny row operation becomes its own GPU kernel launch, so the work fans out into hundreds of launches plus per-iteration Python overhead. Keep the exact same result and rewrite solve() as a vectorized expression over the full tensor. GPUFlight counts the kernel launches inside the judged region, so it rewards collapsing the loop into one batched operation.

Languages: PyTorch

Lesson

PyTorch dispatches computation to backend kernels; views can change metadata without launching a kernel. When you loop in Python and do one small operation per iteration, you pay a launch (and Python dispatch cost) for each one, and the GPU spends more time being fed than computing. The fix is to express the whole computation as operations over the full tensor: broadcasting `x * scale + bias` applies to all rows at once in a couple of launches. The math is identical; only the number of launches changes. GPUFlight measures kernel-launch count inside the harness's judged window, so the batched version passes while the per-row loop does not.

A Python vectorized expression dispatches tensor operations. Multiplication and addition may launch separate kernels, each containing many CUDA blocks. Compilation/fusion is a separate follow-up.

What you will learn

Distinguish views, dispatched operations and observed GPU kernel launches.
See how a Python loop of small ops multiplies launch and dispatch overhead.
Use broadcasting to apply an operation to a whole tensor in one shot.
Read kernel-launch count as a signal that work is properly batched.

Starter code

    # Naive: a separate op per row means one kernel launch per row, plus
    # Python per-iteration overhead. Correct, but it does not scale.
    out = torch.empty_like(x)
    for i in range(x.shape[0]):
        out[i] = x[i] * scale + bias
    return out