 |
MimIR
MimIR is my Intermediate Representation
|
|
MimIR
MimIR is my Intermediate Representation
|
Bufferized tensor operations: a generic map_reduce loop-nest abstraction plus the pointwise and indexing operations (broadcast, pad, concat, gather, scatter), 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. Every op is in destination-passing style: it takes the output buffer dst right after the memory, writes into it, and returns it. Allocation is thus explicit and belongs to the caller (buffer.alloc), and an in-place update is a mere choice of dst. 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 btensor, so btensor is the lower of the two layers and owns the definition.
The loop-nest core; 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 schedule record a nest builder is handed:
Transforms update the field they decide (s#vdim ← d) and pass the rest through. tensor.Sched is an alias for this type: btensor owns it for the same reason it owns Ring.
The type of a schedule chooser: for whatever nest type N the lowering instantiates, it is handed the target's builder mk and returns the nest to run.
The canonical chooser: it hands the builder s itself.
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.
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.
Copies input into dst, 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).
The loop implementation for the loops enclosing a schedule's vector loop: affine.For with its body wrapped in ll.novec, so that LLVM vectorizes the nominated loop and none of the strided loops around it.
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.
Destination-passing map-reduce: apply the selected nest to ordinary fold and write-back functions.
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