Dependencies

plugin core;
plugin math;
plugin cps; // for %cps.cps2ds
plugin affine;
plugin buffer;

Types

Matrices are n-dimensional dependent arrays, represented by the shared %buffer.Buf type. The element operations (%buffer.alloc / %buffer.read / %buffer.write / %buffer.constant / %buffer.shape) and their lowering to %mem.Ptr (%buffer.lower_ptr) live in the buffer plugin.

High-level matrix operations

%matrix.prod

Matrix product of an m × k and a k × l matrix.

axm %matrix.prod: [m k l: Nat, [p e: Nat]]
                → [%mem.M 0, %buffer.Buf (2, (m, k), %math.F (p, e)), %buffer.Buf (2, (k, l), %math.F (p, e))]
                → [%mem.M 0, %buffer.Buf (2, (m, l), %math.F (p, e))], normalize_prod;

%matrix.transpose

Transposes a k × l matrix.

axm %matrix.transpose: [[k l: Nat], T: *]
                     → [%mem.M 0, %buffer.Buf (2, (k, l), T)]
                     → [%mem.M 0, %buffer.Buf (2, (l, k), T)], normalize_transpose;

%matrix.sum

Sums up all elements of a matrix and returns a scalar.

axm %matrix.sum: [n: Nat, S: «n; Nat», [p e: Nat]]
               → [%mem.M 0, %buffer.Buf (n, S, %math.F (p, e))]
               → [%mem.M 0, %math.F (p, e)];

%matrix.map_reduce

A general iteration/reduction scheme over tensors, inspired by and generalizing einsum. It maps a combination function comb over the elements addressed by a set of index tuples and reduces the results with an accumulator starting from zero; e.g. einsum(idx, inputs) = map_reduce(0, +, product, inputs). %matrix.prod, %matrix.transpose, and %matrix.sum are all expressed via map_reduce (see the unfolding functions below). The counts, dimensions, and indices are assumed to be statically known, to avoid inline meta-programming.

