The tensor Plugin

See also

-X tensor:<arg>

mim::plug::tensor

SSA-world tensor operations: a generic map_reduce loop-nest abstraction plus derived operations (reshape, slice, convolution, pooling, dot product, ...), lowered via affine loops and bufferized onto the buffer layer.

Dependencies

plugin tuple;
plugin refly;
plugin core;
import compile;
plugin option;
plugin vec;
plugin cps;
plugin affine;
plugin buffer;
plugin btensor;
 
use core.ops.n;

Types

Ring

Represents an algebraic ring. E.g., let nat_ring = (Nat, 0, core.nat.add, core.nat.mul);.

An alias for btensor.Ring. The buffer-world btensor plugin sits below tensor — lower_to_mem bufferizes onto btensor.* ops - so the definition lives there and both layers share one notion of a ring.

anx Ring = btensor.Ring;

Constructors

Build a tensor from nothing but its shape.

generate

Constructs a tensor by evaluating body once for every logical output coordinate. Coordinates use I64 so the body can do ordinary scalar arithmetic before converting them to bounded tensor indices.

output = empty(s_out)
for index in ndindex(s_out):
    output[index] = body(index)
return output

axm generate: {T: *, r: Nat}
            → [s_out: «r; Nat»]
            → [body: [«r; I64»] → T]
            → «s_out; T»;

splat

Constructs a tensor whose every element is value. Keeping this explicit lets bufferization pick a fill operation instead of first building an index lambda or a monolithic nested-array literal.

output = empty(s_out)
for index in ndindex(s_out):
    output[index] = value
return output

axm splat: {T: *, r: Nat}
         → [s_out: «r; Nat», value: T]
         → «s_out; T»;

Accessors

Element access and shape reflection.

get

Extracts the element at index from arr: get ((a, b, c), arr) ≡ arr#a#b#c.

