28. Tensor Contractions in Mountain Goat Itself¶

Mountain goats on the mountain drawn as a blueprint

What you will understand: how Mountain Goat gets tensors of rank 3 and more, and how a general tensor contraction becomes a short Mountain Goat program with no C++ in it. Chapter 27 did the axis bookkeeping (permuting and reshaping) in C++ and left only the matrix product to the language; its "Limits" section said a language-level reshape and permute were not built. This chapter builds them: two new compiler operations, mg.reshape and mg.permute, and three new built-ins, reshape, permute and contract. Everything is checked against the definition of a contraction (plain nested loops, no matrix product) on thirteen shape combinations, and against the worked examples of the Tensor Contractions (CPU) appendix of the companion book CUDA From First Principles.

What you need to know first: Chapter 27 (what a contraction is and why it is a matrix product in disguise: read its primer first), Chapter 22 (how an operation is added to the dialect, the lowering and the front end) and the language tour.

Compile and run

cd docs/part28/code && ./build.sh                 # once: the newest compiler (adds mg.reshape and mg.permute)
./mgc run examples/03_double_contraction.mg       # compile and run one example
./mgc mlir examples/04_axis_at_the_front.mg       # show the MLIR the front end produces
./run_examples.sh > examples_out.txt              # every example and every error example on this page
./check_contractions.py                           # 13 shape combinations against the definition, exact
./mutation.sh > mutation_out.txt                  # breaks the compiler on purpose, 21 ways (about ten minutes)
cd ../../part15/code && ./run_lit.sh              # the whole test suite

Every listing and output on this page comes from these commands.

Why this page shows FAIL lines, and why that is good

Two places show failures on purpose. The error examples are programs the compiler must reject: each ends with exit status 1 and produces no program. The mutation output comes from a negative control: the compiler deliberately broken and run against the tests, where a "caught" line (a test failed) is the wanted result and "NOT CAUGHT" would be a test gap.

What was missing¶

Until now every Mountain Goat value was a two-dimensional matrix. Chapter 27's recipe, permute, reshape, multiply, reshape, needs tensors of any rank and the two rearranging steps. Chapter 27 kept those in C++ and noted that the language itself could do them with three additions: tensors of rank greater than two, a reshape, and a permute. Here they are.

The three new built-ins¶

Written	Meaning	Checked at compile time
`reshape(x, d0, d1, ...)`	the same numbers, viewed with shape `d0 x d1 x ...` (row-major order; no data moves)	`x` has a static shape; the element counts are equal
`permute(x, p0, p1, ...)`	reorder the axes: result axis `i` is input axis `pi`	`p` lists each axis exactly once
`contract(a, (i, j, ...), b, (k, l, ...))`	sum the products over the paired axes: `a`'s axis `i` with `b`'s axis `k`, and so on	both lists the same length; every axis in range and not repeated; every pair the same size

Literals stay matrices, so a rank-3 value is made by reshaping one: reshape([[1, 2, ..., 12], [13, ..., 24]], 2, 3, 4) is a 2x3x4 tensor whose first block is the numbers 1 to 12. Only five things can be done with a value of rank other than 2: reshape it, permute it, contract it, bind it with let, or print it. Arithmetic, transpose, reductions and function arguments need a matrix, and say so (see error 7).

One by one¶

1. Reshape and permute¶

################ mgc run 01_reshape_and_permute.mg
  | # Tensors of rank 3 and more. A literal is always a matrix; reshape(x, d0, d1, ...) views its numbers with a new shape (row-major, no data moves).
  | let a = reshape([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], [13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24]], 2, 3, 4)   # 2 blocks of 3x4
  | print a
  | # permute(x, p0, p1, p2): result axis i is input axis pi. Swapping the first two axes gives a 3x2x4 tensor.
  | print permute(a, 1, 0, 2)
  | # A "rotation" of the axes (a 3-cycle): result axes are input axes 2, 0, 1, so the shape 2x3x4 becomes 4x2x3.
  | print permute(a, 2, 0, 1)
  | # reshape undoes nothing and copies nothing: the same 24 numbers in a 4x6 matrix.
  | print reshape(a, 4, 6)
  ----- output
Unranked Memref base@ = 0x… rank = 3 offset = 0 sizes = [2, 3, 4] strides = [12, 4, 1] data = 
[[[1,    2,    3,    4], 
  [5,    6,    7,    8], 
  [9,    10,    11,    12]], 
 [[13,    14,    15,    16], 
  [17,    18,    19,    20], 
  [21,    22,    23,    24]]]
Unranked Memref base@ = 0x… rank = 3 offset = 0 sizes = [3, 2, 4] strides = [8, 4, 1] data = 
[[[1,    2,    3,    4], 
  [13,    14,    15,    16]], 
 [[5,    6,    7,    8], 
  [17,    18,    19,    20]], 
 [[9,    10,    11,    12], 
  [21,    22,    23,    24]]]
