Week 3 · Lesson 4

Convolutions

A convolution slides a small kernel along an input and takes a dot product at every position. With range and end you can write one by hand: two loops for a 1D convolution, four for a 2D one. The only new difficulty is padding, which reads outside the input, and you already have the tool for that: prove the index in range, then zero what doesn't belong.

A valid 1D convolution

Take an input x of length 6 and a kernel w of length 3. Output position i is the dot product of the kernel with the window of x that starts at i:

y[i] = x[i]*w[0] + x[i+1]*w[1] + x[i+2]*w[2]

"Valid" means the window never leaves the input, so there are 6 - 3 + 1 = 4 outputs. Step through it with x = [1 2 3 4 5 6] and w = [1 0 -1].

In .uir that's a loop over i with a loop over k inside it. The accumulator is the output element itself: reset y[i] to zero with a store that depends on i, then add into the version after that store, once per k.

%i   = range 4
%yv  = shrink %y %i shape=[]
%z   = store %yv 0.0          # reset y[i], once per i
%y0  = after %y %z
%k   = range 3
%ik  = add %i %k
%xv  = shrink %x %ik shape=[] # x[i + k]: interval [0, 3] + [0, 2] = [0, 5]
%xe  = load %xv
%wv  = shrink %w %k shape=[]
%we  = load %wv
%p   = mul %xe %we
%yv0 = shrink %y0 %i shape=[]
%cur = load %yv0              # runs inside the k loop: its only use is there
%sum = add %cur %p
%s   = store %yv0 %sum
%ek  = end %s %k
%ei  = end %ek %i
%y1  = after %y %ei

No masking is needed: the range intervals alone prove i + k stays in [0, 5]. The oracle renders it as two loops, with the reset store in the outer one:

for (int32_t r6 = 0; r6 < 4; r6++) {
  b2[r6] = 0x0.0p+0f;
  for (int32_t r11 = 0; r11 < 3; r11++) {
    float v12 = b2[r6];
    int32_t v13 = (int32_t)((uint32_t)r6 + (uint32_t)r11);
    float v15 = b0[v13];
    float v17 = b1[r11];
    float v18 = (v15 * v17);
    float v19 = (v12 + v18);
    b2[r6] = v19;
  }
}

Write the valid convolution

Slot 0 is x (6 i32), slot 1 is w (3 i32), slot 2 is y (4 i32). Make the final contents of slot 2 equal y[i] = x[i]*w[0] + x[i+1]*w[1] + x[i+2]*w[2] for every i in 0..3. The starter only uses w[0]. The tests set slots 0 and 1. i32 sums wrap as usual.

Zero padding

Usually you want as many outputs as inputs. That means the window hangs off both ends: for a kernel of 3, y[0] needs x[-1] and y[5] needs x[6]. "Zero padding" says those reads count as zero.

y[i] = x[i-1]*w[0] + x[i]*w[1] + x[i+1]*w[2]     with x = 0 outside [0, 6)

Write the index as j = i + k - 1. With i in [0, 5] and k in [0, 2], j is in [-1, 6], and a shrink at j fails the bounds proof from lesson 2.3. You need two separate things:

  1. An offset that's always legal, even when it points at the wrong element. max j 0 is in [0, 6], still one too many. mod (max j 0) 6 is in [0, 5]: provably inside x.
  2. A mask that throws the wrong element away. j is a real element exactly when 0 ≤ j and j < 6. When it isn't, use 0.0 instead of the loaded value.
%j  = add %ik -1              # i + k - 1, in [-1, 6]
%lo = cmplt %j 0              # off the left end
%hi = cmplt %j 6              # not off the right end
%ok = where %lo false %hi     # 0 <= j < 6
%jp = max %j 0                # [0, 6]
%jm = mod %jp 6               # [0, 5]: provably in bounds
%xv = shrink %x %jm shape=[]
%xe = load %xv                # a real element, maybe the wrong one
%xz = where %ok %xe 0.0       # the padding value where j is off the end
The mask doesn't prove anything

It's tempting to think where %ok ... makes the raw index safe, because the bad element is never used. But the load still happens, and the bounds proof only looks at the offset's interval. A where doesn't narrow an interval (lesson 2.3). The offset itself has to be provably legal, which is why you need both the clamp and the mask.

Same-size convolution with zero padding

Slot 0 is x (6 f32), slot 1 is w (3 f32), slot 2 is y (6 f32). Make the final contents of slot 2

y[i] = x[i-1]*w[0] + x[i]*w[1] + x[i+1]*w[2]

for every i in 0..5, where any x outside [0, 6) counts as 0.0. Every view must be provably in bounds, so mask the index instead of shrinking at i + k - 1 directly. Add terms in order k = 0, 1, 2, starting from 0.0, so the f32 rounding matches.

2D convolution on flat buffers

