Case Study 2: Optimizing the Boundary

"Wrapping the kernel in Fortran was the easy 90%. Making the wrapper fast was the 90% that was left."

Executive Summary

Where the first case study ported a loop, this one engineers the interface around one. You have a working f2py-driven 2-D heat solver — the Project Checkpoint architecture, Fortran kernel plus Python driver — and it is disappointing: barely faster than pure NumPy, sometimes slower. Nothing is wrong with the Fortran. Everything is wrong with the boundary: how the arrays are laid out, how often they are copied, how often the code crosses from Python into Fortran and back. This study is a staged optimization of that boundary, each stage a single, principled change with a measured payoff, ending with a solver whose speed is set by its arithmetic rather than its plumbing. It is the performance discipline of Part VII applied early, to the one place a two-language program most often leaks its speed.

Skills applied: the two-language workflow (§15.1); f2py intent and !f2py directives (§15.2); the C-vs-F-contiguous copy and intent(inout) (§15.3); building, optimizing, and validating a kernel across changes (§15.5); honest, same-answer-preserving benchmarking.

Background

The starting point is the naive but reasonable driver: an intent(out) kernel returning a fresh field each step, called once per timestep from a Python loop.

! heat_kernel.f90 — the naive step: returns a NEW field each call
module heat_kernel
  use, intrinsic :: iso_fortran_env, only: dp => real64
  implicit none
contains
  subroutine step(u, u_new, alpha, dt, n, m)
    integer,  intent(in)  :: n, m
    real(dp), intent(in)  :: u(n, m)
    real(dp), intent(out) :: u_new(n, m)
    real(dp), intent(in)  :: alpha, dt
!f2py intent(hide), depend(u) :: n = shape(u, 0), m = shape(u, 1)
    integer  :: i, j
    real(dp) :: lap
    u_new = u
    do j = 2, m - 1
      do i = 2, n - 1
        lap = u(i-1,j) + u(i+1,j) + u(i,j-1) + u(i,j+1) - 4.0_dp*u(i,j)
        u_new(i,j) = u(i,j) + alpha*dt*lap
      end do
    end do
  end subroutine step
end module heat_kernel

# the naive driver — correct, but leaking speed at the boundary
u = np.zeros((n, n))                 # <-- default C-contiguous!
u[0, :] = 100.0
for _ in range(nsteps):
    u = heatlib.heat_kernel.step(u, alpha, dt)

It is correct — the field evolves exactly as it should. It is also slow for reasons that have nothing to do with the five-point stencil. We fix them one at a time.

Phase 1 — Establish an Honest Baseline

You cannot optimize what you have not measured, and you cannot trust an optimization that changed the answer. So the harness is fixed for the whole study: run the naive driver, record the time, and snapshot the final field as the reference every later version must reproduce.

import numpy as np, time, heatlib
n, nsteps, alpha, dt = 400, 2000, 1.0, 0.1

u = np.zeros((n, n)); u[0, :] = 100.0
t0 = time.perf_counter()
for _ in range(nsteps):
    u = heatlib.heat_kernel.step(u, alpha, dt)
baseline_time = time.perf_counter() - t0
reference = u.copy()                 # every optimized version must match this

We will not invent a number for baseline_time — you will measure it — but we can reason about its composition, and that reasoning is the whole study. Each step call, with a C-contiguous u, forces f2py to copy the entire n×n field into Fortran order on the way in and copy the n×n result back on the way out. That is $2n^2$ doubles copied per step, $2n^2 \cdot \texttt{nsteps}$ over the run — pure overhead, independent of the arithmetic. For n = 400, nsteps = 2000, that is $2 \cdot 400^2 \cdot 2000 = 6.4 \times 10^8$ doubles, about 5 gigabytes of copying that does no physics.

The diagnosis: the naive version spends much of its life moving bytes between two memory layouts. The stencil is not the bottleneck; the boundary is.

Phase 2 — Fix the Layout (order='F')

The first change is one keyword and it is the largest single win. Create the field F-contiguous so it already matches Fortran's column-major layout, and the per-call copy vanishes — f2py hands Fortran a pointer straight into NumPy's buffer.

u = np.zeros((n, n), dtype=np.float64, order='F')   # F-contiguous from birth
u[0, :] = 100.0
for _ in range(nsteps):
    u = heatlib.heat_kernel.step(u, alpha, dt)       # step returns F-contiguous too
assert np.allclose(u, reference)                     # same answer — nothing broke

The step result is itself F-contiguous, so the next call is also zero-copy; the whole loop crosses the boundary without copying. The $2n^2$-per-step tax from Phase 1 is simply gone. Notice the discipline: the assert np.allclose(u, reference) proves the optimization changed the speed and not the science — the single most important habit in performance work.

⚡ Performance Note. In a two-language numerical program that is "mysteriously only a little faster than Python," a missing order='F' is the first thing to check, before touching the kernel. Layout is cheaper to fix than arithmetic and usually the bigger lever.

Phase 3 — Eliminate the Per-Call Allocation (intent(inout))

With copying gone, the next cost is subtler: the intent(out) kernel allocates a brand-new n×n array on every call to hold u_new, and Python's garbage collector must later reclaim the old one. Two thousand allocate-and-free cycles of a megabyte-plus array is real work. Update the field in place instead.

! the in-place step: no new array; u is modified directly
subroutine step_inplace(u, alpha, dt, n, m)
  integer,  intent(in)    :: n, m
  real(dp), intent(inout) :: u(n, m)          ! in place -> f2py requires F-contiguous
  real(dp), intent(in)    :: alpha, dt