Unranked Memref base@ = 0x… rank = 3 offset = 0 sizes = [4, 2, 3] strides = [6, 3, 1] data = 
[[[1,    5,    9], 
  [13,    17,    21]], 
 [[2,    6,    10], 
  [14,    18,    22]], 
 [[3,    7,    11], 
  [15,    19,    23]], 
 [[4,    8,    12], 
  [16,    20,    24]]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [4, 6] strides = [6, 1] data = 
[[1,   2,   3,   4,   5,   6], 
 [7,   8,   9,   10,   11,   12], 
 [13,   14,   15,   16,   17,   18], 
 [19,   20,   21,   22,   23,   24]]
  ----- exit status 0

The first print shows the 2x3x4 tensor: two blocks of three rows of four. permute(a, 1, 0, 2) swaps the first two axes: result axis 0 is input axis 1, so the shape becomes 3x2x4, and the first result block holds row 0 of each original block (1 2 3 4 and 13 14 15 16). permute(a, 2, 0, 1) is a rotation of the axes (a 3-cycle): the shape 2x3x4 becomes 4x2x3. This example is worth keeping in mind because a swap undone is the same swap, but a 3-cycle is not: reading the permutation backwards gives a different tensor, and the tests below depend on that. The last line reshapes the same 24 numbers into a 4x6 matrix; it copies nothing.

2. A matrix product is a contraction¶

################ mgc run 02_matmul_is_a_contraction.mg
  | # A matrix product is the simplest contraction: sum over axis 1 of A (the columns) against axis 0 of B (the rows). The appendix's first example.
  | let a = [[1, 2], [3, 4], [5, 6]]
  | let b = [[1, 0, 2, 1], [0, 1, 1, 2]]
  | print contract(a, (1), b, (0))     # 3x4, equal to a @ b
  | print a @ b                         # the same numbers, written with the operator
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [3, 4] strides = [4, 1] data = 
[[1,   2,   4,   5], 
 [3,   4,   10,   11], 
 [5,   6,   16,   17]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [3, 4] strides = [4, 1] data = 
[[1,   2,   4,   5], 
 [3,   4,   10,   11], 
 [5,   6,   16,   17]]
  ----- exit status 0

contract(a, (1), b, (0)) sums over a's axis 1 against b's axis 0, and gives exactly a @ b. These are the appendix's numbers (its first worked example). The compiler produces no reshape and no permute for it: one mg.matmul (a test asserts that).

3. The appendix's double contraction¶

################ mgc run 03_double_contraction.mg
  | # A double contraction: [2,3,4] x [3,4,5], summing over A's axes 1 and 2 against B's axes 0 and 1.
  | # The data is the appendix's: A[i] = (i % 7) - 3 and B[i] = (i % 5) - 2, numbered in row-major order. Each literal is a matrix whose rows are
  | # the first axis of the tensor; reshape then views it with its real shape.
  | let a = reshape([[-3, -2, -1, 0, 1, 2, 3, -3, -2, -1, 0, 1], [2, 3, -3, -2, -1, 0, 1, 2, 3, -3, -2, -1]], 2, 3, 4)
  | let b = reshape([[-2, -1, 0, 1, 2, -2, -1, 0, 1, 2, -2, -1, 0, 1, 2, -2, -1, 0, 1, 2], [-2, -1, 0, 1, 2, -2, -1, 0, 1, 2, -2, -1, 0, 1, 2, -2, -1, 0, 1, 2], [-2, -1, 0, 1, 2, -2, -1, 0, 1, 2, -2, -1, 0, 1, 2, -2, -1, 0, 1, 2]]
  | print contract(a, (1, 2), b, (0, 1))      # (2) free axis of A, (5) free axis of B: a 2x5 result; expected [[10, 5, 0, -5, -10], [2, 1, 0, -1, -2]]
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 5] strides = [5, 1] data = 
[[10,   5,   0,   -5,   -10], 
 [2,   1,   0,   -1,   -2]]
  ----- exit status 0

A is 2x3x4 with A[i] = (i % 7) - 3, B is 3x4x5 with B[i] = (i % 5) - 2, both numbered in row-major order. Pairing A's axes {1, 2} with B's {0, 1} leaves A's axis 0 and B's axis 2: a 2x5 result, and it is the appendix's numpy.tensordot answer, [[10, 5, 0, -5, -10], [2, 1, 0, -1, -2]]. (The literals in this file were generated by examples/gen_big_literals.py, which is checked in.)

4. A contracted axis at the front¶

################ mgc run 04_axis_at_the_front.mg
  | # The contracted axis is the FIRST axis of A, so A must be permuted before it can be flattened: this is the case that makes the language
  | # move data. A is 2x3x4, B is 2x5; summing over A's axis 0 against B's axis 0 leaves A's axes 1, 2 and B's axis 1: a 3x4x5 result.
  | let a = reshape([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], [-1, -2, -3, -4, -5, -6, -7, -8, -9, -10, -11, -12]], 2, 3, 4)
  | let b = [[1, 0, 2, 0, 3], [0, 1, 0, 2, 0]]
  | print contract(a, (0), b, (0))
  ----- output
Unranked Memref base@ = 0x… rank = 3 offset = 0 sizes = [3, 4, 5] strides = [20, 5, 1] data = 
[[[1,    -1,    2,    -2,    3], 
  [2,    -2,    4,    -4,    6], 
  [3,    -3,    6,    -6,    9], 
  [4,    -4,    8,    -8,    12]], 
 [[5,    -5,    10,    -10,    15], 
  [6,    -6,    12,    -12,    18], 
  [7,    -7,    14,    -14,    21], 
  [8,    -8,    16,    -16,    24]], 
 [[9,    -9,    18,    -18,    27], 
  [10,    -10,    20,    -20,    30], 
  [11,    -11,    22,    -22,    33], 
  [12,    -12,    24,    -24,    36]]]
  ----- exit status 0

Here A's contracted axis is its first axis, so A cannot simply be read as a matrix: it must be permuted first. This is the one case in the family where the language moves data. The MLIR the front end writes shows the four steps (this is what mgc mlir prints, comments removed):

%t4 = mg.permute %t2 {permutation = array<i64: 1, 2, 0>} : tensor<2x3x4xf64> -> tensor<3x4x2xf64>
%t5 = mg.reshape %t4 : tensor<3x4x2xf64> -> tensor<12x2xf64>
%t6 = mg.matmul %t5, %t3 : tensor<12x2xf64>, tensor<2x5xf64> -> tensor<12x5xf64>
%t7 = mg.reshape %t6 : tensor<12x5xf64> -> tensor<3x4x5xf64>

Permute A so the free axes (1, 2) come first and the contracted axis 0 last, view the result as a 12x2 matrix, multiply by B (already a 2x5 matrix), view the 12x5 product as 3x4x5. That is Chapter 27's recipe, written by the compiler instead of by hand.

5. Axes in reverse order¶

################ mgc run 05_axes_in_reverse_order.mg
  | # The ORDER of the paired axes matters: the first axis listed for A meets the first listed for B. Here A's axes 2 and 1 pair with B's axes 0 and 1.
  | # A is 2x3x4 and B is 4x3x2, so the pairs have sizes 4=4 and 3=3. Listed the other way, (1, 2) against (0, 1), the sizes would not match.
  | let a = reshape([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1]], 2, 3, 4)
  | let b = reshape([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24]], 4, 3, 2)
  | print contract(a, (2, 1), b, (0, 1))      # a 2x2 result
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 2] strides = [2, 1] data = 
[[1090,   1168], 
 [90,   96]]
  ----- exit status 0