A 2D input of H rows and W columns lives in a flat row-major buffer, so element (r, c) is at r*W + c (lesson 0.4). A valid 2D convolution with a KH by KW kernel has four loops: output row and column, kernel row and column.

y[oi*OW + oj] = sum over ki, kj of x[(oi + ki)*W + (oj + kj)] * w[ki*KW + kj]
OH = H - KH + 1,  OW = W - KW + 1

Every offset is a sum of products of ranges and constants, so the interval analysis proves it automatically. For a 4 by 4 input and a 2 by 2 kernel, oi is in [0, 2] and ki in [0, 1], so oi + ki is in [0, 3], and the same goes for oj + kj. That puts (oi + ki)*4 + (oj + kj) in [0, 3*4 + 3] = [0, 15]. The oracle's C for that program hoists each index computation to the shallowest loop it can live in:

for (int32_t r7 = 0; r7 < 3; r7++) {               // oi
  int32_t v8 = (int32_t)((uint32_t)r7 * (uint32_t)3);
  for (int32_t r9 = 0; r9 < 3; r9++) {             // oj
    int32_t v10 = (int32_t)((uint32_t)v8 + (uint32_t)r9);
    b2[v10] = 0;
    for (int32_t r15 = 0; r15 < 2; r15++) {        // ki
      int32_t v16 = (int32_t)((uint32_t)r7 + (uint32_t)r15);
      int32_t v17 = (int32_t)((uint32_t)v16 * (uint32_t)4);
      int32_t v18 = (int32_t)((uint32_t)r15 * (uint32_t)2);
      for (int32_t r19 = 0; r19 < 2; r19++) {      // kj
        int32_t v20 = b2[v10];
        ...

Lesson 5 explains where each of those lines goes and why.

2D convolution in C

Write void conv2d(const int *x, int H, int W, const int *k, int KH, int KW, int *y): a valid 2D convolution on flat row-major arrays. x is H by W, k is KH by KW, and y is (H-KH+1) by (W-KW+1). Overwrite every element of y. The values are small, so plain int arithmetic is fine here.

A convolution is a matrix multiply in disguise

Look at the valid 1D convolution again. Every output is a dot product of the kernel with one window. Write the windows as the rows of a matrix:

windows (4 x 3)      w (3)        y (4)
[1 2 3]              [ 1]         [-2]
[2 3 4]      @       [ 0]    =    [-2]
[3 4 5]              [-1]         [-2]
[4 5 6]                           [-2]

That windows matrix doesn't need to exist in memory. Row i, column k of it is just x[i + k], the same element read by several rows. The whole trick of a convolution is that index map: (i, k) → i + k. Week 5 introduces movement ops (reshape, permute, pad) that describe exactly these index maps on shapes, so you'll write conv as a few movement ops followed by a multiply and a reduce, and the compiler will produce these loops for you. Padding becomes a pad op, and the masking you wrote by hand becomes something the compiler emits.

Check yourself

A valid 1D convolution has an input of length 10 and a kernel of length 3. How many outputs are there?

The window has to fit entirely inside the input. Where can its first element start?

8: the window can start at 0 through 7, so 10 - 3 + 1 outputs.

With x = [1 2 3 4 5 6] and w = [1 0 -1], what is y[2] for the valid convolution?

The window for output 2 is x[2], x[3], x[4].

-2: 3*1 + 4*0 + 5*(-1). Every output of this kernel is the difference of two elements two apart, and x rises by 1 per element, so all four outputs are -2.

For the padded convolution, j = i + k - 1 with i in [0, 5] and k in [0, 2]. Which offset can you shrink a 6-element x at?

Work out each interval. C's mod takes the sign of the left operand.

A view of one element at offset [lo, hi] needs lo ≥ 0 and hi + 1 ≤ 6.

mod j 6 is in [-5, 5], because j can be negative. max j 0 is in [0, 6], one past the end. Only mod (max j 0) 6 is in [0, 5]. The oracle rejects the other three. A where never narrows an interval.

In the convolution, why does the reset store write to y[i] instead of to a one-element alloc?

Which ranges does store %a 0.0 depend on, if %a is a view at a constant offset?

Placement follows open ranges. A reset that depends on i (through its offset) runs once per i, and one that depends on nothing runs once in total. Accumulating straight into the output at an offset built from the outer ranges gives you the reset for free.

A 2D input is H = 5 rows by W = 7 columns, flat and row-major. Which offset is output (oi, oj)'s kernel tap (ki, kj)?

Row-major means you multiply the row by the number of columns.

Row oi + ki, column oj + kj, and a row is 7 elements long, so (a). Option (b) multiplies by the number of rows instead of the row length, and (c) is the column-major formula, which walks down columns.

The valid convolution of a length-6 input with a length-3 kernel equals windows @ w. What's the shape of windows? (rows, columns)

One row per output, one column per kernel tap.

4 by 3: four outputs, three taps. Row i, column k is x[i + k], so the 12 entries are only 6 distinct input elements.