Week 4 · Lesson 4

Fuse or materialize

The syllabus says it plainly: "We have to confront questions about materialization here." Rangeify answers "what is your element at index i?" one loop at a time. When a value is used in two loops, it gets asked twice. You can let it compute the answer twice, or store the answers once and read them back. That choice is materialization.

A value with two uses

Here is a sum-of-squares normalization: square every element, divide each square by the total of the squares. The squares, %d, have two uses.

%v  = load %x
%d  = mul %v %v              # the squares
%s  = reduce %d op=add       # use 1: the total
%r  = recip %s
%e  = expand %r shape=[4]
%w  = mul %d %e              # use 2: each square / total
%st = store %y %w

The reduce becomes one loop and the store becomes another (lesson 4.3). Each loop asks %d for its element at the loop's own index, so %d is asked once per element in each loop. Nothing in rangeify connects the two questions.

Recompute: fusion

Without stage, the spec lets your compiler choose either way. The oracle always recomputes, and stage forces a temp. Recomputing means rangeify answers each question by computing the element where it's needed. The square is fused into both loops. No temp memory, no extra pass, but the multiply happens twice per element.

Put %q = stage %d in front of both uses, and rangeify does something different: it fills a temp buffer with every element of %d in a loop of its own, then answers both questions with a load from that buffer. The value is materialized. One multiply per element, but a write and two reads of a 4-element temp.

without stage (recompute)

float a3[1] = {0};
float v7 = b0[0];
float v8 = (v7 * v7);
a3[0] = v8;
for (int32_t r14 = 0; r14 < 3; r14++) {
  float v15 = a3[0];
  int32_t v16 = (int32_t)((uint32_t)r14 + (uint32_t)1);
  int32_t v17 = (v16 % 4);
  float v19 = b0[v17];
  float v20 = (v19 * v19);   /* square */
  float v21 = (v15 + v20);
  a3[0] = v21;
}
float v25 = a3[0];
float v26 = (1.0f / v25);
for (int32_t r27 = 0; r27 < 4; r27++) {
  float v30 = b0[r27];
  float v31 = (v30 * v30);   /* square, again */
  float v32 = (v31 * v26);
  b1[r27] = v32;
}

with %q = stage %d on both uses (materialize)

float a3[4] = {0};
for (int32_t r4 = 0; r4 < 4; r4++) {
  int32_t v6 = (r4 % 4);
  float v8 = b0[v6];
  float v9 = (v8 * v8);      /* square, once */
  a3[r4] = v9;               /* into the temp */
}
float a12[1] = {0};
float v16 = a3[0];
a12[0] = v16;
for (int32_t r22 = 0; r22 < 3; r22++) {
  float v23 = a12[0];
  int32_t v24 = (int32_t)((uint32_t)r22 + (uint32_t)1);
  int32_t v25 = (v24 % 4);
  float v27 = a3[v25];       /* read the temp */
  float v28 = (v23 + v27);
  a12[0] = v28;
}
float v32 = a12[0];
float v33 = (1.0f / v32);
for (int32_t r34 = 0; r34 < 4; r34++) {
  float v37 = a3[r34];       /* read the temp */
  float v38 = (v37 * v33);
  b1[r34] = v38;
}

Two loops become three. Fusion multiplies twice per element. Materializing multiplies once but writes four floats and reads them twice. The % 4 lines come from decoding a flat loop counter into an index. Here they never change the value, and lesson 5.5 shows how a simplifier removes them.

Two outputs from one value

Slot 0 is x (4 f32). Write x * x into slot 1 (4 f32) and the sum of x * x into slot 2 (1 f32), using stage so both outputs read one materialized copy of the squares.

The answer is the same either way

You can choose freely because the choice can't change the result. An elementwise op computes the same bits for the same inputs wherever it runs, and every reduce folds its elements in the same order, left to right from the first. So the recomputed square and the stored square are the same float, and both sums add the same floats in the same order.

The three versions of the program above (no stage, stage on one use, stage on both) compile to kernels that, with x = [1.0 2.0 3.0 4.0], leave exactly the same bits in y: 3d088889 3e088889 3e99999a 3f088889, which is 1/30, 4/30, 9/30, 16/30 rounded to f32.

Materialization is a performance decision

Because results can't change, stage never fixes a wrong answer, and leaving it out never breaks a right one. The grader in week 4 checks only bits and size. Week 6 starts timing kernels, and that's when the choice starts to cost or save you something.

A cost model in words

Recomputing costs arithmetic: every extra use repeats the whole chain that produces the value. Materializing costs memory traffic: one write of n elements, one read of n per use, and n elements of temp storage.

Lesson 0.2 is the reason memory usually loses. A multiply takes about a nanosecond, and a trip to DRAM takes about 100 nanoseconds. For a value as cheap as v * v, recomputing it in each loop is almost always faster than writing it out and reading it back. Materializing starts to pay when the value is expensive to compute (a long chain of ops, or a reduce of its own) and has many uses.

Staging a plain load is pure loss. stage (load %x) copies x into a temp and then reads the copy instead of the original. Same reads, plus a write.
Two reduces over one vector already work without a temp: each reduce is its own loop reading x. The oracle even shares the load of x[0] that seeds both folds, because hash consing makes it one node.
A value used once is never worth staging: there's nothing to share.

stage wraps a value, not a name

%q = stage %d is a new value. Uses of %q read the temp, and uses of %d still recompute. Stage only the store's use and leave the reduce reading %d, and the oracle's C squares in the fill loop and again in the reduce loop: three loops, and the reduce saved nothing. To share, every use has to go through the staged value.

Variance

Slot 0 is x (8 f32). Write its variance into the one-element buffer in slot 1, computed exactly as

mean = (sum of x) * 0.125
c    = x + expand(-mean)          # -mean is mean * -1.0
var  = (sum of c * c) * 0.125

Both sums are reduce op=add. Results are compared bit for bit.

Which nodes are worth staging

A graph is a list of (id, srcs) pairs. Return, sorted, every non-leaf id that has two or more distinct consumers. A node read twice by the same consumer (like mul %c %c) is still one use, since rangeify builds that element once per index. That's the candidate list a materialization pass starts from.

Check yourself

Without stage, how many times per element does the oracle's C compute v * v for the sum-of-squares normalization above?

How many loops ask %d for an element?

Twice: once in the reduce loop and once in the store loop. Each loop asks for %d's element at its own index, and without a temp the square is recomputed for each.

With %q = stage %d on both uses, how many for loops does the oracle render?

The reduce and the store are still loops. What does stage add?

Three: the loop that fills the temp with the squares, the reduce loop and the store loop.

Adding stage in front of a value changes which of these?

The temp holds f32 values, and the value was already f32.

Storing an f32 into an f32 buffer and reading it back gives the same bits, and reduces fold in the same order either way. The C changes (an extra loop and a temp) but the result can't.

Which stage is always a loss?

What does the temp contain, and where did those floats come from?

A staged load is a copy of memory you could read directly. It adds a write and a temp without saving any arithmetic.

You write %q = stage %d, use %q in the store, and leave the reduce reading %d. What happens?

%q and %d are two different nodes.

Stage makes a new value. Only its uses read the temp. The reduce still asks %d directly, so the square is computed in the fill loop and again in the reduce loop.

On a typical machine, why does recomputing v * v usually beat materializing it?

Lesson 0.2: compare a cycle with a trip to memory.

Arithmetic is cheap and memory traffic is not. Materializing a value pays only when computing it is expensive enough to outweigh the extra writes and reads. Week 4 doesn't time anything yet, but week 6 does.