The order in which the pairs are listed matters: the first axis listed for A meets the first listed for B. Here (2, 1) against (0, 1) pairs A's axis 2 (size 4) with B's axis 0 (size 4) and A's axis 1 (size 3) with B's axis 1 (size 3). Listing them as (1, 2) against (0, 1) would pair sizes 3 and 4 with 4 and 3 and be an error.

6. Vectors: dot product, outer product, vector times matrix¶

################ mgc run 06_vectors_dot_and_outer.mg
  | # Rank 1 tensors exist too (reshape a one-row matrix to a single dimension). Three contractions:
  | let u = reshape([[1, 2, 3]], 3)
  | let v = reshape([[4, 5, 6, 7]], 4)
  | print contract(u, (0), u, (0))     # dot product: every axis summed away, 1+4+9 = 14 (returned as a 1x1 matrix)
  | print contract(u, (), v, ())       # outer product: nothing summed, a 3x4 table of products
  | print contract(reshape([[1, 2, 3, 4]], 4), (0), reshape([[1, 0, 1], [0, 1, 1], [1, 1, 0], [2, 0, 1]], 4, 3), (0))   # vector times matrix
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data = 
[[14]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [3, 4] strides = [4, 1] data = 
[[4,   5,   6,   7], 
 [8,   10,   12,   14], 
 [12,   15,   18,   21]]
Unranked Memref base@ = 0x… rank = 1 offset = 0 sizes = [3] strides = [1] data = 
[12,  5,  7]
  ----- exit status 0

A rank-1 tensor is made by reshaping a one-row matrix to a single dimension. Contracting every axis (the dot product) gives the scalar 14, returned as a 1x1 matrix (the language has no rank-0 values). Contracting nothing (the outer product) is a table of products. The third result is a rank-1 vector, which prints with sizes = [3].

7. A rank 4 result¶

################ mgc run 07_rank4_result.mg
  | # A rank 4 result: A is 2x2x3 and B is 3x2x2; contracting only the 3's (A's axis 2 with B's axis 0) leaves A's two free axes and B's two: 2x2x2x2.
  | let a = reshape([[1, 2, 3, 4, 5, 6], [7, 8, 9, 10, 11, 12]], 2, 2, 3)
  | let b = reshape([[1, 0, 2, 1], [0, 1, 1, 2], [1, 1, 0, 0]], 3, 2, 2)
  | print contract(a, (2), b, (0))
  ----- output
Unranked Memref base@ = 0x… rank = 4 offset = 0 sizes = [2, 2, 2, 2] strides = [8, 4, 2, 1] data = 
[[[[4,     5], 
   [4,     5]], 
  [[10,     11], 
   [13,     14]]], 
 [[[16,     17], 
   [22,     23]], 
  [[22,     23], 
   [31,     32]]]]
  ----- exit status 0

A is 2x2x3 and B is 3x2x2; contracting only the 3's leaves A's two free axes and B's two, 2x2x2x2.

8. Inside a function¶

################ mgc run 08_in_a_function.mg
  | # A contraction inside a function. Function parameters are matrices, so the function reshapes them; the shapes are checked once, when it is defined.
  | # It computes the double contraction of a 2x3x4 tensor (passed as a 2x12 matrix) with a 3x4x5 tensor (passed as a 3x20 matrix), for any values.
  | def double_contraction(a: tensor[2x12], b: tensor[3x20]) = contract(reshape(a, 2, 3, 4), (1, 2), reshape(b, 3, 4, 5), (0, 1))
  | let a = [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2]]
  | let b = [[1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5], [1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5], [1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5]]
  | print double_contraction(a, b)        # every output is a sum of 12 products: 12*(i+1)*(k+1) at row i, column k, because b repeats 1..5 along its last axis
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 5] strides = [5, 1] data = 
[[12,   24,   36,   48,   60], 
 [24,   48,   72,   96,   120]]
  ----- exit status 0

Function parameters are still matrices, so the function reshapes them. The shapes are checked once, when the function is defined, and the compiled function works for any values. The output is 12*(i+1)*(k+1) at row i, column k (here b repeats 1..5 along its last axis and a is all ones in row 0, all twos in row 1, and each output sums 12 products), which you can verify by hand.

Mistakes the compiler reports¶

