Case Study 1: Vectorizing a Scalar Loop

"The loop you did not write is the loop you cannot get wrong."

Executive Summary

You have inherited a small analysis routine that computes, for a table of measurements, the mean of each column and then each value's deviation from its column mean — a temperature anomaly, in the language of climate science. It works, but it is written in the element-at-a-time style of a language that does not understand arrays: three nested loops, a running-sum accumulator, an index bookkeeping you have to check by eye. This study reads that code, understands exactly what it does, and then ports it to itself — rewrites it in modern Fortran's whole-array style — verifying that the answer is bit-for-bit the same while the code shrinks to three lines that read like the definition of the quantity. The skill you are practicing is the one that matters most in this chapter: seeing a loop and recognizing the array operation hiding inside it.

Skills applied: reading and analyzing existing scalar code; reductions along a dimension, sum(a, dim=) (§5.4); array sections and whole-array subtraction (§5.2–5.3); spread for shape-matching; column-major loop order (§5.6).

Background

The data is a table: rows are weather stations, columns are months. Entry data(i, j) is the temperature recorded by station i in month j. Our tiny example has three stations and four months:

station \ month Jan Feb Mar Apr
1 10 12 5 21
2 20 14 10 22
3 30 16 15 23

Two quantities are wanted. First, the column mean — the average across stations for each month. Second, the anomaly — each reading minus its own month's mean, which tells you how far above or below normal that station-month was. By hand, the month means are 20, 14, 10, 22, and, for instance, station 3 in January is 30 - 20 = +10: ten degrees above that month's average.

Phase 1 — Read the Data Layout

Before touching the computation, notice the layout, because it drives everything. The table is a rank-2 array data(station, month) — station is the first index, month the second. By §5.6, Fortran stores this column by column, so all three stations of January sit contiguously in memory, then all three of February, and so on. A "column" in our table is one month's readings across stations, and that is exactly the natural unit for the column mean. The layout and the computation are aligned; keep them that way.

Phase 2 — The Scalar Code, Understood

Here is the routine as you found it. Read it and narrate what each loop does before reading our narration:

! The inherited scalar version: correct, but element-at-a-time.
do j = 1, nm                          ! for each month (column)
  s = 0.0_dp
  do i = 1, ns                        ! sum this month's readings over stations
    s = s + data(i, j)
  end do
  colmean(j) = s / real(ns, dp)       ! the month mean
  do i = 1, ns                        ! subtract the mean from each reading
    anom(i, j) = data(i, j) - colmean(j)
  end do
end do

Three loops, doing three jobs: accumulate a column sum, divide it to get the mean, then subtract the mean back through the column. It is correct, and — credit where due — its loop order is already cache-friendly: the outer loop is over months (columns) and the inner work runs down each column over the first index, with the memory grain. But it is fifteen lines of index bookkeeping for an idea you can state in one sentence, and every hand-written index (data(i, j), anom(i, j)) is a chance to write j where you meant i.

Phase 3 — Port It to Whole-Array Form

Now rewrite it as arrays. The column sum over stations is a reduction along the first dimension — sum(data, dim=1) — which collapses the (ns, nm) table to a length-nm row of monthly totals. Divide by the number of stations and you have every month mean in one statement:

colmean = sum(data, dim=1) / real(ns, dp)      ! rank-1: one mean per month

The anomaly subtracts each month's mean from each reading. Two clean ways to say it. The explicit one loops over months and subtracts a scalar from a column section:

do j = 1, nm
  anom(:, j) = data(:, j) - colmean(j)         ! whole column minus its scalar mean
end do

The fully vectorized one uses spread to replicate colmean into a full (ns, nm) matrix — one identical row per station — so the subtraction is a single whole-array operation:

anom = data - spread(colmean, dim=1, ncopies=ns)   ! no loop at all

spread(colmean, dim=1, ncopies=ns) takes the length-nm vector of means and stacks ns copies of it as rows, producing exactly the (ns, nm) array whose every row is the month means — precisely what you must subtract. The complete program:

