One Line, an Order of Magnitude: Why the Speedups Live in Data Movement

Two worlds, one lesson

I ran two small experiments on the same GPU — an NVIDIA L4 — from two lanes that rarely share a page. One is machine learning: a JAX kernel accelerated with @jit. The other is old-school scientific computing: a 1-D heat-diffusion stencil in C, offloaded with an OpenACC directive. Different languages, different toolchains, different decades of pedigree.

Both got roughly an order of magnitude faster from a one-line change. And in both cases the extra speed did not come from doing less arithmetic — the math was byte-for-byte identical before and after. It came from moving less memory.

That is the whole post. The two stories are worth telling because they make the compute-layer lens concrete from opposite directions: a GPU has so much arithmetic throughput that the interesting variable is almost always data movement, not compute.

Experiment 1 — JAX @jit: 33.9× by fusing the memory passes

The JAX demo (projects/08-jax-example) is a small tour of what makes JAX JAX: pure functions plus composable transforms (jit, grad, vmap) lowered to XLA. The piece that matters here is jit on a chain of elementwise operations.

Eager, the chain runs op by op. Each operation reads its input array out of HBM, does one cheap arithmetic step per element, and writes a full-size intermediate array back to HBM — which the next op immediately reads again. For an elementwise pipeline the arithmetic is trivial and the array traffic is everything. The GPU's thousands of ALUs sit mostly idle, starved, waiting on memory.

@jit hands the whole chain to XLA, which fuses it: one compiled kernel that streams each element through the entire pipeline in registers and writes the result once. Same math, but one pass over HBM instead of one pass per op.

backend=gpu
1. jit — XLA fuses the op chain into one compiled program
   eager 46.4 ms   jit 1.37 ms   speedup 33.9x

33.9× on the L4 (the same demo on a laptop CPU shows only ~1.6× — the ALUs there aren't starved enough for fusion to matter much). That gap is the point: the faster and wider the hardware, the more a un-fused memory-bound chain leaves on the floor, and the more fusion buys you. This is exactly why torch.compile exists on the PyTorch side — same idea, same win. XLA just makes it the default way of thinking.

Experiment 2 — OpenACC #pragma acc data: ~40× by staying resident

The directive-offload demo (projects/09-directive-offload) is from the other world entirely: taking an existing C/Fortran HPC code and GPU-accelerating it by annotating loops with #pragma acc — no CUDA rewrite — using NVIDIA's HPC compilers (nvc/nvfortran, NVHPC 26.5). The example is heat1d: a 1-D Jacobi stencil marched for 2000 time steps.

The naive offload puts a #pragma acc parallel loop on the stencil update and calls it done. It's correct, and it's slow — because at each of the 2000 iterations the compiler copies the arrays host→device before the kernel and device→host after. The kernel itself is microseconds; the GPU spends its life waiting on PCIe round trips.

The fix is one line: wrap the entire time loop in a #pragma acc data copy(...) region. Now the data lands on the GPU once, all 2000 iterations run there, and the result comes back once. Identical kernels, identical numerics.

~40× from one line. Without the region, every iteration copies both arrays both ways — 2 arrays × 2 directions × 2000 steps = 8000 transfers. With it, 2. The compiler will even tell you: -Minfo=accel prints every implicit copy it inserted, and NVCOMPILER_ACC_NOTIFY=1 logs each kernel launch and data movement at runtime. Nsight Systems shows the same story as a timeline — a wall of tiny copies collapsing to two.

Why they're the same story

Strip away the languages and both experiments reduce to the same picture of a GPU:

• Compute is cheap and abundant. Thousands of ALUs, enormous FLOP ceilings. Neither experiment was ever limited by how much arithmetic it did.
• Data movement is the tax. HBM bandwidth inside the card; PCIe bandwidth across to the host. Every array you materialize, every buffer you copy, is paid for in the currency the accelerator is actually short on.
• The one-line fixes are both "move the data less." Fusion (JAX/XLA) removes the intra-device HBM round trips between fused ops. A data region (OpenACC/NVHPC) removes the host↔device PCIe round trips between iterations. Same disease, two altitudes on the memory hierarchy.

This is why "just add a GPU" so often disappoints, and why profiling beats guessing. The speedup you're leaving on the table usually isn't a faster kernel — it's a copy you didn't need to make. The engineer's job is to find the data region: the scope across which data can stay put. In XLA the compiler finds it for you when you mark a function @jit; in OpenACC you mark it yourself with #pragma acc data. Either way, the analysis is identical — follow the bytes, not the FLOPs.

Run it yourself

Both are open source and CPU-runnable for the numerics (the directives and jit are simply ignored or trivial on a CPU); a GPU is only needed to reproduce the order-of-magnitude gaps. The harness and companion tutorials ("JAX — how and why, with or without PyTorch" and "directive-based GPU offload") live in the llm-masters repository.

# JAX: eager vs jit (backend prints gpu on a GPU box)
python projects/08-jax-example/jax_demo.py

# OpenACC: the data-region lesson (needs NVHPC + a GPU to offload)
cd projects/09-directive-offload
make verify                                      # serial numeric check, any machine
nvc -acc -gpu=cc89 -Minfo=accel heat1d.c -o heat1d   # cc89 = L4; cc80 A100, cc90 H100

The takeaway compresses to one sentence: on a modern accelerator, an order of magnitude is usually hiding in the memory traffic, not the math — and the fix is often a single line that changes where the data lives, not what the code computes.