################ mgc run errors/e1_axis_size_mismatch.mg
  | # The appendix's trap. A is 3x2 and B is 4x5; contracting A's axis 1 (size 2) with B's axis 0 (size 4) is impossible. The compiler says so before any code exists.
  | let a = reshape([[1, 2, 3, 4, 5, 6]], 3, 2)
  | let b = reshape([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20]], 4, 5)
  | print contract(a, (1), b, (0))
  ----- output
errors/e1_axis_size_mismatch.mg:4: error: contract: axis 1 of the first tensor has size 2 but axis 0 of the second has size 4
  ----- exit status 1

################ mgc run errors/e2_lists_differ_in_length.mg
  | # One axis of A against two of B: the lists must pair up.
  | let a = [[1, 2], [3, 4]]
  | print contract(a, (1), a, (0, 1))
  ----- output
errors/e2_lists_differ_in_length.mg:3: error: contract: the two axis lists must have the same length
  ----- exit status 1

################ mgc run errors/e3_axis_out_of_range.mg
  | # A matrix has axes 0 and 1; there is no axis 5.
  | let a = [[1, 2], [3, 4]]
  | print contract(a, (5), a, (0))
  ----- output
errors/e3_axis_out_of_range.mg:3: error: contract: axis 5 is out of range for the first tensor (shape 2x2)
  ----- exit status 1

################ mgc run errors/e4_axis_listed_twice.mg
  | # An axis can be summed once, so listing it twice is a mistake.
  | let a = [[1, 2], [3, 4]]
  | print contract(a, (1, 1), a, (0, 1))
  ----- output
errors/e4_axis_listed_twice.mg:3: error: contract: axis 1 is listed twice for the first tensor
  ----- exit status 1

################ mgc run errors/e5_reshape_element_count.mg
  | # reshape keeps the numbers, so the counts must match: 6 numbers cannot fill a 4x2 shape.
  | let a = [[1, 2, 3], [4, 5, 6]]
  | print reshape(a, 4, 2)
  ----- output
errors/e5_reshape_element_count.mg:3: error: cannot reshape 2x3 (6 elements) to 4x2 (8 elements)
  ----- exit status 1

################ mgc run errors/e6_not_a_permutation.mg
  | # permute needs every axis exactly once: a rank 3 tensor has axes 0, 1, 2 and this list repeats one and omits another.
  | let a = reshape([[1, 2, 3, 4, 5, 6, 7, 8]], 2, 2, 2)
  | print permute(a, 0, 0, 2)
  ----- output
errors/e6_not_a_permutation.mg:3: error: permute: a rank-3 tensor needs each of the axes 0..2 listed exactly once, got 0, 0, 2
  ----- exit status 1

################ mgc run errors/e7_arithmetic_on_rank3.mg
  | # Arithmetic is defined on matrices. A rank 3 value must be reshaped (or contracted) first.
  | let a = reshape([[1, 2, 3, 4, 5, 6, 7, 8]], 2, 2, 2)
  | print a + a
  ----- output
errors/e7_arithmetic_on_rank3.mg:3: error: '+' works on matrices (rank 2), but this value has shape 2x2x2 (rank 3); use reshape, permute or contract
  ----- exit status 1

################ mgc run errors/e8_dynamic_shape.mg
  | # reshape needs every size known now. batch has a '?' row count, so the compiler cannot tell whether it has the right number of elements.
  | def flatten(batch: tensor[?x4]) = reshape(batch, 2, 2, 2)
  | print flatten([[1, 2, 3, 4]])
  ----- output
errors/e8_dynamic_shape.mg:2: error: reshape needs a static shape; this value has a '?' dimension (?x4)
  ----- exit status 1