axm %matrix.map_reduce:
    [ n: Nat, S: «n; Nat», T: *, // output shape
      m: Nat,                    // number of inputs
      NI: «m; Nat»,              // input dimension counts
      TI: «m; *»,                // input element types
      SI: «i: m; «NI#i; Nat»»    // input shapes
    ]
  → [ mem:   %mem.M 0,
      zero:  T,                                              // initial accumulator
      comb:  Fn [%mem.M 0, T, «i: m; TI#i»] → [%mem.M 0, T], // inner combination
      input: «i: m; [«NI#i; Nat», %buffer.Buf (NI#i, SI#i, TI#i)]»
    ]
  → [%mem.M 0, %buffer.Buf (n, S, T)], normalize_map_reduce;

%matrix.map_reduce_aff

The buffer-world counterpart of %tensor.map_reduce: the affine-indexed generalisation of map_reduce on %buffer.Buf operands, threading %mem.M. So/Sr are the (logical) output shape and full loop bounds; acc_out/accs are affine maps from the Ro+Rr loop vector to the output/input coordinates. The combiner is mem-threaded, like %matrix.map_reduce.

axm %matrix.map_reduce_aff: {nis: Nat}
                          → {To: *, Ro Rr: Nat}
                          → [So: «Ro; Nat», Sr: «%core.nat.add (Ro, Rr); Nat»]
                          → {Tis: «nis; *», Ris: «i: nis; Nat», Sis: «i: nis; «Ris#i; Nat»»}
                          → [f: Fn [%mem.M 0, To, «i: nis; Tis#i»] → [%mem.M 0, To], init: To]
                          → [acc_out: [«%core.nat.add (Ro, Rr); %affine.index»] → «Ro; %affine.index»]
                          → [accs: «i: nis; ([«%core.nat.add (Ro, Rr); %affine.index»] → «Ris#i; %affine.index»)»]
                          → [%mem.M 0, is: «i: nis; %buffer.Buf (Ris#i, Sis#i, Tis#i)»]
                          → [%mem.M 0, %buffer.Buf (Ro, So, To)];

%matrix.broadcast

The buffer-world counterpart of %tensor.broadcast: expand input from logical shape s_in to s_out (each dimension is either equal or broadcast from size 1). The input/output buffer shapes (ri/si, ro/so) are the actual — possibly size-1-folded — buffer shapes and are inferred from the operands; s_in/s_out are the logical (unfolded) shapes.

axm %matrix.broadcast: {T: *, ri: Nat, si: «ri; Nat», ro: Nat, so: «ro; Nat», r: Nat}
                     → [s_in: «r; Nat», s_out: «r; Nat»]
                     → [%mem.M 0, %buffer.Buf (ri, si, T)]
                     → [%mem.M 0, %buffer.Buf (ro, so, T)];

%matrix.pad

The buffer-world counterpart of %tensor.pad: out[…o…] = input[…o−lo…] inside the interior, value (mode 0, constant) or the clamped edge element (mode ≠ 0, replicate) outside. s_out#d = lo#d + s_in#d + hi#d is passed explicitly. The buffer shapes are stated as the logical shapes; %buffer.Buf normalizes literal size-1 axes away, so they agree with the folded buffer handles, and the loop generation folds its read/write indices accordingly.

axm %matrix.pad: {T: *, r: Nat}
              → [s_in s_out: «r; Nat», mode: Nat, lo hi: «r; Nat»]
              → [%mem.M 0, input: %buffer.Buf (r, s_in, T), value: T]
              → [%mem.M 0, %buffer.Buf (r, s_out, T)];

%matrix.concat

The buffer-world counterpart of %tensor.concat: joins nis buffers along axis ax. out[…o…] = is#k[… o with o#ax ↦ o#ax − off#k …], where k is the input whose range o#ax lands in (off are the prefix sums of the per-input extents along ax). s_out (the summed extent along ax, the shared extents elsewhere) is passed explicitly.

axm %matrix.concat: {T: *, nis r: Nat}
                  → [ax: Idx r]
                  → {Sis: «i: nis; «r; Nat»»}
                  → [s_out: «r; Nat»]
                  → [%mem.M 0, is: «i: nis; %buffer.Buf (r, Sis#i, T)»]
                  → [%mem.M 0, %buffer.Buf (r, s_out, T)];

Unfolding functions

These express the high-level operations above in terms of %matrix.map_reduce.

%matrix.map_reduce_prod

Follows the principle ij ← ik, kj (out[i, j] = Σ_k in1[i, k] · in2[k, j]) using multiplication as combination and addition as reduction.

fun %matrix.map_reduce_prod (m k l: Nat, pe: «2; Nat») (mem: %mem.M 0, M: %buffer.Buf (2, (m, k), %math.F pe), N: %buffer.Buf (2, (k, l), %math.F pe))@tt: [%mem.M 0, %buffer.Buf (2, (m, l), %math.F pe)] =
    let R = %math.F pe;
    let zero_real = %math.conv.f2f pe 0.0:%math.F64;
    return (%matrix.map_reduce
        (2, (m, l), R, 2, (2, 2), (R, R), ((m, k), (k, l)))
        (mem, zero_real,
            fn (mem: %mem.M 0, acc: R, ab: «2; R»): [%mem.M 0, R] =
                return (mem, %math.arith.add 0 (acc, %math.arith.mul 0 ab)),
            (((0, 2), M), ((2, 1), N))));

%matrix.map_reduce_transpose

Transpose a matrix by iterating the indices in swapped order.

fun %matrix.map_reduce_transpose ((k l: Nat), T: *) (mem: %mem.M 0, M: %buffer.Buf (2, (k, l), T))@tt: [%mem.M 0, %buffer.Buf (2, (l, k), T)] =
    let zero = ⊥:T;
    return (%matrix.map_reduce
        (2, (l, k), T, 1, 2, T, (k, l))
        (mem, zero,
            // We ignore the (zero) accumulator and just return the read value.
            fn (mem: %mem.M 0, acc a: T): [%mem.M 0, T] = return (mem, a),
            ((1, 0), M)));

%matrix.map_reduce_sum

Sums up all elements of a matrix and returns a scalar.

fun %matrix.map_reduce_sum(n: Nat, S: «n; Nat», pe: «2; Nat») (mem: %mem.M 0, M: %buffer.Buf (n, S, %math.F pe))@tt: [%mem.M 0, %math.F pe] =
    let R = %math.F pe;
    let zero_real = %math.conv.f2f pe 0.0:%math.F64;
    let idxs = ‹i: n; %core.nat.add (1, %core.bitcast Nat i)›;
    let (mem, res) = %matrix.map_reduce
        (1, 1, R, 1, n, R, S)
        (mem, zero_real,
            fn (mem: %mem.M 0, acc: R, a: R): [%mem.M 0, R] =
                return (mem, %math.arith.add 0 (acc, a)),
            (idxs, M));
    return (mem, %core.bitcast R res);

Phases

%matrix.lower_matrix_high_level_map_reduce

Rewrites the high-level matrix operations (%matrix.prod, %matrix.transpose, %matrix.sum) into %matrix.map_reduce via the map_reduce_* unfolding functions above.

axm %matrix.lower_matrix_high_level_map_reduce: %compile.Phase;

%matrix.lower_matrix_medium_level

Lowers %matrix.map_reduce into %affine.For loops, making the iteration scheme explicit.

axm %matrix.lower_matrix_medium_level: %compile.Phase;

%matrix.lower_aff

Lowers the buffer-world ops (%matrix.map_reduce_aff, %matrix.broadcast, %matrix.pad, %matrix.concat) to %affine.For loops over %buffer element operations.

axm %matrix.lower_aff: %compile.Phase;

Table of Contents