Week 3 · Lesson 2
Memory in loops
A loop body runs again and again over the same storage. Week 2's rule, that every load names the version it reads, still holds, but "the version" now changes every iteration. This lesson covers the accumulator, the one placement rule that makes it work, and the case where a load reads the contents from before the loop.
An accumulator
Summing a buffer means keeping a running total in memory. Here it lives in a one-element alloc:
%in = buffer 4 dtype=i32 slot=0 data=[1 2 3 4]
%out = buffer 1 dtype=i32 slot=1
%acc = alloc 1 dtype=i32 # starts as 0
%k = range 4
%a = shrink %acc 0 shape=[]
%cur = load %a # the running total so far
%iv = shrink %in %k shape=[]
%x = load %iv # in[k]
%sum = add %cur %x
%s = store %a %sum # acc = acc + in[k]
%e = end %s %k
%acc1 = after %acc %e # acc after the whole loop
%a1 = shrink %acc1 0 shape=[]
%tot = load %a1 # 10
%o = shrink %out 0 shape=[]
%so = store %o %tot
%out1 = after %out %soLook at %cur. It doesn't depend on %k, so by the open-ranges rule from lesson 3.1 it would be computed once, before the loop. That would read 0 once, and every iteration would store 0 + in[k], leaving 4 instead of 10. Loads need one more rule.
A load is placed like any other value: inside every loop whose range it depends on, and no deeper. One exception makes accumulators work. Take the innermost loop that contains every use of the load. If that loop is deeper than the load's own place, and it stores into the version the load reads, the load moves into it: it runs on every iteration and sees the current contents (the previous iteration's store, or the contents before the loop on the first iteration).
A use counts as inside a loop only if it's placed there itself: because it depends on that loop's range, or because it's a load this same rule moved there.
%cur's only use, %sum, is inside the k loop, and the k loop stores into %acc, the version %cur reads. So %cur is loaded fresh each iteration. Step through it:
%tot reads %acc1, the version after the loop. Its use is outside the loop, so it runs once, after the loop closes. Here's the C the oracle renders, trimmed of the wraparound casts:
int32_t a4[1] = {0};
for (int32_t r7 = 0; r7 < 4; r7++) {
int32_t v8 = a4[0]; /* %cur: inside the loop */
int32_t v10 = b0[r7]; /* in[k] */
a4[0] = v8 + v10;
}
b1[0] = a4[0]; /* %tot: after the loop */Row sums
Slot 0 is a 3-by-4 i32 matrix, row-major. Write each row's sum into slot 1: out[i] = A[i*4+0] + ... + A[i*4+3]. Slot 1 starts with junk, so every out[i] has to be reset before its row is summed. The starter forgets.
Split uses read the old contents
If a load is used both inside and outside a loop, the innermost loop holding every use is the one around the loop. The load runs once, before the loop, and every use sees the contents from before the loop:
%b = buffer 2 dtype=i32 slot=0 data=[5 0]
%acc = shrink %b 0 shape=[]
%cur = load %acc # used inside AND outside the loop
%k = range 3
%one = add %k 1
%sum = add %cur %one # use 1: inside the loop
%s = store %acc %sum
%e = end %s %k
%b1 = after %b %e
%dst = shrink %b1 1 shape=[]
%s2 = store %dst %cur # use 2: after the loop
%b2 = after %b1 %s2Every iteration computes 5 + (k + 1) from the same 5. The last one, k = 2, leaves b[0] = 8, and b[1] gets the old 5. Drop the second use and %cur moves into the loop: b[0] becomes 5 + 1 + 2 + 3 = 11.
Because UOps are hash consed, "the load of b[0]" is one node. You can't have it be loop-carried for one use and pre-loop for another. And writing load %acc a second time doesn't help: two identical loads are the same node. If you need both the running value and the starting value, load through two different view nodes of the same element, for example one at offset 0 and one at an offset the graph sees as different, such as mul %k 0. With that change, the program above gives b[0] = 11 (loop-carried) and b[1] = 5 (the starting value).
This loop is meant to add 1 to b[0] three times, starting from 10. It writes 11:
%b = buffer 1 dtype=i32 slot=0 data=[10]
%i = range 3
%z = mul %i 0
%v = shrink %b %z shape=[] # the store's view depends on %i
%v0 = shrink %b 0 shape=[] # the load's view doesn't
%x = load %v0
%y = add %x 1 # depends on no range
%s = store %v %y
%e = end %s %i
%b1 = after %b %e%y depends on no range, so it's placed outside the loop, and so is its load: both run once, and every iteration stores the same 11. Load through the store's own view (%x = load %v) and the load depends on %i, runs every iteration, and the result is 13.
Resets and chains of versions inside a loop
The row-sum exercise needs a reset per row. The reset store has to depend on %i so it runs inside the i loop, and it does: it stores into shrink %out %i. Then the j loop accumulates into the version after the reset:
%i = range 3
%o = shrink %out %i shape=[]
%z = store %o 0 # reset out[i]: runs once per i
%out0 = after %out %z # a version that lives inside the i loop
%j = range 4
...
%o0 = shrink %out0 %i shape=[]
%cur = load %o0 # loop-carried in the j loop
%s = store %o0 %sum
%ej = end %s %j
%ei = end %ej %i
%out1 = after %out %ei # the output: no open ranges%out0 has {i} open, so it's local to one iteration of the i loop. Only an after with nothing open counts as the buffer's output, which is why %out1 builds on %out, not on %out0. When the effect is an end, the rule is that the loop must store into the same memory after names.
The largest of five
Slot 0 holds five f32s. Put their maximum into slot 1's only element. Slot 1 starts with junk, so start the running max from something no input can beat.
Check yourself
The split-use program above, without the second store. What's the final b[0]?
%b = buffer 2 dtype=i32 slot=0 data=[5 0]
%acc = shrink %b 0 shape=[]
%cur = load %acc
%k = range 3
%one = add %k 1
%sum = add %cur %one
%s = store %acc %sum
%e = end %s %k
%b1 = after %b %eIs %cur used anywhere outside the loop now?
Its only use is inside the k loop, which stores into the version it reads, so it's loop-carried: 5 + 1 + 2 + 3 = 11.
And the full split-use program, with the second store into b[1]. What are the final contents?
Where does %cur run when one of its uses is outside the loop?
%cur runs once, before the loop, and reads 5. Each iteration stores 5 + (k + 1), so the last leaves 8. b[1] gets the pre-loop 5: [8 5].
In the accumulator example, where does %tot = load %a1 run?
Which version does it read, and where is its use?
%tot reads %acc1, the version after the loop, and its use (the store into %out) is outside every loop. So it runs once, after the loop's closing brace.
A reset store store %o 0 with %o = shrink %acc 0 shape=[] sits next to an i loop and a j loop. How often does it run?
What are the store's open ranges?
The store depends on no range, so it has nothing open and runs once, outside both loops. To reset once per row, the store has to depend on %i, for example by storing into shrink %out %i.
Which ops may consume an effect (a store or an end)?
Effects are void.
Effects have dtype void, and only after (to name the memory after it) and end (to repeat it) take them.
Why does %out1 = after %out %ei build on %out and not on %out0, the version after the reset?
What are %out0's open ranges?
%out0 depends on the reset store, which depends on %i. A version with an open range exists once per iteration, so it can't be what the buffer holds after the loop. The output is the after with nothing open, built on the version before the loops.