 |
MimIR
MimIR is my Intermediate Representation
|
|
MimIR
MimIR is my Intermediate Representation
|
Bufferized tensor operations: a generic map_reduce loop-nest abstraction plus derived operations (matrix product, transpose, pad, concat, ...), lowered to affine.For loop nests over buffer element operations.
This is the bufferized mirror of the tensor plugin: every op here is the buffer.Buf counterpart of a tensor.* op, takes the same groups in the same order, and threads mem.M on top. tensor.lower_to_mem bufferizes the SSA-world tensor.* ops onto these. Only the differences from that counterpart are documented below, and shapes are stated as the logical ones — buffer.Buf normalizes literal size-1 axes away and the loop generation folds its read/write indices to match.
Matrices are n-dimensional dependent arrays, represented by the shared buffer.Buf type. The element operations (buffer.alloc / buffer.read / buffer.write / buffer.lit / buffer.shape) and their lowering to mem.Ptr (buffer.lower_ptr) live in the buffer plugin.
Represents an algebraic ring. E.g., let nat_ring = (Nat, 0, core.nat.add, core.nat.mul);.
tensor.Ring is an alias for this type: tensor bufferizes onto *, so btensor is the lower of the two layers and owns the definition.
The loop-nest core every other operation is lowered onto: product_2d and transpose_2d are plain lams over map_reduce, the remaining axioms are expanded by lower_map_reduce directly.
The mem-threaded CPS identity (with an empty epilogue-input pack) — the neutral post epilogue normalize_map_reduce delegates with.
map_reduce_post without the epilogue and with the neutral schedule — the buffer-world counterpart of tensor.map_reduce. Its normalizer delegates to map_reduce_post with post = id (so Tp = To), no epilogue inputs (nps = 0), no vector dim, and no unrolling.
Unlike tensor.map_reduce this cannot be a delegating lam: applying a lam instantiates its later group domains eagerly, and the is domain here contains the dependent axm application buffer.Buf (Ris#i, Sis#i, Tis#i), whose signature cannot be checked over the still-unsolved {Tis, Ris, Sis} holes — the tensor-level counterpart is a plain array and checks structurally. An axm application, in contrast, stays stuck until the checker has solved the holes from the operands.
The buffer-world counterpart of tensor.map_reduce_post and the workhorse behind every operation below: the affine-indexed generalisation of map_reduce on buffer.Buf operands, threading mem.M. Same groups in the same order as there, so only the differences:
The canonical schedule value: a chooser over the target's loop-nest builder, picking
tensor.mk_sched is an alias for this lam: btensor owns it for the same reason it owns Ring.
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. tensor.proj_map is an alias for this lam.
Pure re-indexing: every element of the result is one element of an input, so the operation lives entirely in the access map.
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. The exception to the rule above: ri/si and ro/so are the actual — possibly size-1-folded — buffer shapes, inferred from the operands, while s_in/s_out are the logical ones.
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 and 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, since concat_shape lives in the tensor plugin.
Pads each axis: out[…o…] = input[…o−lo…] inside the interior, value (mode 0, constant) or the clamped edge element (mode ≠ 0, replicate) outside. s_out cannot be computed here — pad_shape lives in the tensor plugin — so it is passed as its own group in front of the operands, as in concat.
Ignores the (⊥) accumulator and just returns the read value — the mem-threaded counterpart of tensor_copy.
Permutes the dimensions of a 2-dimensional matrix. As there, the shape is a single «2; Nat» vector instead of two separate Nats and {T} precedes {s}; like product_2d this is a lam over map_reduce rather than an axm.
Transpose a matrix by reading the input at the swapped loop variables; there is no reduction loop.
Operations that fold a reduction dim away.
Fused multiply-add over the ring R; ab#0 is the t1 element, ab#1 the t2 element. The mem-threaded counterpart of the tensor plugin's conv_fun.
Computes the matrix product of two 2-dimensional matrices in the ring R. [R: Ring] is the leading explicit group, as there: a ring cannot be recovered from R#T by inference.
Unlike tensor.product_2d this is a lam wrapping map_reduce directly, not an axm paired with an _impl annex and a lowering phase: nothing in this plugin inspects the high-level form — no fusion phase over *, and the normalizer was a no-op — so the extra staging bought nothing. transpose_2d follows the same pattern.
The loop vector is (i, j, k) and follows the principle ij ← ik, kj ( \(\mathit{out}_{i,j} = \sum_k \mathit{t1}_{i,k} \cdot \mathit{t2}_{k,j}\)), so k is the single reduction loop.
Data-dependent access: an index buffer picks the coordinate along dim.
For each output coordinate o, reads the source at o with axis dim replaced by index[o]. The index element type proves that selected coordinates are in bounds along dim.
Starts with a copy of input, then applies updates in row-major index order; for duplicate destinations, the last update wins.
Loop nests are first-class IR values here: the schedule of a map_reduce_post is not data interpreted by a lowering — it is a chooser over the target's loop-nest builder (see tensor.mk_sched), carried as an operand; lower_map_reduce binds it to mr_nest and merely applies the result. The builders below construct every loop structure in the frontend IR (loop-implementation policies, window unrolling, the row-accumulator interchange); partial evaluation unfolds them once the ranks are known.
The type of anything with affine.For's calling convention over [mem, A] accumulators — a "loop implementation" policy value.
The default loop implementation: affine.For itself, eta-wrapped.
The loop implementation for a schedule's vector dim: the same affine.For, its body wrapped in the identity annotation ll.vec — the nomination to run this loop as the innermost unit-stride vectorized loop. affine.lower_for moves the wrapper onto the loop's exit condition for the backend (and drops it when the ll module is not loaded).
One nest of loop-implemented loops over the dim range [lo, hi) of the bounds Sr, threading an accumulator A and the shared «r; I32» loop-vector prefix; calls inner at the bottom with the prefix updated at the range's positions. Composed by a right fold — no recursion, and the flat I32 prefix avoids dependently-typed loop-variable tuples.
Like range_nest, but every dim in the range is fully unrolled: the inner continuation is chained once per iteration point by a right fold over the dim's (compile-time) extent — straight-line code with constant coordinates, no loops, and structurally terminating (no recursion, no PE filters: the fold itself unfolds once the extent is a literal).
The type of scheduled loop nests for a map_reduce over the domain Sr (ro leading parallel dims, rr trailing reduction dims, fold element type To): given the fold seed init, the per-point fold step cell (fold the elements at one loop vector into the accumulator) and the per-cell write-back wb (consume one folded value at its output coordinates), run the nest.
The scheduled nest builder, selecting one of the four nests below from the schedule:
Per-cell fold order is untouched by any of these transforms.
Lowers every operation of this plugin to affine.For loops over buffer element operations, binding each op's sched chooser to mr_nest.
1.18.0