In order: a contracted pair of different sizes (the appendix's trap, now caught when the program is compiled, before any code exists); axis lists of different lengths; an axis number that does not exist; an axis listed twice; a reshape whose element counts differ; a permutation that repeats one axis and omits another; arithmetic on a rank 3 value (it must be reshaped or contracted first); and reshape of a value with a ? size (the compiler cannot know whether the counts match, so it refuses; a ? dimension is only ever checked at run time, Chapter 14). Each exits with status 1 and gives the file and line.

Compare with Chapter 27's treatment of the same trap. There, with the check on, C++ threw an exception; with the check skipped, Mountain Goat's run-time inner-dimension check aborted the program. Here, because every size in a contraction is a number in the source, the mismatch is found before the program exists. The "unchecked loop that quietly reads half of one input" the appendix warns about cannot be written in this language at all.

How it works¶

Three layers, the same three as every earlier operation.

The dialect (MgOps.td, MgDialect.cpp). mg.reshape takes a static tensor and returns a static tensor with the same number of elements. mg.permute carries a permutation attribute. Each has a verifier: reshape rejects unequal element counts and dynamic shapes; permute rejects a permutation that does not list each axis once, a wrong length, and a result shape that is not the input's axes in the new order. tests/verifiers.mlir feeds each one bad MLIR directly, so a different front end could not build an invalid operation.

The lowering (LowerToAffine.cpp). The two operations are lowered very differently, and the difference is the whole performance story:

mg.reshape becomes a memref.reinterpret_cast: the same buffer with a new list of sizes and the row-major strides. No element moves. This is safe because a Mountain Goat value never changes after it is made.
mg.permute becomes a loop nest with one affine.for per axis, running over the result and storing into each result element the input element whose axis perm[k] has the loop's k-th index. This is a real copy of the whole tensor.

The front end (mgfront.py). contract is not a new operation: it is desugared into the operations above, in this order. Work out A's free axes fa and contracted axes ia (as listed), and B's ib and fb. Permute A to fa + ia, reshape it to M x K (where M is the product of A's free sizes and K the product of the contracted sizes). Permute B to ib + fb and reshape it to K x N. One mg.matmul. Reshape the M x N product to A's free sizes followed by B's. When a permutation is the identity, or a reshape does not change the shape, the front end emits nothing, so a plain matrix product is just an mg.matmul and example 3 contains no mg.permute. If every axis is contracted the result is the 1x1 matrix noted above.

Tests¶

Six new lit files in test/contraction28/ (the suite is now 111 tests):

Test	What it checks
`examples`	the eight examples print the values on this page
`errors`	the eight mistakes give the exact messages shown, with file and line
`verifiers`	five bad uses of `mg.reshape`/`mg.permute` as raw MLIR are rejected with the exact diagnostic
`ir`	example 4 turns into the four-step MLIR above; example 3 has no `mg.permute`; example 2 has no `mg.permute` and no `mg.reshape`
`against-definition`	`check_contractions.py`: thirteen shape combinations, whole numbers and quarters, exactly equal to a loop-by-loop definition written separately in Python

The thirteen combinations are the nine Chapter 27 used (matrix product, double contraction, axis in the front of A, axes in reverse order, axes in the middle of both, vector times matrix, dot product, outer product, rank 3 by rank 3 with one shared axis) plus four new ones: a 3-cycle of axes on both sides, scattered contracted axes with their order swapped, a rank-5 operand, and size-1 axes. The numbers are small whole numbers, then multiples of a quarter, which are exact in binary, so "equal" means every digit and no tolerance is involved.

-- whole numbers
  ok   matrix product                             A[3, 2] B[2, 4] -> [3, 4]
  ok   double contraction                         A[2, 3, 4] B[3, 4, 5] -> [2, 5]
  ok   contracted axis in the FRONT of A          A[2, 3, 4] B[2, 5] -> [3, 4, 5]
  ok   contracted axes in reverse order           A[2, 3, 4] B[4, 3, 5] -> [2, 5]
  ok   contracted axes in the middle of both      A[5, 2, 3] B[4, 3, 2, 6] -> [5, 4, 6]
  ok   vector times matrix                        A[4] B[4, 6] -> [6]
  ok   dot product (every axis contracted)        A[2, 3, 4] B[2, 3, 4] -> [1, 1]
  ok   outer product (nothing contracted)         A[3] B[4] -> [3, 4]
  ok   rank 3 x rank 3, one shared axis           A[3, 4, 5] B[5, 2, 3] -> [3, 4, 2, 3]
  ok   a 3-cycle of axes on both sides            A[2, 3, 4] B[3, 2, 5] -> [4, 5]
  ok   contracted axes scattered, order swapped   A[3, 2, 4, 2] B[4, 5, 3] -> [2, 2, 5]
  ok   rank 5 operand                             A[2, 2, 2, 2, 3] B[3, 2, 2] -> [2, 2, 2, 2]
  ok   size-1 axes                                A[1, 3, 1] B[3, 1] -> [1, 1, 1]
-- quarters
  ok   matrix product                             A[3, 2] B[2, 4] -> [3, 4]
  ok   double contraction                         A[2, 3, 4] B[3, 4, 5] -> [2, 5]
  ok   contracted axis in the FRONT of A          A[2, 3, 4] B[2, 5] -> [3, 4, 5]
  ok   contracted axes in reverse order           A[2, 3, 4] B[4, 3, 5] -> [2, 5]
  ok   contracted axes in the middle of both      A[5, 2, 3] B[4, 3, 2, 6] -> [5, 4, 6]
  ok   vector times matrix                        A[4] B[4, 6] -> [6]
  ok   dot product (every axis contracted)        A[2, 3, 4] B[2, 3, 4] -> [1, 1]
  ok   outer product (nothing contracted)         A[3] B[4] -> [3, 4]
  ok   rank 3 x rank 3, one shared axis           A[3, 4, 5] B[5, 2, 3] -> [3, 4, 2, 3]
  ok   a 3-cycle of axes on both sides            A[2, 3, 4] B[3, 2, 5] -> [4, 5]
  ok   contracted axes scattered, order swapped   A[3, 2, 4, 2] B[4, 5, 3] -> [2, 2, 5]
  ok   rank 5 operand                             A[2, 2, 2, 2, 3] B[3, 2, 2] -> [2, 2, 2, 2]
  ok   size-1 axes                                A[1, 3, 1] B[3, 1] -> [1, 1, 1]
all cases agree with the definition

Are the tests good enough? Break the compiler 21 ways.¶

#!/bin/sh
# READ THIS FIRST: this script breaks the Chapter 28 code on purpose, one way at a time, and re-runs the Chapter 28 tests (test/contraction28).
#   "caught by ..." (a test failed) is the EXPECTED, wanted result: it shows the tests can notice that mistake.
#   "NOT CAUGHT" would be a test gap. Part 1 edits the front end (mgfront.py) in place; part 2 rebuilds mg-opt with one bug injected into the
#   dialect or the lowering (a full rebuild each, in work/mut/N). Every edited file is restored afterwards. Output: mutation_out.txt
HERE=$(cd "$(dirname "$0")" && pwd); LIT=$HERE/../../part15/code/run_lit.sh; W=$HERE/work/mut; rm -rf $W; mkdir -p $W
FILTER=':: contraction28/'
echo "NOTE: this script deliberately breaks the compiler code. 'caught' lines are EXPECTED: they show the tests can detect the mistake."
F=$HERE/mgfront.py; cp $F $W/mgfront.orig
restore() { cp $W/mgfront.orig $F; }
trap restore EXIT
fails_of() { grep '^FAIL' | sed 's/.*:: //; s/ (.*//' | tr '\n' ' '; }
out=$(MG_TEST_TMP=$W/out sh $LIT --filter "$FILTER" 2>&1); echo "baseline (nothing broken): $(echo "$out" | grep -q '^FAIL' && echo 'UNEXPECTED FAILURE' || echo 'tests pass (expected)')"
echo; echo "== Part 1: the front end (contract desugaring and its checks)"
fe() {  # label sed-expr
  restore; sed -i "$2" $F
  if cmp -s $F $W/mgfront.orig; then echo "$1: MUTATION DID NOT APPLY"; return; fi
  f=$(MG_TEST_TMP=$W/out sh $LIT --filter "$FILTER" 2>&1 | fails_of)
  if [ -n "$f" ]; then echo "$1: caught by $f"; else echo "$1: NOT CAUGHT (a test gap)"; fi
  restore
}
fe "A is permuted contracted-axes-first (free and contracted swapped)"   's/self.permute_by(a, sa, fa + ia)/self.permute_by(a, sa, ia + fa)/'
fe "B is permuted free-axes-first (free and contracted swapped)"         's/self.permute_by(b, sb, ib + fb)/self.permute_by(b, sb, fb + ib)/'
fe "B's contracted axes taken in increasing order, not as paired"        's/self.permute_by(b, sb, ib + fb)/self.permute_by(b, sb, sorted(ib) + fb)/'
fe "the output lists B's free axes before A's"                           's/out = tuple(sa\[d\] for d in fa) + tuple(sb\[d\] for d in fb)/out = tuple(sb[d] for d in fb) + tuple(sa[d] for d in fa)/'
fe "the size check of a contracted pair is removed"                      's/if sa\[x\] != sb\[y\]:/if False:/'
fe "a repeated axis is no longer rejected"                               's/if ax in axes\[:k\]:/if False:/'
fe "an out-of-range axis is no longer rejected"                          's/if ax >= len(s):/if False:/'
fe "the equal-length check of the two axis lists is removed"             's/if len(ia) != len(ib):/if False:/'
fe "the reshape element-count check is removed (the dialect verifier is the last line of defence)" 's/if n_in != n_out:/if False:/'
fe "the permutation check is removed (the dialect verifier is the last line of defence)" 's/if sorted(perm) != list(range(len(s))):/if False:/'
fe "arithmetic on a rank 3 value is no longer rejected by the front end" 's/if len(s) != 2:$/if False:/'
fe "a dot product is not turned into a 1x1 matrix"                       's/if not out: return r, (1, 1) /if not out: return self.reshape_to(r, (m, n), ()) /'
fe "an unneeded permute is emitted (identity permutation not skipped)"   's/if list(perm) == list(range(len(s))): return v, s/pass/'
fe "an unneeded reshape is emitted (same shape not skipped)"             's/if tuple(dims) == tuple(s): return v, s/pass/'
fe "a '?' dimension is accepted by reshape and contract"                 's/if None in s: raise MgError(self.line, f"{name} needs a static shape/if False: raise MgError(self.line, f"{name} needs a static shape/'
restore

echo; echo "== Part 2: the dialect verifiers and the lowering (a full rebuild per mutation)"
n=0
try() {  # label file sed-expr
  n=$((n+1)); d=$W/$n; mkdir -p $d/build; cp -r $HERE/tree $d/tree; rm -rf $d/tree/build
  tgt=$d/tree/lib/$2; cp $tgt $d/before; sed -i "$3" $tgt
  if cmp -s $d/before $tgt; then echo "$1: MUTATION DID NOT APPLY"; return; fi
  ( cd $d/build && cmake $d/tree -DMLIR_DIR=/usr/lib/llvm-18/lib/cmake/mlir -DLLVM_DIR=/usr/lib/llvm-18/lib/cmake/llvm -DCMAKE_BUILD_TYPE=Release -DLLVM_ENABLE_ASSERTIONS=OFF >/dev/null 2>&1 && mkdir -p include/mg && make -j8 >/dev/null 2>&1 )
  [ -x $d/build/mg-opt ] || { echo "$1: BUILD FAILED"; return; }
  f=$(MG_OPT=$d/build/mg-opt MG_TEST_TMP=$d/out sh $LIT --filter "$FILTER" 2>&1 | fails_of)
  if [ -n "$f" ]; then echo "$1: caught by $f"; else echo "$1: NOT CAUGHT (a test gap)"; fi
}
try "permute reads through the INVERSE permutation"              LowerToAffine.cpp 's/for (size_t k = 0; k < perm.size(); ++k) inIdx\[perm\[k\]\] = ivs\[k\];/for (size_t k = 0; k < perm.size(); ++k) inIdx[k] = ivs[perm[k]];/'
try "permute does not reorder at all (copies as is)"             LowerToAffine.cpp 's/inIdx\[perm\[k\]\] = ivs\[k\];/inIdx[k] = ivs[k];/'
try "reshape view uses the wrong strides (scaled by the wrong dimension)" LowerToAffine.cpp 's/strides\[i\] = strides\[i + 1\] \* shape\[i + 1\];/strides[i] = strides[i + 1] * shape[i];/'
try "the reshape element-count check is removed from the verifier" MgDialect.cpp 's/if (in.getNumElements() != out.getNumElements())/if (false)/'
try "the permute duplicate-axis check is removed from the verifier" MgDialect.cpp 's/if (p < 0 || p >= in.getRank() || seen\[p\])/if (p < 0 || p >= in.getRank())/'
try "the permute result-shape check is removed from the verifier"   MgDialect.cpp 's/if (out.getDimSize(i) != in.getDimSize(perm\[i\]))/if (false)/'
restore
echo; echo "restored: $(cmp -s $F $W/mgfront.orig && echo 'mgfront.py is back to its original' || echo 'RESTORE FAILED')"

NOTE: this script deliberately breaks the compiler code. 'caught' lines are EXPECTED: they show the tests can detect the mistake.
baseline (nothing broken): tests pass (expected)

== Part 1: the front end (contract desugaring and its checks)
A is permuted contracted-axes-first (free and contracted swapped): caught by contraction28/ir.mlir contraction28/examples.mlir contraction28/against-definition.mlir 
B is permuted free-axes-first (free and contracted swapped): caught by contraction28/ir.mlir contraction28/examples.mlir contraction28/against-definition.mlir 
B's contracted axes taken in increasing order, not as paired: caught by contraction28/against-definition.mlir 
the output lists B's free axes before A's: caught by contraction28/ir.mlir contraction28/examples.mlir contraction28/against-definition.mlir 
the size check of a contracted pair is removed: caught by contraction28/errors.mlir 
a repeated axis is no longer rejected: caught by contraction28/errors.mlir 
an out-of-range axis is no longer rejected: caught by contraction28/errors.mlir 
the equal-length check of the two axis lists is removed: caught by contraction28/errors.mlir 
the reshape element-count check is removed (the dialect verifier is the last line of defence): caught by contraction28/errors.mlir 
the permutation check is removed (the dialect verifier is the last line of defence): caught by contraction28/errors.mlir 
arithmetic on a rank 3 value is no longer rejected by the front end: caught by contraction28/errors.mlir 
a dot product is not turned into a 1x1 matrix: caught by contraction28/examples.mlir contraction28/against-definition.mlir 
an unneeded permute is emitted (identity permutation not skipped): caught by contraction28/ir.mlir 
an unneeded reshape is emitted (same shape not skipped): caught by contraction28/ir.mlir 
a '?' dimension is accepted by reshape and contract: caught by contraction28/errors.mlir 

== Part 2: the dialect verifiers and the lowering (a full rebuild per mutation)
permute reads through the INVERSE permutation: caught by contraction28/examples.mlir contraction28/against-definition.mlir 
permute does not reorder at all (copies as is): caught by contraction28/examples.mlir contraction28/against-definition.mlir 
reshape view uses the wrong strides (scaled by the wrong dimension): caught by contraction28/examples.mlir contraction28/against-definition.mlir 
the reshape element-count check is removed from the verifier: caught by contraction28/verifiers.mlir 
the permute duplicate-axis check is removed from the verifier: caught by contraction28/verifiers.mlir 
the permute result-shape check is removed from the verifier: caught by contraction28/verifiers.mlir 

restored: mgfront.py is back to its original

All 21 are caught. Four things are worth noticing.

The 3-cycle matters. "permute reads through the INVERSE permutation" is caught only because examples and cases use permutations that are not their own inverse (the 3-cycle in example 1, and the rotated axes in the reference cases). A swap-only test would have let it through.
One mutation is caught by only one test. "B's contracted axes taken in increasing order, not as paired" changes nothing in any of the eight examples (their axes are already in increasing order for B) and is caught only by against-definition, whose cases include axes listed in reverse. That is the same lesson as Chapter 27's, and the reason the comparison with the definition exists.
Two mutations change no output, and are still caught. Emitting an identity permute or a same-shape reshape copies or views data for nothing, so no value test could notice. The ir test asserts that example 3 contains no mg.permute and that example 2 contains neither operation, which turns "no data moves when none needs to" into something that can fail.
The front end and the verifiers guard the same mistakes. With the front end's reshape or permutation check deleted, the errors test still fails (it expects the front end's message), while verifiers shows the dialect rejects the bad operation by itself. One of the first runs of this script contained a mutation that did not compile ("BUILD FAILED"), which the script reports rather than hides; it was rewritten and the whole script re-run, and the output above is from that complete run.

Limits and what is not established¶

Static shapes only. reshape, permute and contract need every size known when compiling; a ? size is rejected. The tensors of rank other than 2 cannot be function parameters or results, only intermediate values (a function reshapes its matrix arguments, and a result is a matrix).
No arithmetic on rank 3 and up. Elementwise +, *, relu, reductions and @ still need matrices. A rank 3 value must be reshaped, or contracted, first. Adding rank-N elementwise operations is straightforward and not done.
No scalars of rank 0. A full contraction (a dot product) returns a 1x1 matrix. A 1x1 matrix is the language's scalar-shaped value; it is not the compile-time number a literal like 2 is.
No batch axes. As in Chapter 27, a contraction pairs axes and sums them away. An axis present in both inputs and the output (batched matrix multiply, einsum) is not supported. Chapter 29's attention needs batches of one head and avoids the question.
The permutation copies. permute is a real copy of the whole tensor, and a contraction like example 4 pays for it. No timing was done here, the copy was not fused into the product, and no claim is made about the cost on large tensors. The reshape is free.
Only small shapes were run, all f64, the largest a few hundred elements (the rank-5 case has 48 elements in its first operand).
Exactness depends on the order of additions, which the desugaring keeps the same as the definition's (the flattened contracted index runs through the pairs in the order listed). That is why the equality is exact for these cases. It is an argument plus 26 passing runs, not a proof for every shape, and a compiler setting that reorders floating-point additions would break it. The test data are exact in binary, so it would not show that.
The language's contract is not einsum. No repeated-subscript notation, no diagonals (an axis listed twice is rejected), no trace.
The appendix's own program was not run (see Chapter 27). The agreement shown is with the appendix's printed values and with the definition written for these chapters.
GPU: not covered; the GPU path is compile-only (Chapters 10 to 12) and mg.reshape/mg.permute have no GPU lowering.
Platform: x86-64 Linux, LLVM 18.1.3.

Reproducing¶

cd docs/part28/code && ./build.sh
./run_examples.sh > examples_out.txt
./check_contractions.py > check_contractions_out.txt
./mutation.sh > mutation_out.txt
cd ../../part15/code && ./run_lit.sh          # 111 tests

Chapter summary¶

Mountain Goat now has tensors of any rank, with reshape, permute and contract built in. A contraction is written in the language, with no C++ around it.
reshape is free (a new view of the same buffer); permute copies; contract is a compile-time rewrite into permute, reshape, one matrix product, reshape.
It agrees exactly with the definition of a contraction on thirteen shape combinations and with the appendix's worked values.
Mistakes the appendix had to check at run time (a mismatched contracted pair, a repeated axis, an out-of-range axis) are compile errors here.
The tests are shown able to fail: 21 injected bugs in the front end, the verifiers and the lowering.

Self-check questions¶

Each answer is collapsed; try the question first.

What is the shape of contract(a, (0, 2), b, (1, 0)) for a of shape [3, 5, 4] and b of shape [4, 3, 6], and which matrices are multiplied?

Answer

A's axis 0 (size 3) pairs with B's axis 1 (size 3), and A's axis 2 (size 4) with B's axis 0 (size 4). A's free axis is axis 1 (size 5); B's is axis 2 (size 6). The output is [5, 6]. A is permuted to [5, 3, 4] (free axis first, then the contracted axes in the order listed: 0 then 2, i.e. sizes 3 and 4) and read as 5x12; B is permuted to [3, 4, 6] (axis 1, then axis 0, then the free axis 2) and read as 12x6. The product is 5x6.
Why does example 2's contraction compile to a single mg.matmul, while example 4's needs an mg.permute?

Answer

In example 2, A's free axis is already first and its contracted axis last, and B's contracted axis is already first, so the row-major data of each is already a matrix and nothing needs to move. In example 4, A's contracted axis is first, so its numbers are not in the order a matrix product needs (free axes bundled as rows, contracted as columns), and the data must be copied into that order.
Which of reshape and permute copies data, and why can the other not?

Answer

permute copies. After permuting, the same logical tensor must be laid out in memory with a different axis order, and this chapter's matrix product reads its operands row-major, so the numbers have to be moved. reshape only changes how a flat, row-major buffer is cut into axes, which does not change where any number is stored, so it is a new view of the same memory.
Why must a 3-cycle (not a swap) be among the permutation tests?

Answer

A swap is its own inverse: applying it forwards or backwards gives the same tensor. A lowering that reads the permutation the wrong way round would therefore pass every swap-only test. A 3-cycle is not its own inverse, so only it exposes that mistake. The mutation "permute reads through the inverse permutation" is caught only because of this.
Why is the mismatch of the appendix's trap a compile error here but a run-time abort in Chapter 27?

Answer

Chapter 27's C++ built the matrices from sizes known only when the C++ ran, so the sizes reached the compiled matrix product as run-time numbers and only its run-time check could stop them. Here the sizes are literal numbers in the source, so contract can compare them while compiling. A ? size could not be checked at compile time; that is why reshape and contract refuse values with a ? dimension.
Why does the dot product return [[14]] (a 1x1 matrix) rather than 14?

Answer

The language has no rank-0 tensors: a number such as 14 in the source is a compile-time scalar, folded away before code is generated, while a computed value must be a tensor with a shape. The dot product is computed at run time, so it is a 1x1 matrix.
What does reshape(x, 2, 3, 4) do when x is 4x6, and what is printed?

Answer

It views the 24 numbers of x in row-major order as two blocks of three rows of four. Nothing is copied. For x holding 1 to 24 row by row, the first block is rows 1 2 3 4, 5 6 7 8, 9 10 11 12. (4x6 and 2x3x4 both have 24 elements; reshape(x, 4, 2) would be an error: 6 and 8 elements differ.)
The mutation "an unneeded permute is emitted (identity permutation not skipped)" does not change any program's output. Why is it still worth catching?

Answer

An identity permute is a full copy of a tensor for nothing: it would slow every contraction down without changing its answer, and no test of values could ever notice. The ir test asserts that example 3 has no mg.permute at all, which is how a performance property (no data moves when none needs to) is turned into something a test can fail on.
Why can tensors of rank other than 2 not be function parameters yet, and what is the workaround?

Answer

The type syntax tensor[RxC], the C++ header generator (mgc lib, which wraps a function's arguments as mg::Matrix) and the dynamic-shape machinery were all built for rank 2, and extending them was not part of this chapter's scope. The workaround is the one example 8 uses: pass a matrix (a 2x12 for a 2x3x4 tensor) and reshape it inside the function. The reshape is free, and the compiled function works for any values of that shape.
What would be needed to add a rank-N +?

Answer

An elementwise operation already exists in the dialect for any equal ranks (mg.add's verifier compares rank and each dimension), but its lowering is a two-loop nest and the front end rejects non-matrices. Needed: a lowering that builds a loop nest of the operand's rank (as mg.permute's already does), front-end shape rules for ranks above 2 (equal shapes; whether size-1 axes stretch), and tests with a definition-style reference as for contract.