!f2py intent(hide), depend(u) :: n = shape(u, 0), m = shape(u, 1)
  real(dp) :: old(n, m), lap
  integer  :: i, j
  old = u                                      ! one snapshot per call (unavoidable)
  do j = 2, m - 1
    do i = 2, n - 1
      lap = old(i-1,j) + old(i+1,j) + old(i,j-1) + old(i,j+1) - 4.0_dp*old(i,j)
      u(i,j) = old(i,j) + alpha*dt*lap
    end do
  end do
end subroutine step_inplace

u = np.zeros((n, n), dtype=np.float64, order='F')
u[0, :] = 100.0
for _ in range(nsteps):
    heatlib.heat_kernel.step_inplace(u, alpha, dt)   # no return; u changes in place
assert np.allclose(u, reference)

Now Python allocates the field once and Fortran writes into it repeatedly. This is exactly the case where intent(inout)'s strictness (§15.3) is a feature: pass a C-contiguous array here and f2py raises rather than silently copying, which would defeat the very allocation we are trying to save. There is still one old = u snapshot inside the kernel — the stencil genuinely needs the previous field — but that copy is in fast Fortran-to-Fortran memory, not across the language boundary.

Phase 4 — Amortize the Crossing (batch the steps)

The last overhead is the boundary crossing itself: every step_inplace call pays a fixed cost to marshal arguments from Python into Fortran, however small. At 2000 calls it is minor; at millions it is not. The fix is to let one Fortran call do many timesteps, so Python crosses the boundary once per batch rather than once per step.

! advance k timesteps inside Fortran; cross the boundary once per batch
subroutine advance(u, alpha, dt, ksteps, n, m)
  integer,  intent(in)    :: n, m, ksteps
  real(dp), intent(inout) :: u(n, m)
  real(dp), intent(in)    :: alpha, dt
!f2py intent(hide), depend(u) :: n = shape(u, 0), m = shape(u, 1)
  real(dp) :: old(n, m), lap
  integer  :: i, j, t
  do t = 1, ksteps
    old = u
    do j = 2, m - 1
      do i = 2, n - 1
        lap = old(i-1,j) + old(i+1,j) + old(i,j-1) + old(i,j+1) - 4.0_dp*old(i,j)
        u(i,j) = old(i,j) + alpha*dt*lap
      end do
    end do
  end do
end subroutine advance

u = np.zeros((n, n), dtype=np.float64, order='F')
u[0, :] = 100.0
batch = 100
for _ in range(nsteps // batch):
    heatlib.heat_kernel.advance(u, alpha, dt, batch)   # 100 steps per crossing
assert np.allclose(u, reference)

The physics is identical — advance(u, …, k) is just step_inplace done k times — so the reference field still matches. What changed is that Python now crosses into Fortran nsteps/batch times instead of nsteps. The trade-off is granularity: with batching, Python only regains control every batch steps, so if the driver wants to plot or save every 50 steps, the batch size and the output cadence must agree. This is the two-language workflow negotiating with itself — the more work you hand Fortran per call, the less the boundary costs, but the coarser Python's control becomes.

Phase 5 — Know When to Stop

Stack the changes and reason about where the time now goes, relative to the naive baseline:

(The column is deliberately qualitative — the real factors are yours to measure, and depend on n, nsteps, and your machine. The ordering of the wins is the robust lesson.)

After these three changes the program spends its time in the stencil arithmetic, which is exactly where a fast program should spend it — and where Part VII goes next (loop order for column-major access, vectorization, then parallelism). That is the signal to stop tuning the boundary: once the boundary is no longer the bottleneck, further boundary cleverness buys nothing, and the next win is inside the kernel or across many cores. Optimizing past the bottleneck is how you trade readability for a speedup you cannot measure.

Discussion Questions

1. Rank the three optimizations (layout, in-place, batching) by expected payoff for a small grid (n = 32, nsteps = 10^6) and again for a large grid (n = 4000, nsteps = 100). Do the rankings differ, and why? (Think about which cost scales with $n^2$ and which with nsteps.)
2. Batching coarsens Python's control over the run. Design an interface that keeps most of batching's speed and lets the driver save a frame every 50 steps. What does advance need to return or accept?
3. Phase 3's intent(inout) makes a C-contiguous array a hard error instead of a silent copy. Argue that this stricter behavior is safer for a production solver, not merely faster.

Your Turn: Extensions

• Option A. Implement all four versions, wrap them in one module, and produce the real version of the Phase 5 table for your machine at two grid sizes. Does the ordering of wins hold?
• Option B. Add a max_change output to advance (the largest single-cell update over the batch) so the Python driver can stop early when the field has reached steady state. Which intent, and does it disturb the zero-copy property of u?
• Option C. Take the fully optimized kernel into Chapter 33 territory: sketch where an OpenMP directive would go on the stencil loop, and argue why the boundary optimizations here are a prerequisite for parallel speedup to be visible (if the boundary dominates, the parallel kernel's gains are hidden).

Key Takeaways

• In a two-language program, the boundary — layout, copies, allocations, crossings — is where speed most often leaks, and none of those costs live in the kernel you were tempted to tune first.
• Fix them in payoff order: layout (order='F') kills the per-call copy, intent(inout) kills the per-call allocation, batching kills the per-crossing overhead. Each is one principled change.
• Every optimization must carry an assert np.allclose(..., reference): a faster program that computes a different answer is not an optimization, it is a regression with good marketing.
• Stop when the boundary is no longer the bottleneck. The next speedup lives inside the kernel (Part VII) or across cores (Part VIII), not in more boundary cleverness.