index precedes arr so that {r, s} can be inferred: index: «i: r; Idx (s#i)» pins the rank and every extent, after which arr: «s; T» peels off exactly that many levels of its nested array type. The other way round arr would be checked while r is still unknown and the peeling depth would be ambiguous.

axm get: {T: *, r: Nat, s: «r; Nat»} → [index: «i: r; Idx (s#i)», arr: «s; T»] → T, normalize_get;

set

Inserts the element x in arr at index: set ((a, b, c), arr, x) ≡ Insert arr a (Insert (arr#a) b (Insert (arr#a#b) c x)).

index precedes arr for the same reason as in get.

axm set: {T: *, r: Nat, s: «r; Nat»} → [index: «i: r; Idx (s#i)», arr: «s; T», x: T] → «s; T», normalize_set;

shape

Returns the shape of arr as a rank-r array of Nats: (s#0, …, s#(r-1)).

The rank r is explicit: normalize_shape peels r levels off the nested array type, which inference cannot do backwards since «s; T» does not reverse-unify. Size-1 axes collapse out of the type, so arr must have no literal-1 dims among the first r.

axm shape: {A: *} → [r: Nat] → [arr: A] → «r; Nat», normalize_shape;

Map-Reduce

The loop-nest core every other operation is lowered onto: lower_tensor re-applies the matching annex from Implementations, while broadcast, concat, gather, pad and scatter are expanded by lower_map_reduce directly.

id

The CPS identity (with an empty epilogue-input pack) — the neutral post epilogue.

anx fun id {T: *} (x: T, _: [])@tt: T = return x;

map_reduce_post

The workhorse behind every tensor operation: a reduction fold over an affinely-indexed loop nest. The loop vector (o…, r…) has length Rn; its leading Ro dims are the parallel output loops, the trailing ones the reduction loops that are folded away.

• Sr, So: the full loop bounds and the result shape (rank Ro). So may differ from the leading Ro bounds of Sr — that is what lets map_out write to a re-shaped output, e.g. a transpose.
• f: folds one element of each input into the accumulator, seeded per output cell with init.
• post: maps the folded accumulator, plus one element of each of the nps epilogue inputs post_is, to the stored element. It runs once per output cell, between fold and write-back — the slot fuse_tensor composes trailing elementwise maps into (bias + activation after a convolution or matrix product).
• map_out: full loop vector ↦ the Ro write coordinates in «So». The reduction is folded away before the write, so it may only depend on the leading Ro indices.
• maps#i: full loop vector ↦ the Ris#i read coordinates of input i. Transposes, slices and broadcasts are expressed here.
• post_maps#j: the Ro output-cell coordinates (the write coordinates map_out produces) ↦ the read coordinates of epilogue input j; epilogue inputs are read once per output cell, not per loop iteration.
• sched: the op's schedule, selected in the frontend via the Tileable machinery below and carried as an ordinary operand through fusion — a target-agnostic chooser (canonically a mk_sched value) that, handed the target's loop-nest builder, returns the nest to run. The lowering only binds it and applies the result to the op's decision-free fold step and write-back, so unrolling, interchange and vectorization are plain IR, not lowering behavior.

The total loop count Rn is stated instead of the reduction count Rr = Rn - Ro: a core.nat.add (Ro, Rr) in the domain cannot be inverted by type inference and every call site would have to pass the whole meta group explicitly with @, whereas Sr pins Rn structurally.

nis, nps, Ro, Rn, Ris and Rps must be literals before the tensor::Lower phase; the bounds So and Sr may stay symbolic.

axm map_reduce_post: {nis nps: Nat}
                   → {To Tp: *, Ro Rn: Nat, TSched: □}
                   → [So: «Ro; Nat», Sr: «Rn; Nat», sched: TSched]
                   → {Tis: «nis; *», Ris: «i: nis; Nat», Sis: «i: nis; «Ris#i; Nat»»,
                      Tps: «nps; *», Rps: «j: nps; Nat», Sps: «j: nps; «Rps#j; Nat»»}
                   → [f: Fn [To, «i: nis; Tis#i»] → To, init: To, post: Fn [To, «j: nps; Tps#j»] → Tp]
                   → [map_out: «Rn; affine.Index» → «Ro; affine.Index»]
                   → [     maps: «i: nis; «Rn; affine.Index» → «Ris#i; affine.Index»»,
                      post_maps: «j: nps; «Ro; affine.Index» → «Rps#j; affine.Index»»]
                   → [is: «i: nis; «Sis#i; Tis#i»», post_is: «j: nps; «Sps#j; Tps#j»»]
                   → «So; Tp»;

SchedT

The type of a schedule value: a chooser over the target's loop-nest builder — for whatever nest type N the lowering instantiates, it is handed the target's builder mk and returns the nest to run.

let SchedT = [N: *] → ([vdim unroll: Nat] → N) → N;

mk_sched

The canonical schedule, picking

• vdim, the dim to run innermost as the unit-stride vector loop — a parallel dim runs over a row of accumulators, the innermost reduction dim selects the reduction-vectorized nest (see btensor.mr_nest), and vdim ≥ Rn means "none";
• unroll, the number of trailing reduction dims to unroll into straight-line code.

An alias for btensor.mk_sched, which btensor owns for the same reason it owns btensor.Ring.

anx mk_sched = btensor.mk_sched;

map_reduce

map_reduce_post without the epilogue and with the neutral schedule: delegates with post = id (so Tp = To), no epilogue inputs (nps = 0), no vector dim (vdim = Rn, the out-of-range sentinel), and no unrolling. Rn is stated instead of Rr for the same reason as there.

anx lam map_reduce {nis: Nat} {To: *, Ro Rn: Nat}
                   (So: «Ro; Nat», Sr: «Rn; Nat»)
                   {Tis: «nis; *», Ris: «i: nis; Nat», Sis: «i: nis; «Ris#i; Nat»»}
                   (f: Fn [To, «i: nis; Tis#i»] → To, init: To)
                   (map_out: [«Rn; affine.Index»] → «Ro; affine.Index»)
                   (maps: «i: nis; ([«Rn; affine.Index»] → «Ris#i; affine.Index»)»)
                   (is: «i: nis; «Sis#i; Tis#i»»): «So; To» =
    map_reduce_post
        (So, Sr, mk_sched (Rn, 0))
        (f, init, id @To) map_out (maps, ()) (is, ());

proj_map

subs ↦ λ o. ‹j; o#(subs#j)›: the access map that reads input axis j at loop variable subs#j, e.g. ((0, 2), (2, 1)) for the two inputs of a matrix product. An alias for btensor.proj_map, which btensor owns for the same reason it owns btensor.Ring.

anx proj_map = btensor.proj_map;

Shape Operations

Pure re-indexing: every element of the result is one element of an input, so the operation lives entirely in the access map.

broadcast

Expands the dimensions of input to fit s_out: for all i, either s_in#i = s_out#i or s_in#i = 1, and the size-1 dimensions are expanded to fit s_out.

Lowered directly to packs by the LowerMapReduce phase (lower_broadcast): a size-1 axis broadcast to n becomes a ‹n; …› pack rather than a materialized/looped copy, which map_reduce cannot express.

axm broadcast: {T: *, r: Nat}
             → [s_in s_out: «r; Nat», input: «s_in; T»]
             → «s_out; T», normalize_broadcast;

broadcast_in_dim

Transposes and expands the dimensions of input to fit s_out. Each input dimension is mapped to the output dimension given by index; the holes are filled by broadcasting.

Todo

We could probably implement this in terms of broadcast and transpose directly.

axm broadcast_in_dim: {T: *, r_in r_out: Nat}
                    → [s_in: «r_in; Nat», s_out: «r_out; Nat», input: «s_in; T», index: «r_in; Idx r_out»]
                    → «s_out; T», normalize_broadcast_in_dim;

concat_acc

Accumulator for the axis-sum: adds row#ax to the running total.

lam concat_acc {r: Nat} (ax: Idx r) (acc: Nat, row: «r; Nat»): Nat = acc + row#ax;

concat_shape

(ax, Sis) ↦ s_out: sums the inputs' extents along ax, keeps the shared extents on every other axis.

lam concat_shape {nis r: Nat} (ax: Idx r, Sis: «nis; «r; Nat»»): «r; Nat» =
    let sum_ax = vec.fold.l (concat_acc ax) (0, Sis);
    let row_0  = Sis#(core.idx 0 nis 0);
    ‹d: r; core.select (core.icmp.e (d, ax), sum_ax, row_0#d)›;

concat

Joins nis inputs along axis ax. All inputs share rank r and agree on every axis except ax; the output extent along ax is the sum of the inputs' extents there (s_out#ax = Σ_i Sis#i#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).

The output shape is deduced from ax and the per-input shapes Sis (see concat_shape), so it is not passed explicitly.

axm concat: {T: *, nis r: Nat}
          → [ax: Idx r]
          → {Sis: «i: nis; «r; Nat»»}
          → [is: «i: nis; «Sis#i; T»»]
          → «concat_shape (ax, Sis); T», normalize_concat;

flip

Reverses every axis: out[…d…] = input[… (s_in#d − 1) − d …].

axm flip: {T: *, r: Nat}
        → [s_in: «r; Nat»]
        → [input: «s_in; T»]
        → «s_in; T», normalize_flip;

pad_shape

lam pad_shape {r: Nat} (lo s_in hi: «r; Nat»): «r; Nat» =
    ‹d: r; lo#d + s_in#d + hi#d›;

pad

Pads each axis: s_out#d = lo#d + s_in#d + hi#d, rank preserved. out[…d…] = input[… o#d − lo#d …] whenever every axis lands inside the input region; otherwise the value depends on mode:

• mode = 0 (constant): out-of-region cells take the scalar value.
• mode = 1 (replicate): out-of-region cells take the nearest edge element (the read index is clamped to [0, s_in#d − 1] per axis); value is ignored.

The output shape is deduced from lo, s_in and hi (see pad_shape), so it is not passed explicitly.

// `(lo, s_in, hi) ↦ ‹d; lo#d + s_in#d + hi#d›`: the padded shape.
 
axm pad: {T: *, r: Nat}
       → [s_in: «r; Nat»]
       → [mode: Nat, lo hi: «r; Nat»]
       → [input: «s_in; T», value: T]
       → «pad_shape (lo, s_in, hi); T», normalize_pad;

repeat

Tiles input to the larger shape s_out (each s_out#d must be a multiple of s_in#d): out[…d…] = input[… d mod s_in#d …].

axm repeat: {T: *, r: Nat}
          → [s_in: «r; Nat»]
          → [s_out: «r; Nat»]
          → [input: «s_in; T»]
          → «s_out; T», normalize_repeat;

reshape

Row-major reshape: prod(s_in) must equal prod(s_out). Loops over the result and reads the input at the delinearized linear index, so it works for arbitrary rank changes (e.g. flatten/unflatten).

s_in is explicit because s_out has a different rank and so cannot pin the input's rank/element-type during inference.

axm reshape: {T: *, r_in r_out: Nat}
           → [s_in: «r_in; Nat»]
           → [s_out: «r_out; Nat»]
           → [input: «s_in; T»]
           → «s_out; T», normalize_reshape;

slice

Strided slice / narrow: out[…d…] = input[… start#d + step#d · d …]. Same rank in/out; s_out gives the sliced extents (the caller must ensure start#d + step#d · (s_out#d − 1) < s_in#d).

axm slice: {T: *, r: Nat}
         → [s_in: «r; Nat»]
         → [start step s_out: «r; Nat»]
         → [input: «s_in; T»]
         → «s_out; T», normalize_slice;

transpose_shape

(s, permutation) ↦ ‹i; s#(permutation⁻¹#i)›: the permuted shape.

lam transpose_shape {r: Nat} (s: «r; Nat», permutation: «r; Idx r»): «r; Nat» =
  let shape_permutation = ‹i: r; option.unwrap_unsafe (vec.first core.icmp.e (permutation, i))›;
  ‹i: r; s#(shape_permutation#i)›;

transpose

Permutes the dimensions of input according to permutation.

permutation precedes input, and {s} is a separate implicit group introduced after it, so that {T, r, s} can be inferred: permutation: «r; Idx r» pins the rank r, and only then is the {s} Hole created — with the now-literal type «r; Nat», which lets input: «s; T» peel off exactly r levels of its nested array type. Merging the two groups would not do: the domain «s; T» is built once, while r is still unknown.

axm transpose: {T: *, r: Nat}
             → [permutation: «r; Idx r»]
             → {s: «r; Nat»}
             → [input: «s; T»]
             → «transpose_shape (s, permutation); T»;

transpose_2d

Permutes the dimensions of a 2-dimensional tensor.

axm transpose_2d: {T: *}
                → {s: «2; Nat»}
                → [input: «s; T»]
                → «s#tt, s#ff; T»;

Element-wise Operations

Apply a scalar function at every coordinate.

binary

Maps a binary function over a pair of tensors.

axm binary: {Ti1 Ti2 To: *}
          → [app: [Ti1, Ti2] → To]
          → {r: Nat, s: «r; Nat»}
          → [is: [«s; Ti1», «s; Ti2»]]
          → «s; To»;

map

Maps a function over a collection of tensors.

axm map: {T: *, ni: Nat, Is: «ni; *»}
       → [app: «i: ni; Is#i» → T]
       → {r: Nat, s: «r; Nat»}
       → [is: «i: ni; «s; Is#i»»]
       → «s; T»;

select

Maps core.select over tensors.

axm select: {T: *}
          → {r: Nat, s: «r; Nat»}
          → [is: [«s; Bool», «s; T», «s; T»]]
          → «s; T»;

unary

Maps a unary function over a tensor.

axm unary: {Ti To: *}
         → [app: Ti → To]
         → {r: Nat, s: «r; Nat»}
         → [i: «s; Ti»]
         → «s; To»;

Reductions

Operations that fold one or more reduction dims away.

bmm

Computes the batch matrix multiplication (BMM) of two 3-dimensional tensors — a special case of dot_product.

axm bmm: [R: Ring]
       → {B M K N: Nat}
       → [t1: «B, M, K; R#T», t2: «B, K, N; R#T»]
       → «B, M, N; R#T»;

conv_shape

(n, cout, (H, W), (KH, KW), stride, dilation, padding) ↦ (n, cout, OH, OW), where \(O_d = (s_d + 2 \cdot \mathit{padding}_d - (\mathit{dilation}_d \cdot (k_d - 1) + 1)) / \mathit{stride}_d + 1\): the NCHW output shape.

pub lam conv_shape (n cout: Nat, s k stride dilation padding: «2; Nat»): «4; Nat» =
    let hw = ‹d: 2; (s#d + 2 * padding#d - (dilation#d * (k#d - 1) + 1)) / stride#d + 1›;
    (n, cout, hw#0_2, hw#1_2);

conv

2-D convolution (cross-correlation) in the ring R, groups = 1, no bias:

\[ \mathit{out}_{n,co,oh,ow} = \sum_{ci,kh,kw} \mathit{padded}_{n,ci,ih,iw} \cdot \mathit{weight}_{co,ci,kh,kw}, \qquad ih = \mathit{stride}_0 \cdot oh + \mathit{dilation}_0 \cdot kh, \qquad iw = \mathit{stride}_1 \cdot ow + \mathit{dilation}_1 \cdot kw \]

• Layout: input: «N, Cin, H, W; R#T», weight: «Cout, Cin, KH, KW; R#T», output: «N, Cout, OH, OW; R#T».
• padding zero-pads the spatial axes (composed via pad), so the windowed read is a pure affine map.
• The output shape follows from the extents and window parameters (see conv_shape) and is not passed explicitly.

axm conv: [R: Ring]
        → {n cin cout h w kh kw: Nat}
        → [stride dilation padding: «2; Nat»]
        → [input: «n, cin, h, w; R#T», weight: «cout, cin, kh, kw; R#T»]
        → «conv_shape (n, cout, (h, w), (kh, kw), stride, dilation, padding); R#T»;

dot_general_shape

The result shape: the batching extents, then the left operand's free extents, then the right operand's; the batching and contracting extents are checked to agree.

lam dot_general_shape {r1 r2: Nat} {nc nb: Nat}
  (c1: «nc; Idx r1», c2: «nc; Idx r2», b1: «nb; Idx r1», b2: «nb; Idx r2»)
  (s1: «r1; Nat», s2: «r2; Nat») =
      let bs_check = ‹i: nb; refly.check (s1#(b1#i) == s2#(b2#i), s1#(b1#i), "batching dims don't match")›;
      let cs_check = ‹i: nc; refly.check (s1#(c1#i) == s2#(c2#i), s1#(c1#i), "contracting dims don't match")›;
      let bs       = ‹i: nb; s1#(b1#i)›;
      let bc_1     = vec.cat (b1, c1);
      let s1_res   = vec.diff (s1, bc_1);
      let bc_2     = vec.cat (b2, c2);
      let s2_res   = vec.diff (s2, bc_2);
      let s12_res  = vec.cat (s1_res, s2_res);
      let s_out    = vec.cat (bs, s12_res);
      s_out;

dot_product

Returns the generalized dot product of a and b:

• R: the ring in which the dot product is performed
• r1/r2: the ranks of the two inputs
• nc/nb: the number of contracting/batching dimensions
• c1/c2: the contracting dimensions of the left/right input
• b1/b2: the batching dimensions of the left/right input
• s1/s2: the shape of the left/right input
• a/b: the left/right input

axm dot_product: [R: Ring]
               → {r1 r2: Nat}
               → {nc nb: Nat}
               → [c1: «nc; Idx r1», c2: «nc; Idx r2», b1: «nb; Idx r1», b2: «nb; Idx r2»]
               → {s1: «r1; Nat», s2: «r2; Nat»}
               → [a: «s1; R#T», b: «s2; R#T»]
               → «dot_general_shape (c1, c2, b1, b2) (s1, s2); R#T»;

pool

2-D pooling: reduces each (kernel#0, kernel#1) window of input: «N, C, H, W; T» with the fold g seeded by init to output: «N, C, OH, OW; T», each channel independently.

• padding pads the spatial axes with init (composed via pad), so padded cells are inert.
• Max-pooling is g = max, init = −∞; sum-/average-pooling is g = +, init = 0 — divide by the window size afterwards for the average.
• Unlike conv there is no weight, so the window size kernel is an explicit operand, and the output spatial size s_out = (OH, OW) is passed explicitly (no Nat division to compute it).

axm pool: {T: *}
        → [g: [T, T] → T, init: T]
        → {n c h w: Nat}
        → [kernel stride dilation padding s_out: «2; Nat»]
        → [input: «n, c, h, w; T»]
        → «n, c, s_out#0_2, s_out#1_2; T»;

product_2d

Computes the matrix product of two 2-dimensional tensors — a special case of dot_product.

axm product_2d: [R: Ring]
              → {m k l: Nat}
              → [t1: «m, k; R#T», t2: «k, l; R#T»]
              → «m, l; R#T»;

Indexing

Data-dependent access: an index tensor picks the coordinate along dim.

gather

Reads one source element for every index element:

• the output coordinate is also the index coordinate, except on dim, where the value loaded from index selects the source coordinate;
• s_idx#d <= s_src#d is required for every d != dim;
• dim: Idx r and the index element type encode the axis and the indexed-source bounds in the type.

for out_index in ndindex(s_idx):
    selected = index[out_index]
    source_index = tuple(
        selected if d == dim else out_index[d]
        for d in range(r)
    )
    output[out_index] = input[source_index]

axm gather: {T: *, r: Nat}
          → [s_src: «r; Nat», s_idx: «r; Nat»]
          → [dim: Idx r]
          → [input: «s_src; T», index: «s_idx; Idx (s_src#dim)»]
          → «s_idx; T»;

scatter

Starts with input and visits index in row-major order. Each corresponding updates element is stored at the input coordinate selected by index on dim. Duplicate destinations use the last visited update. For d != dim, s_idx#d <= s_src#d; on every axis, s_idx#d <= s_updates#d.

output = input.copy()
for update_coord in ndindex(s_idx):
    selected = index[update_coord]
    destination = tuple(
        selected if d == dim else update_coord[d]
        for d in range(r)
    )
    output[destination] = updates[update_coord]

axm scatter: {T: *, r: Nat}
           → [s_src: «r; Nat», s_idx: «r; Nat», s_updates: «r; Nat»]
           → [dim: Idx r]
           → [input: «s_src; T», index: «s_idx; Idx (s_src#dim)», updates: «s_updates; T»]
           → «s_src; T»;

Reflection

Compile-time queries about an operand that a schedule selector (see Tiling) branches on.

fastest_axis

Layout reflection for schedule selectors: which axis of the rank-r tensor t is the fastest-varying (unit-stride) axis of the tensor actually read, once fuse_tensor's read-through has absorbed a pure re-indexed read behind t — behind a transpose the source's trailing axis maps back to a different axis of t.

• a plain tensor (anything that is not such a read, a still-symbolic operand included): its own last axis r − 1.
• an undecidable read (composed arithmetic like reshape, a nested read chain): the sentinel r.

This reflects only operand structure the frontend cannot inspect itself; the schedule decision built on the answer stays in the IR next to the op (see dot_product_impl). tensor::Lower applies it where the op's operands are concrete.

axm fastest_axis: {T: *} → [r: Nat, t: T] → Nat, normalize_fastest_axis;

if_static

Binding-time dispatch for schedule selectors, so that one selector serves compile-time and runtime extents without staging annotations at the call site:

• k a literal: the normalizer picks s.
• k symbolic: stays stuck, keeping the choice open until instantiation.
• still stuck at lowering time: the lowerings residualize to d, since an undecided k is then a runtime value.

axm if_static: {T: *} → [k: Nat, s d: T] → T, normalize_if_static;

Tiling

Tiling, strip-mining and schedule selection are implemented entirely in the frontend IR, with partial evaluation doing the rewriting; no C++ transformation code.

Tileable

A Tileable existentially packages a computation over an affinely-indexed iteration domain together with its schedule (cf. MLIR's TilingInterface):

• Out: the result type constructor — what materializing over a domain with ro parallel dims of extents so yields.
• the iteration domain Sr (leading Ro parallel dims, trailing Rr reduction dims) and the result shape So — MLIR's getIterationDomain + getLoopIteratorTypes.
• oax: which output axis each parallel dim writes — the piece of the write map the domain transforms need (strip_mine_par splits the output along oax#d).
• the schedule vdim/unroll (see mk_sched): the chooser retile packs into the op.
• retile: rebuild the computation over a reindexed domain, given the new domain and schedule, iota (maps new loop vectors to old loop vectors) and omega (maps (old output coordinates, new loop vector) to new output coordinates).

The op's internals (combiner, init, epilogue, inputs, access maps) stay hidden in the closure.

let Tileable = [
    Out: [ro: Nat, so: «ro; Nat»] → *,
    Ro: Nat, Rr: Nat,
    So: «Ro; Nat»,
    oax: «Ro; Nat»,
    Sr: «Ro + Rr; Nat»,
    vdim: Nat, unroll: Nat,
    retile: [Ro_n Rr_n: Nat]
          → [So_n: «Ro_n; Nat», Sr_n: «Ro_n + Rr_n; Nat», vdim_n unroll_n: Nat]
          → [iota:  «Ro_n + Rr_n; affine.Index» → «Ro + Rr; affine.Index»,
             omega: [«Ro; affine.Index», «Ro_n + Rr_n; affine.Index»] → «Ro_n; affine.Index»]
          → Out (Ro_n, So_n),
];

comp_map

Composition helper for access maps: (m ∘ g).

lam comp_map {A B C: *} (m: B → C, g: A → B) (o: A): C = m (g o);

interchange

Permute the loop dims — new loop j iterates old dim perm#j. Loop order and write position are decoupled by the write map, so So is untouched. The permutation must keep parallel dims parallel (perm#j < Ro iff j < Ro), enforced at compile time; the schedule's vdim moves to its new position.

anx lam interchange (t: Tileable, perm: «t#Ro + t#Rr; Nat»): Tileable =
    let r = t#Ro + t#Rr;
    let perm_ok = ‹j: r; let jj = core.bitcast Nat j;
                         refly.check (core.icmp.e @2 (jj < t#Ro, perm#j < t#Ro),
                                       perm#j,
                                       "interchange: permutation must not mix parallel and reduction dims")›;
    let sr_n   = ‹i: r; t#Sr#(core.idx 0 r (perm_ok#i))›;
    let vdim_n = match vec.first `== (perm_ok, t#vdim) with
                     | j: Idx r => core.bitcast Nat j
                     | _: []    => t#vdim;
    // The parallel dims permute among themselves (checked above), so the new dim j writes the
    // output axis its old dim perm#j wrote.
    let oax_n = ‹j: t#Ro; t#oax#(core.idx 0 (t#Ro) (perm_ok#(core.idx 0 r (core.bitcast Nat j))))›;
    // iota: old dim i is found at the new position j with perm#j = i.
    lam iota_p (o: «r; affine.Index»): «r; affine.Index» =
        ‹i: r; o#(option.unwrap_unsafe (vec.first `== (perm_ok, core.bitcast Nat i)))›;
    lam retile_n (ro_m rr_m: Nat)
                 (so_m: «ro_m; Nat», sr_m: «ro_m + rr_m; Nat», vdim_m unroll_m: Nat)
                 (iota:  «ro_m + rr_m; affine.Index» → «r; affine.Index»,
                  omega: [«t#Ro; affine.Index», «ro_m + rr_m; affine.Index»] → «ro_m; affine.Index»)
               : t#Out (ro_m, so_m) =
        t#retile (ro_m, rr_m) (so_m, sr_m, vdim_m, unroll_m) (comp_map (iota_p, iota), omega);
    (t#Out, t#Ro, t#Rr, t#So, oax_n, sr_n, vdim_n, t#unroll, retile_n);

materialize

Materialize a Tileable over its current domain and schedule (identity reindexing).

anx lam materialize (t: Tileable): t#Out (t#Ro, t#So) =
    lam id_i (o: «t#Ro + t#Rr; affine.Index»): «t#Ro + t#Rr; affine.Index» = o;
    lam id_o (oo: «t#Ro; affine.Index», o: «t#Ro + t#Rr; affine.Index»): «t#Ro; affine.Index» = oo;
    t#retile (t#Ro, t#Rr) (t#So, t#Sr, t#vdim, t#unroll) (id_i, id_o);

mr_tileable

The map_reduce_post instance, the SSA analog of MLIR's getTiledImplementation: retile rebuilds the same op over the new domain, composing the access maps with iota and the write map with omega.

• No extract/insert_slice machinery — the tiled reads hit the original inputs through the reindexed access maps.
• Epilogue-free (nps = 0): impls build raw ops and the epilogue arrives later via fuse_tensor, because epilogue-input access maps take output-cell coordinates, whose rank a domain transform may change.
• oax states which output axis map_out writes each parallel dim to — the identity for the impls' projection write maps.

anx lam mr_tileable
                {nis: Nat}
                {Ro Rr: Nat}
                {To Tp: *}
                (So: «Ro; Nat», oax: «Ro; Nat», Sr: «Ro + Rr; Nat»)
                {Tis: «nis; *», Ris: «i: nis; Nat», Sis: «i: nis; «Ris#i; Nat»»}
                (f: Fn [To, «i: nis; Tis#i»] → To, init: To, post: Fn [To, []] → Tp)
                (map_out: [«Ro + Rr; affine.Index»] → «Ro; affine.Index»)
                (maps: «i: nis; ([«Ro + Rr; affine.Index»] → «Ris#i; affine.Index»)»)
                (is: «i: nis; «Sis#i; Tis#i»»)
              : Tileable =
    lam out_t (ro: Nat, so: «ro; Nat»): * = «so; Tp»;
    lam retile (Ro_n Rr_n: Nat)
               (So_n: «Ro_n; Nat», Sr_n: «Ro_n + Rr_n; Nat», vdim_n unroll_n: Nat)
               (iota:  «Ro_n + Rr_n; affine.Index» → «Ro + Rr; affine.Index»,
                omega: [«Ro; affine.Index», «Ro_n + Rr_n; affine.Index»] → «Ro_n; affine.Index»)
             : «So_n; Tp» =
        lam map_out_n (o: «Ro_n + Rr_n; affine.Index»): «Ro_n; affine.Index» =
            omega (map_out (iota o), o);
        map_reduce_post @(nis, 0) @(To, Tp, Ro_n, Ro_n + Rr_n, SchedT)
            (So_n, Sr_n, mk_sched (vdim_n, unroll_n))
            @(Tis, Ris, Sis, (), (), ()) (f, init, post)
            map_out_n
            (‹i: nis; comp_map (maps#i, iota)›, ())
            (is, ());
    (out_t, Ro, Rr, So, oax, Sr, Ro + Rr, 0, retile);

unroll_trailing

Fully unroll the trailing d reduction dims into the loop body — e.g. a convolution's kh·kw window (d = 2). Their trip counts must be literals at lowering time.

lam unroll_trailing (t: Tileable, d: Nat): Tileable =
    (t#Out, t#Ro, t#Rr, t#So, t#oax, t#Sr, t#vdim, d, t#retile);

vectorize

Run dim d as the innermost, unit-stride vector loop (per-cell fold order is unchanged). A parallel d runs inside the reduction loops over a row of accumulators; the innermost reduction dim as d selects btensor.mr_nest's reduction-vectorized nest instead.

lam vectorize (t: Tileable, d: Nat): Tileable =
    (t#Out, t#Ro, t#Rr, t#So, t#oax, t#Sr, d, t#unroll, t#retile);

Strip-Mining

Strip-mining is a pure domain transform (cf. TilingInterface's op[N] → op[N/ts, ts]) that retile absorbs entirely: iota recombines the split loop vector, and for a parallel dim omega splits the write coordinate. The result is an ordinary map_reduce_post over a higher-rank domain with composed access maps — schedule operand and lowering untouched. Classic cache tiling is composition: strip-mine, then interchange the tile loops outward.

check_div

Checks that ts exactly divides the extent sd (discharges symbolically for factored extents).

lam check_div (sd ts: Nat): Nat =
    refly.check (sd / ts * ts == sd, sd, "strip_mine: tile size must divide the extent");

shift_ax

Axis renumbering after inserting a new axis at position a.

lam shift_ax (a x: Nat): Nat = core.select (x < a, x, x + 1);

split_bounds

Split bounds vector v (length n) at dim d into (nt, ts); the result arity m (= n+1) is passed explicitly, in the syntactic form the call site's types require (the checker does not re-associate core.nat.add symbolically).

lam split_bounds (m n: Nat, v: «n; Nat», d ts nt: Nat): «m; Nat» =
    ‹i: m; let ii = core.bitcast Nat i;
           core.select (ii < d, v#(core.idx 0 n ii),
           core.select (ii == d, nt,
           core.select (ii == d + 1, ts,
           v#(core.idx 0 n (ii - 1)))))›;

split_iota

Recombine a split loop vector (m = n+1): old#d = new#d · ts + new#(d+1).

lam split_iota (m n d ts: Nat) (o: «m; affine.Index»): «n; affine.Index» =
    ‹j: n; let jj = core.bitcast Nat j;
           core.select (jj < d, o#(core.idx 0 m jj),
           core.select (jj == d,
                         affine.op.add (affine.semiop.mul (o#(core.idx 0 m d), ts),
                                         o#(core.idx 0 m (d + 1))),
           o#(core.idx 0 m (jj + 1))))›;

split_omega

Split an output coordinate vector (m = n+1): new#a = old#a ⌊/⌋ ts, new#(a+1) = old#a mod ts.

lam split_omega (m n a ts: Nat) (oc: «n; affine.Index»): «m; affine.Index» =
    ‹j: m; let jj = core.bitcast Nat j;
           core.select (jj < a, oc#(core.idx 0 n jj),
           core.select (jj == a, affine.semiop.floordiv (oc#(core.idx 0 n a), ts),
           core.select (jj == a + 1, affine.semiop.rem (oc#(core.idx 0 n a), ts),
           oc#(core.idx 0 n (jj - 1)))))›;

strip_mine_par

Split parallel dim d (d < Ro) into two adjacent parallel dims. The result shape splits along the output axis d writes (oax#d), so the materialized tensor is packed («…, N/ts, ts, …» along that axis); a row-major reshape restores the original shape — and in a fused graph fuse_tensor's read-through absorbs it into the consumers' access maps.

anx lam strip_mine_par (t: Tileable, d ts: Nat): Tileable =
    let r    = t#Ro + t#Rr;
    let ro_n = t#Ro + 1;
    let r_n  = t#Ro + 1 + t#Rr;
    let a    = t#oax#(core.idx 0 (t#Ro) d);
    let sd   = refly.check (t#So#(core.idx 0 (t#Ro) a) == t#Sr#(core.idx 0 r d),
                             t#Sr#(core.idx 0 r d),
                             "strip_mine_par: output extent (at oax#d) must match domain extent");
    let nt   = check_div (sd, ts) / ts;
    // The split dim writes the split axis pair; all other dims write their old axis, renumbered
    // around the insertion at `a`. The unit-stride half is the inner one: a vdim at `d` moves to
    // `d + 1`.
    let oax_n = ‹j: ro_n; let jj = core.bitcast Nat j;
                          core.select (jj < d, shift_ax (a, t#oax#(core.idx 0 (t#Ro) jj)),
                          core.select (jj == d, a,
                          core.select (jj == d + 1, a + 1,
                          shift_ax (a, t#oax#(core.idx 0 (t#Ro) (jj - 1))))))›;
    let vdim_n = core.select (t#vdim < t#Ro, core.select (t#vdim < d, t#vdim, t#vdim + 1), r_n);
    lam retile_n (ro_m rr_m: Nat)
                 (so_m: «ro_m; Nat», sr_m: «ro_m + rr_m; Nat», vdim_m unroll_m: Nat)
                 (iota:  «ro_m + rr_m; affine.Index» → «r_n; affine.Index»,
                  omega: [«ro_n; affine.Index», «ro_m + rr_m; affine.Index»] → «ro_m; affine.Index»)
               : t#Out (ro_m, so_m) =
        lam iota_c  (o: «ro_m + rr_m; affine.Index»): «r; affine.Index» =
            split_iota (r_n, r, d, ts) (iota o);
        lam omega_c (oo: «t#Ro; affine.Index», ov: «ro_m + rr_m; affine.Index»): «ro_m; affine.Index» =
            omega (split_omega (ro_n, t#Ro, a, ts) oo, ov);
        t#retile (ro_m, rr_m) (so_m, sr_m, vdim_m, unroll_m) (iota_c, omega_c);
    (t#Out, ro_n, t#Rr,
     split_bounds (ro_n, t#Ro, t#So, a, ts, nt), oax_n,
     split_bounds (r_n, r, t#Sr, d, ts, nt), vdim_n, t#unroll, retile_n);

strip_mine_red

Split reduction dim d (Ro ≤ d < Ro + Rr) of extent N into two adjacent reduction dims (N/ts, ts). If the split lands inside the trailing unroll window, both halves unroll — the tap count is unchanged.

anx lam strip_mine_red (t: Tileable, d ts: Nat): Tileable =
    let r   = t#Ro + t#Rr;
    let r_n = t#Ro + (t#Rr + 1);
    let nt  = check_div (t#Sr#(core.idx 0 r d), ts) / ts;
    let vdim_n   = core.select (t#vdim < t#Ro, t#vdim, r_n);
    let unroll_n = core.select (d < r - t#unroll, t#unroll, t#unroll + 1);
    lam retile_n (ro_m rr_m: Nat)
                 (so_m: «ro_m; Nat», sr_m: «ro_m + rr_m; Nat», vdim_m unroll_m: Nat)
                 (iota:  «ro_m + rr_m; affine.Index» → «r_n; affine.Index»,
                  omega: [«t#Ro; affine.Index», «ro_m + rr_m; affine.Index»] → «ro_m; affine.Index»)
               : t#Out (ro_m, so_m) =
        t#retile (ro_m, rr_m) (so_m, sr_m, vdim_m, unroll_m)
                 (comp_map (split_iota (r_n, r, d, ts), iota), omega);
    (t#Out, t#Ro, t#Rr + 1, t#So, t#oax,
     split_bounds (r_n, r, t#Sr, d, ts, nt), vdim_n, unroll_n, retile_n);

Schedules

A Schedule is just a Tileable transformer, a selector a compile-time inspection of an op's parameters returning the appropriate one — combinator composition, dispatched as an extract over thunks so that partial evaluation resolves it entirely. This is where "when is a transformation appropriate" lives: in the IR next to the op, not in a lowering heuristic.

Schedule

A transformer of Tileables.

let Schedule = Tileable → Tileable;

sched_id

lam sched_id (t: Tileable): Tileable = t;

unroll_budget

The greedy trailing-suffix unroll scan (the old lowering's budget scan): the number of trailing reduction dims whose tap product stays within budget. Folded innermost-first over the reduction dims with state (count, taps, alive); a dim whose extent is still open at lowering time stops the scan (if_static residualizes its step to "over budget").

lam unroll_budget (t: Tileable, budget: Nat): Nat =
    let r = t#Ro + t#Rr;
    lam step (d: Nat, st: «3; Nat»): «3; Nat» =
        let taps  = st#1_3 * t#Sr#(core.idx 0 r d);
        let alive = st#2_3 * if_static (taps, (0, 1)#(taps <= budget), 0);
        (st#0_3 + alive, core.select (alive == 1, taps, st#1_3), alive);
    (vec.fold.r @(«3; Nat», Nat) @(t#Rr) step (‹i: t#Rr; t#Ro + core.bitcast Nat i›, (0, 1, 1)))#0_3;

vectorize_if

Vectorize the parallel dim d only when cond holds; otherwise the sentinel "no vector dim".

lam vectorize_if (t: Tileable, cond: Bool, d: Nat): Tileable =
    vectorize (t, (t#Ro + t#Rr, d)#cond);

conv_schedule

The convolution schedule, selected from the op's parameters — the old lowering's heuristics, now plain IR next to the op:

• unroll the trailing reduction dims greedily while their tap product fits the register budget (the kh·kw window; a small cin joins it) — constant tap offsets, no window loops;
• run ow (parallel dim 3) innermost with unit stride over a row of accumulators iff it beats the innermost remaining reduction dim on unit-stride/broadcast reads: ow reads the padded input at stride_w and is broadcast in the weight (score [stride_w = 1] + 1); an un-unrolled kw reads the input at dilation_w and the weight at unit stride ([dilation_w = 1] + 1 — the classic nest is already unit-stride, so vectorizing buys nothing); kh/cin score 0; with the whole reduction unrolled there is nothing to beat.

Open extents/strides residualize to no unrolling and no vector dim (if_static).

lam conv_schedule (stride dilation: «2; Nat») (t: Tileable): Tileable =
    let u     = unroll_budget (t, 32);
    let s_ow  = (0, 1)#(stride#1_2 == 1) + 1;
    let s_red = core.select (u == 0, (0, 1)#(dilation#1_2 == 1) + 1, core.select (u < t#Rr, 0, s_ow));
    let ow    = t#Sr#(core.idx 0 (t#Ro + t#Rr) 3);
    let vec   = if_static (stride#1_2 * dilation#1_2 * ow, s_red < s_ow, ff);
    // Rows too short for the vectorizer (ow < 16) are processed in BLOCKS: strip-mine oh by the
    // largest block size ≤ 4 that divides it and sink both halves (vdim = ro − 2 selects mr_nest's
    // blocked branch, which jams the block dim into the vector body — one reduction pass folds the
    // whole row block, sharing the hoisted window weights). Otherwise the split is the neutral
    // ts = 1 (its size-1 axis folds away), so one code path serves every case and no divisibility
    // check fires on an untaken choice.
    let oh    = t#Sr#(core.idx 0 (t#Ro + t#Rr) 2);
    lam divides (d: Nat): Bool = oh / d * d == oh;
    let bs    = ((1, 2)#(divides 2), 4)#(divides 4);
    let block = if_static (oh * ow, core.bit2.and_ 2 (vec, core.bit2.and_ 2 (ow < 16, bs != 1)), ff);
    unroll_trailing (vectorize_if (strip_mine_par (t, 2, (1, bs)#block), vec, (4, 3)#block), u);

dot_schedule

The dot-product schedule for a contraction whose right operand stores the last output dim unit-stride ( \(A \cdot B\)): that dim is invariant in the left operand, so it runs innermost over a row of accumulators (GEMM-style vectorization) and the innermost contraction is register-blocked — strip-mined by 4, the tile unrolled.

• Blocking only needs the divisibility decided, which is weaker than a literal extent: a factored extent («(n+3)/4*4, …», say a padded row) folds the guard to tt under a symbolic n, because core.nat cancels (c*x)/c.
• if_static therefore gates on the tile size rather than on K — a tile that folded to a literal is decided, no matter which kind of extent decided it.
• An undecided divisibility takes the neutral tile (ts = 1), so one selector serves static, symbolic and runtime shapes.

lam dot_schedule (r_out: Nat) (t: Tileable): Tileable =
    let r   = t#Ro + t#Rr;
    let d   = r - 1;
    let K   = t#Sr#(core.idx 0 r d);
    let ts0 = (1, 4)#(K / 4 * 4 == K);
    let ts  = if_static (ts0, ts0, 1);
    unroll_trailing (strip_mine_red (vectorize (t, r_out - 1), d, ts), 1);

dot_schedule_kvec

The dot-product schedule for a contraction whose right operand stores the contraction unit-stride per output cell ( \(x \cdot W^T\) — the fc/linear layout, where fuse_tensor's read-through absorbs the weight transpose into the access map). Vectorizing the last output dim would read a different weight row per lane (gathers), so instead

• the contraction itself runs as the vector loop — unit-stride in both the input row and each weight row, folded with a horizontal sum (see btensor.mr_nest's reduction-vectorized branch, selected by a vdim naming the innermost reduction dim, where unroll names the jam width);
• a 2-D block of output cells is unroll-and-jammed into its body over independent vector accumulators: the last output dim strip-mined by the largest divisor ≤ 4 (jammed lanes share the input-row loads) × the second-to-last by 2 (jammed rows share the weight-row loads — without this the whole weight matrix re-streams once per batch row).

An extent still open at lowering time, an indivisible one, and the missing row dim of a rank-1 output all take the neutral split (1), whose size-1 jam dim folds away.

lam dot_schedule_kvec (r_out: Nat) (t: Tileable): Tileable =
    // A block counts only when it divides AND leaves a non-trivial outer dim (`b < k`): an exact
    // split's leading size-1 axis would type-collapse the packed output («1; T» ≡ T).
    lam blk_by (k b: Nat): Bool = core.bit2.and_ 2 (k / b * b == k, b < k);
    let d   = r_out - 1;
    let n   = t#So#(core.idx 0 (t#Ro) d);
    let bs  = if_static (n, ((1, 2)#(blk_by (n, 2)), 4)#(blk_by (n, 4)), 1);
    let t2  = strip_mine_par (t, d, bs);
    let d2  = r_out - 2;
    let m   = t#So#(core.idx 0 (t#Ro) d2);
    let ms  = (1, if_static (m, (1, 2)#(blk_by (m, 2)), 1))#(2 <= r_out);
    let t3  = strip_mine_par (t2, d2, ms);
    // Dims now (…, m/ms, ms, n/bs, bs | red): swap ms ↔ n/bs so both jam dims are the trailing
    // parallel pair. Identity for a rank-1 output, where the neutral split already sits next to
    // `bs` ((n/bs, 1, bs | red)). `r3` is spelled from t3's own (positional) fields so the pack's
    // length is definitionally the «t#Ro + t#Rr» the dependent `interchange` domain expects —
    // `core.nat.add` does not reassociate an equivalent spelling.
    let r3   = t3#1_9 + t3#2_9;
    let swap = 2 <= r_out;
    let pa   = r_out - 1;
    let pb   = r_out;
    let perm = ‹j: r3; let jj = core.bitcast Nat j;
                       core.select (core.bit2.and_ 2 (swap, jj == pa), pb,
                       core.select (core.bit2.and_ 2 (swap, jj == pb), pa, jj))›;
    unroll_trailing (vectorize (interchange (t3, perm), r3 - 1), 2);

pool_schedule

The pooling schedule: unroll the window while it stays within budget. Pooling reads its single input in every loop dim (no broadcast operand), so ow scores [stride_w = 1] against an un-unrolled kw's [dilation_w = 1] — the row accumulator never wins and the nest stays classic.

lam pool_schedule (stride dilation: «2; Nat») (t: Tileable): Tileable =
    let u     = unroll_budget (t, 32);
    let s_ow  = (0, 1)#(stride#1_2 == 1);
    let s_red = core.select (u == 0, (0, 1)#(dilation#1_2 == 1), core.select (u < t#Rr, 0, s_ow));
    let ow    = t#Sr#(core.idx 0 (t#Ro + t#Rr) 3);
    let vec   = if_static (stride#1_2 * dilation#1_2 * ow, s_red < s_ow, ff);
    unroll_trailing (vectorize_if (t, vec, 3), u);

Implementations

The _impl annexes lower_tensor re-applies to expand the derived operations above, together with the fold steps and access maps they read through. Most are a single map_reduce over a copy step — the operation itself lives in the access map.

The declaration order is load-bearing beyond the usual definition-before-use: lower_tensor resolves an _impl in the world it is building, so an impl must be declared before any annex whose body applies the axiom it implements — reshape_impl before conv_impl and dot_product_impl, which reshape their packed results.

tensor_copy

Ignores the (⊥) accumulator and returns the read element — the fold step of every pure re-indexing operation below.

fun tensor_copy {T: *} (acc: T, y: «1; T»)@tt: T = return (y#0_1);

bid_map

o ↦ (o#(index#i) · [s_in#i ≠ 1])_i: input axis i reads output-loop var o#(index#i) (the transpose), unless it is a broadcast axis (s_in#i = 1), which reads index 0.

lam bid_map {r_in r_out: Nat} (s_in: «r_in; Nat», index: «r_in; Idx r_out»)
            (o: «r_out; affine.Index»): «r_in; affine.Index» =
    ‹i: r_in; affine.semiop.mul (o#(index#i), core.select (s_in#i == 1, 0, 1))›;

broadcast_in_dim_impl

anx lam broadcast_in_dim_impl {T: *, r_in r_out: Nat}
    (s_in: «r_in; Nat», s_out: «r_out; Nat», input: «s_in; T», index: «r_in; Idx r_out»): «s_out; T» =
    map_reduce (s_out, s_out) (tensor_copy, ⊥: T)
        affine.id (bid_map (s_in, index),) input;

flip_map

o ↦ ((s_in#d − 1) − o#d)_d: the reverse read map.

lam flip_map {r: Nat} (s_in: «r; Nat») (o: «r; affine.Index»): «r; affine.Index» =
    ‹d: r; affine.op.sub (affine.lit (s_in#d - 1), o#d)›;

flip_impl

anx lam flip_impl {T: *, r: Nat} (s_in: «r; Nat») (input: «s_in; T»): «s_in; T» =
    map_reduce (s_in, s_in) (tensor_copy, ⊥: T)
        affine.id (flip_map s_in,) input;

gather_pointwise_elem_impl

One output element: reads index at the output coordinate, then the source at that coordinate with axis dim replaced by the loaded value.

anx lam gather_pointwise_elem_impl {T: *, r: Nat}
                                   (s_src: «r; Nat», s_idx: «r; Nat»)
                                   (dim: Idx r)
                                   (out_indices: «r; I64»)
                                   (input: «s_src; T», index: «s_idx; Idx (s_src#dim)»): T =
    let index_coords = ‹d: r; core.conv.u (s_idx#d) (out_indices#d)›;
    let source_axis_index = core.bitcast I64 (get (index_coords, index));
    let source_coords = ‹d: r;
        let coordinate = core.select (core.icmp.e @r (d, dim), source_axis_index, out_indices#d);
        core.conv.u (s_src#d) coordinate›;
    get (source_coords, input);

map_impl

A map_reduce whose fold step drops the accumulator and applies app to the element tuple.

anx lam map_impl {T: *, ni: Nat, Is: «ni; *»}
                     (app: «i: ni; Is#i» → T)
                     {r: Nat, s: «r; Nat»}
                     (is: «i: ni; «s; Is#i» »)
                     : «s; T» =
  fun app_mr (x: T, y: «i: ni; Is#i»)@tt = return (app y);
  // The `{Tis, Ris, Sis}` group stays explicit: with a *symbolic* `ni`, solving `Sis#i ≡ s` for all `i`
  // (i.e. `Sis := ‹ni; s›`) is Miller pattern unification. Adding that rule to `Checker::alpha_impl_` was tried
  // and reverted — even with an occurs check it fires inside `tuple`'s machinery and commits a wrong solution.
  map_reduce (s, s) @(Is, ‹ni; r›, ‹ni; s›) (app_mr, ⊥: T)
      // `‹ni; …›` builds the pack bottom-up, so there is no expected type to drive implicit insertion for the
      // element; `affine.id` needs its rank here.
      affine.id ‹ni; affine.id @r› is;

pool_fun

Wraps a binary fold g as a map_reduce reduction Fn [T, «1; T»] → T.

pub fun pool_fun {T: *} (g: [T, T] → T) (acc: T, y: «1; T»)@tt: T = return (g (acc, y#0_1));

pool2d_in_map

Windowed input read over the loop vector (n, c, oh, ow, kh, kw): (n, c, stride#0·oh + dilation#0·kh, stride#1·ow + dilation#1·kw).

lam pool2d_in_map (stride dilation: «2; Nat») (o: «6; affine.Index»): «4; affine.Index» =
    (o#0_6, o#1_6,
     affine.op.add (affine.semiop.mul (o#2_6, stride#0_2), affine.semiop.mul (o#4_6, dilation#0_2)),
     affine.op.add (affine.semiop.mul (o#3_6, stride#1_2), affine.semiop.mul (o#5_6, dilation#1_2)));

pool_impl

anx lam pool_impl {T: *} (g: [T, T] → T, init: T) {n c h w: Nat}
                      (kernel stride dilation padding s_out: «2; Nat»)
                      (input: «n, c, h, w; T»)
                    : «n, c, s_out#0_2, s_out#1_2; T» =
    let s_in     = (n, c, h, w);
    let pad_lohi = (0, 0, padding#0_2, padding#1_2);
    let padded   = pad s_in (0, pad_lohi, pad_lohi) (input, init);
    let s_padded = shape 4 padded;
 
    // Schedule (frontend, selected from the parameters — see `pool_schedule`).
    materialize (pool_schedule (stride, dilation) (
        mr_tileable @1 @(4, 2)
            ((n, c, s_out#0_2, s_out#1_2), (0, 1, 2, 3), (n, c, s_out#0_2, s_out#1_2, kernel#0_2, kernel#1_2))
            @(T, 4, s_padded)
            (pool_fun g, init, id @T)
            (proj_map @(6, 4) (0, 1, 2, 3))
            (pool2d_in_map (stride, dilation),)
            padded));

repeat_map

o ↦ (o#d mod s_in#d)_d: the wrap-around (tiling) read map.

lam repeat_map {r: Nat} (s_in: «r; Nat») (o: «r; affine.Index»): «r; affine.Index» =
    ‹d: r; affine.semiop.rem (o#d, s_in#d)›;

repeat_impl

anx lam repeat_impl {T: *, r: Nat} (s_in: «r; Nat») (s_out: «r; Nat») (input: «s_in; T»): «s_out; T» =
    map_reduce (s_out, s_out) (tensor_copy, ⊥: T)
        affine.id (repeat_map s_in,) input;

reshape_map

o ↦ delinearize(linearize(o, s_out), s_in): the row-major output→input index map.

anx lam reshape_map {r_in r_out: Nat} (s_in: «r_in; Nat», s_out: «r_out; Nat»)
                (o: «r_out; affine.Index»): «r_in; affine.Index» =
    affine.delinearize (affine.linearize (o, s_out), s_in);

reshape_impl

anx lam reshape_impl {T: *, r_in r_out: Nat}
                         (s_in: «r_in; Nat») (s_out: «r_out; Nat») (input: «s_in; T»): «s_out; T» =
    map_reduce (s_out, s_out) (tensor_copy, ⊥: T)
        affine.id (reshape_map (s_in, s_out),) input;

conv_fun

Fused multiply-add over the ring R; ab#0 is the padded-input element, ab#1 the weight.

fun conv_fun (R: Ring) (acc: R#T, ab: «2; R#T»)@tt: R#T = return (R#add (acc, R#mul (ab#0_2, ab#1_2)));

conv2d_in_map

The «7; affine.Index» loop vector is o = (n, co, oh, ow, ci, kh, kw), so the weight read map (co, ci, kh, kw) and the output write map (n, co, oh, ow) are proj_map projections of it; only this padded-input read needs arithmetic: (n, ci, stride#0·oh + dilation#0·kh, stride#1·ow + dilation#1·kw).

lam conv2d_in_map (stride dilation: «2; Nat») (o: «7; affine.Index»): «4; affine.Index» =
    (o#0_7, o#4_7,
     affine.op.add (affine.semiop.mul (o#2_7, stride#0_2), affine.semiop.mul (o#5_7, dilation#0_2)),
     affine.op.add (affine.semiop.mul (o#3_7, stride#1_2), affine.semiop.mul (o#6_7, dilation#1_2)));

conv_impl

anx lam conv_impl (R: Ring) {n cin cout h w kh kw: Nat}
                      (stride dilation padding: «2; Nat»)
                      (input: «n, cin, h, w; R#T», weight: «cout, cin, kh, kw; R#T»)
                    : «conv_shape (n, cout, (h, w), (kh, kw), stride, dilation, padding); R#T» =
    // Zero-pad only the spatial axes (lo = hi = (0, 0, pad_h, pad_w)); the pad value is the ring zero.
    let s_in     = (n, cin, h, w);
    let pad_lohi = (0, 0, padding#0_2, padding#1_2);
    let padded   = pad s_in (0, pad_lohi, pad_lohi) (input, R#_0);
    let s_padded = shape 4 padded;
    let s_out    = conv_shape (n, cout, (h, w), (kh, kw), stride, dilation, padding);
    let oh       = s_out#2_4;
    let ow       = s_out#3_4;
 
    // The SCHEDULE, selected in the frontend from the parameters (see `conv_schedule`): unroll the
    // kh·kw window into the loop body (constant tap offsets) when it fits the budget, and run `ow`
    // as the innermost, unit-stride loop — for a fixed tap, adjacent output pixels read adjacent
    // input pixels, so the vectorizer emits contiguous vector loads instead of gathers. The
    // schedule may strip-mine `oh` (row blocking), packing the materialized tensor along that axis;
    // the row-major `reshape` restores `s_out` (an identity for the neutral split, and absorbed into
    // the consumers' access maps by `fuse_tensor`'s read-through in a fused graph).
    let sched = conv_schedule (stride, dilation) (
        mr_tileable @2 @(4, 3)
            ((n, cout, oh, ow), (0, 1, 2, 3), (n, cout, oh, ow, cin, kh, kw))
            @((R#T, R#T), (4, 4), (s_padded, (cout, cin, kh, kw)))
            (conv_fun R, R#_0, id @(R#T))
            (proj_map @(7, 4) (0, 1, 2, 3))
            (conv2d_in_map (stride, dilation), proj_map @(7, 4) (1, 4, 5, 6))
            (padded, weight));
    // `sched#3_9` is the scheduled Tileable's (possibly packed) output shape `So`.
    reshape (sched#3_9) s_out (materialize sched);

dot_general_fun

The fold step: a fused multiply-add in the ring R.

pub fun dot_general_fun (R: Ring) (x: R#T, (y z: R#T))@tt: R#T = return (R#add (x, (R#mul (y, z))));

dot_general_pick

Sends input axis i to the loop id it reads — batching dims first, then the contractions, then the remaining free axes.

pub lam dot_general_pick {n: Nat} {na nb nc: Nat} (a: «na; Idx n», b: «nb; Idx n», c: «nc; Idx n») (off_1 off_2: Nat) (i: Idx n): Nat =
  match vec.first core.icmp.e (b, i) with
      | bi: Idx nb => core.bitcast Nat bi
      | _:  []     => match vec.first core.icmp.e (c, i) with
          | ci: Idx nc => off_2 + core.bitcast Nat ci
          | _:  []     => match vec.first core.icmp.e (a, i) with
              | ai: Idx na => off_1 + core.bitcast Nat ai
              | _:  []     => 0;

dot_product_impl

r1/r2 are an explicit group here, unlike in the dot_product axiom: they must be literals already when this lam is β-reduced, since the subscript rows are built with vec.diff/vec.fold.l, whose normalizers decline to fold over a non-literal rank, and nothing re-normalizes the body once a Hole is resolved later. {s1 s2} has no such problem and stays implicit.

tensor::Lower re-applies the axiom's groups positionally, so the two curry shapes still match — except for the leading fastest_2, which has no axiom counterpart: it is fastest_axis of the right operand, pre-applied by Lower where that operand is concrete. The schedule decision built on it lives in the body:

• a contracted fastest axis ( \(x \cdot W^T\), e.g. behind a weight transpose the read-through absorbs) selects dot_schedule_kvec, anything else dot_schedule;
• the selected schedule may strip-mine an output dim (row/output-cell blocking), packing the materialized tensor along that axis, whereupon the row-major reshape restores s_out — the identity for the neutral split, and absorbed into the consumers' access maps by fuse_tensor's read-through in a fused graph.

anx lam dot_product_impl
  (fastest_2: Nat)
  (R: Ring) (r1 r2: Nat) {nc nb: Nat}
  (c1: «nc; Idx r1», c2: «nc; Idx r2», b1: «nb; Idx r1», b2: «nb; Idx r2»)
  {s1: «r1; Nat», s2: «r2; Nat»} (a: «s1; R#T», b: «s2; R#T»)
    : let s_out = dot_general_shape (c1, c2, b1, b2) (s1, s2);
      «s_out; R#T»
    = let s_out    = dot_general_shape (c1, c2, b1, b2) (s1, s2);
      let bc_1     = vec.cat (b1, c1);
      let s1_res   = vec.diff (s1, bc_1);
      let n_s1_res = vec.len s1_res;
      let bc_2     = vec.cat (b2, c2);
      let f        = dot_general_fun R;
      let r_out    = vec.len s_out;
      let a1       = vec.diff (‹i: r1; i›, bc_1);
      let a2       = vec.diff (‹i: r2; i›, bc_2);
      let subs_1   = ‹i: r1; dot_general_pick (a1, b1, c1) (nb, r_out) i›;
      let subs_2   = ‹i: r2; dot_general_pick (a2, b2, c2) (nb + n_s1_res, r_out) i›;
      // The subscript rows use consecutive loop ids: 0 … r_out−1 are the output loops
      // and r_out … r_out+nc−1 the contractions, whose bounds are the contracting extents of `a`.
      let s_red    = ‹i: nc; s1#(c1#i)›;
      let s_full   = vec.cat (s_out, s_red);
      let n_loops  = r_out + nc;
      let t        = mr_tileable @2 @(r_out, nc)
                         (s_out, ‹i: r_out; core.bitcast Nat i›, s_full)
                         @((R#T, R#T), (r1, r2), (s1, s2))
                         (f, R#_0, id @(R#T))
                         (proj_map @(n_loops, r_out) ‹i: r_out; core.bitcast Nat i›)
                         (proj_map @(n_loops, r1) subs_1, proj_map @(n_loops, r2) subs_2)
                         (a, b);
      // Schedule (frontend, selected from the parameters — see `dot_schedule`/`dot_schedule_kvec`):
      // vectorize the contraction iff it is the right operand's effective fastest axis (x·Wᵀ);
      // otherwise the last output dim is unit-stride in that operand (A·B) and runs as the vector
      // loop. Dispatched as an extract over thunks, each materializing its own concretely-scheduled
      // Tileable: behind a stuck extract the record's `Out` field would not reduce, so the
      // reshape's input could not be checked. `sched#3_9` is the scheduled Tileable's (possibly
      // packed) output shape `So`.
      let kvec = option.is_some (vec.first `== (‹i: nc; core.bitcast Nat (c2#i)›, fastest_2));
      lam go_n    (u: []): «s_out; R#T» =
          let sched = dot_schedule r_out t;
          reshape (sched#3_9) s_out (materialize sched);
      lam go_kvec (u: []): «s_out; R#T» =
          let sched = dot_schedule_kvec r_out t;
          reshape (sched#3_9) s_out (materialize sched);
      (go_n, go_kvec)#kvec ();

bmm_impl

The {r1 r2}/{s1 s2} groups are explicit for the same reason as in product_2d_impl.

anx lam bmm_impl (fastest_2: Nat) (R: Ring) {B M K N: Nat} (t1: «B, M, K; R#T», t2: «B, K, N; R#T»): «B, M, N; R#T»
        = dot_product_impl fastest_2 R (3, 3) (2_3, 1_3, 0_3, 0_3) (t1, t2);

product_2d_impl

The {r1 r2} and {s1 s2} groups are passed explicitly — not because inference fails, but because of when it succeeds: dot_product_impl is a curried lam, so applying it to a group β-reduces the body right away. With the ranks still unsolved Holes the vec.diff/vec.fold.l subscript machinery is built over non-literals, its normalizers decline to fold, and nothing re-normalizes the body once the Holes are resolved: the result type-checks but carries the whole vec.* scaffolding inside its access maps (refly.equiv.struc_eq in lit/tensor/dot_product.mim pins this down). Passing literals up front folds it away.

anx lam product_2d_impl (fastest_2: Nat) (R: Ring) {m k l: Nat} (t1: «m, k; R#T», t2: «k, l; R#T»): «m, l; R#T»
        = dot_product_impl fastest_2 R (2, 2) (1_2, 0_2, (), ()) (t1, t2);

scatter_step_impl

One update: reads index and updates at the visited coordinate and sets the element into the accumulated tensor at the destination it selects.

anx lam scatter_step_impl {T: *, r: Nat}
                          (s_src: «r; Nat», s_idx: «r; Nat», s_updates: «r; Nat»)
                          (dim: Idx r)
                          (visit_indices: «r; I64»)
                          (acc: «s_src; T», index: «s_idx; Idx (s_src#dim)», updates: «s_updates; T»): «s_src; T» =
    let index_coords = ‹d: r; core.conv.u (s_idx#d) (visit_indices#d)›;
    let update_coords = ‹d: r; core.conv.u (s_updates#d) (visit_indices#d)›;
    let destination_axis_index = core.bitcast I64 (get (index_coords, index));
    let update = get (update_coords, updates);
    let destination_coords = ‹d: r;
        let coordinate = core.select (core.icmp.e @r (d, dim), destination_axis_index, visit_indices#d);
        core.conv.u (s_src#d) coordinate›;
    set (destination_coords, acc, update);

slice_map

o ↦ (start#d + step#d · o#d)_d: the strided read map.

lam slice_map {r: Nat} (start step: «r; Nat») (o: «r; affine.Index»): «r; affine.Index» =
    ‹d: r; affine.op.add (affine.semiop.mul (o#d, step#d), affine.lit (start#d))›;

slice_impl

anx lam slice_impl {T: *, r: Nat} (s_in: «r; Nat»)
                       (start step s_out: «r; Nat») (input: «s_in; T»): «s_out; T» =
    map_reduce (s_out, s_out) (tensor_copy, ⊥: T)
        affine.id (slice_map (start, step),) input;

transpose_map

o ↦ (o#(perm#0), …, o#(perm#(r−1))): input axis j reads output-loop var o#(perm#j).

lam transpose_map {r: Nat} (perm: «r; Idx r») (o: «r; affine.Index»): «r; affine.Index» = ‹j: r; o#(perm#j)›;

transpose_impl

anx lam transpose_impl {T: *, r: Nat}
                           (permutation: «r; Idx r») {s: «r; Nat»} (input: «s; T»)
                         : «transpose_shape (s, permutation); T»
                         = let out_s = transpose_shape (s, permutation);
                           map_reduce (out_s, out_s) (tensor_copy, ⊥: T)
                               affine.id (transpose_map permutation,) input;

transpose_2d_impl

transpose_impl at the swap permutation.

anx lam transpose_2d_impl {T: *} {s: «2; Nat»} (input: «s; T»): «s#tt, s#ff; T» =
  transpose_impl (tt, ff) input;

binary_impl

map_impl over two tensors.

anx lam binary_impl {Ti1 Ti2 To: *} (app: [Ti1, Ti2] → To) {r: Nat, s: «r; Nat»} (is: [«s; Ti1», «s; Ti2»]): «s; To» =
  map_impl app is;

select_impl

map_impl of core.select over three tensors.

anx lam select_impl {T: *} {r: Nat, s: «r; Nat»} (is: [«s; Bool», «s; T», «s; T»]): «s; T» =
  map_impl core.select is;

unary_impl

map_impl over one tensor.

anx lam unary_impl {Ti To : *} (app: Ti → To) {r: Nat, s: «r; Nat»} (i: «s; Ti»): «s; To» =
  map_impl app i;

Phases

reassoc

Reassociates chains of product_2d with the classic matrix-chain-order dynamic program — but minimizing vector-lane slots instead of scalar multiplications.

• The vector loop runs over a product's trailing extent (dot_schedule) or, for a transposed operand, over its contraction (dot_schedule_kvec), and a literal extent is charged rounded up to whole lanes — -X tensor:reassoc-vec=<n>, 1 to count plain scalar multiplications. A bracketing whose intermediates are too narrow to fill a vector thus pays for the lanes it leaves idle, which is exactly what the multiplication count misses.
• Extents need not be literal: a cost is a polynomial in the symbolic extents, ordered by coefficient-wise ≤, which proves ≤ for every instantiation because extents are non-negative.
• That order is only partial, so a chain can have several bracketings that each win for some instantiation — a batch dimension favouring a different one when small than when large. Up to 4 matrices — -X tensor:reassoc-max=<n>, below 3 to switch it off — those survivors are emitted side by side behind a run-time comparison of their costs, so only the winner's thunk runs; a longer chain is only reassociated where one bracketing provably wins for every extent.
• Common sub-products are hash-consed across the alternatives, so the code growth stays well below \(\mathit{Catalan}(n-1)\).
• A product with more than one consumer stays a chain boundary, as pulling it apart would leave it materialized for its other consumers as well.
• Reassociation is only exact for an associative ring; on floating-point rings it changes the rounding, just like -ffast-math.

axm reassoc: compile.Phase;

lower_tensor

Lowers the derived operations to the core ones by re-applying the matching _impl annexes.

axm lower_tensor: compile.Phase;

lower_map_reduce

Lowers the core operations, generate included, to their underlying primitives: loops, extract, insert, packs, ...

axm lower_map_reduce: compile.Phase;

lower_get_set

Lowers get / set to extract / insert.

axm lower_get_set: compile.Phase;

fuse_tensor

Fuses pure producer map_reduces into their consumers' access maps.

• A producer is only fused where the consumer reads it injectively — its access map uses every loop index, possibly strided/shifted — so fusion never recomputes a producer element; strided reads even skip producer elements entirely.
• In the reverse direction, a trailing elementwise map over a reducing map_reduce is composed into the producer's post epilogue (the GEMM-epilogue pattern), provided the map is its only consumer.

axm fuse_tensor: compile.Phase;

lower_to_mem

Bufferizes the core operations onto the buffer layer: tensor values become buffer.Buf handles, get / set / generate / map_reduce become buffer.read / buffer.write / buffer.alloc threading mem.M. mem.add_mem (scheduled next in the pipeline) threads the memory; run buffer.lower_ptr afterwards to reach mem.Ptr.

axm lower_to_mem: compile.Phase;