program anomalies
  use, intrinsic :: iso_fortran_env, only: dp => real64
  implicit none
  integer,  parameter :: ns = 3, nm = 4        ! 3 stations, 4 months
  real(dp) :: data(ns, nm), anom(ns, nm), colmean(nm)
  integer  :: i

  data = reshape([ 10.0_dp, 12.0_dp,  5.0_dp, 21.0_dp,   &   ! station 1
                   20.0_dp, 14.0_dp, 10.0_dp, 22.0_dp,   &   ! station 2
                   30.0_dp, 16.0_dp, 15.0_dp, 23.0_dp ], &   ! station 3
                 [ns, nm], order=[2,1])

  colmean = sum(data, dim=1) / real(ns, dp)             ! one statement
  anom    = data - spread(colmean, dim=1, ncopies=ns)   ! one statement

  print '(a, 4f7.1)', 'month mean = ', colmean
  print '(a)',        'anomalies (station x month):'
  do i = 1, ns
    print '(4f7.1)', anom(i, :)
  end do
end program anomalies
$ gfortran -std=f2018 -Wall anomalies.f90 -o anom && ./anom
month mean =    20.0   14.0   10.0   22.0
anomalies (station x month):
  -10.0   -2.0   -5.0   -1.0
    0.0    0.0    0.0    0.0
   10.0    2.0    5.0    1.0

Phase 4 — Verify Identical Output

Never trust a "port" you have not checked. Work the numbers by hand and match them to the printout. The month means are (10+20+30)/3, (12+14+16)/3, (5+10+15)/3, (21+22+23)/3 = 20, 14, 10, 22 — matches. Station 2 records exactly the mean every month (20, 14, 10, 22), so its anomalies are all zero — the middle row of zeros confirms it. Station 1 is uniformly below normal (-10, -2, -5, -1) and station 3 uniformly above (+10, +2, +5, +1), and each above/below pair is symmetric because with three values the outer two straddle the mean. Every printed number is accounted for. The scalar and array versions agree.

Phase 5 — Why the Array Version Wins Twice

The rewrite is shorter, but shorter is the least of it. It wins on correctness: there are no hand-written indices to transpose, so the entire class of "i where you meant j" bugs is gone — you cannot misindex a loop you did not write. And it wins on speed potential: sum(data, dim=1) and the whole-array subtract hand the compiler the complete operation over a known shape, with Fortran's no-aliasing guarantee, which is exactly what it needs to vectorize — and, if you ask, to parallelize. The scalar version can be optimized by a good compiler, but the array version states the parallelism instead of hiding it in a loop the compiler must first prove safe to transform. We measure that difference on real hardware in Chapter 27.

One caution, to keep you honest: spread materializes a temporary (ns, nm) array of means. For this table that is nothing; for a field of billions it is real memory traffic, and there the column-section loop (which subtracts a scalar and allocates nothing) can be the better choice. Readable-and-fast is the goal, but "fast" always deserves a measurement — the discipline of Chapter 28.

Discussion Questions

  1. The scalar version's loop order was already cache-friendly (inner work over the first index). Rewrite it deliberately wrong — swap so the inner work runs across rows — and explain, in memory terms, why it would be slower even though the answer is unchanged.
  2. sum(data, dim=1) reduces over stations; what does sum(data, dim=2) reduce over, and what would it mean physically for this table?
  3. The spread form allocates a temporary; the section-loop form does not. Describe a data size at which you would switch from the elegant form to the frugal one, and how you would decide (rather than guess).

Your Turn: Extensions

  • Option A. Add a row analysis: compute each station's annual mean with sum(data, dim=2) / real(nm, dp), and its anomaly from the global mean. Verify by hand on the 3×4 table.
  • Option B. Replace spread with the explicit column-section loop and confirm the output is identical. Which do you find more readable? Which would you ship for a million-row table, and why?
  • Option C. Generalize the program to read ns and nm and the data from a file at run time using an allocatable data(:,:) (§5.5). What must change, and what stays exactly the same? (This is the bridge to Chapter 7.)

Key Takeaways

  • A nested loop that accumulates and then transforms a column is almost always a reduction (sum(a, dim=)) plus a whole-array operation — learn to see the array expression inside the loop.
  • Porting scalar code to array form removes hand-written indices, and with them an entire class of bugs; the array form is not just shorter but safer.
  • Whole-array expressions state structure the compiler can vectorize; scalar loops make it reconstruct that structure first. Same answer, more room to go fast.
  • Elegance has a cost worth watching: intrinsics like spread allocate temporaries. Prefer the readable form, but let a measurement, not a feeling, decide when to trade it for the frugal one.