The affine Plugin

See also

mim::plug::affine

Dependencies

import mem;
import core;
import compile;
plugin refly;

Operations

For

This operation ranges from (including) begin to (excluding) end using step as stride. In addition, this loop manages n loop accumulators whose initial values init must be given. On each iteration, the given body is invoked and receives

• the current iteration index,
• the current values of the loop accumulators acc, and
• a yield continuation to continue with the next iteration, passing the new accumulator values.

After termination of the loop, exit is invoked with the final accumulator values.

axm For: {m n: Nat, Ts: «n; *»}
       → [Cn [iter: Idx m, acc: «i: n; Ts#i», yield: Cn «i: n; Ts#i»]] // body
       → [Cn «i: n; Ts#i»]                                             // exit
       → Cn [begin _end step: Idx m, init: «i: n; Ts#i»];

A For whose body is wrapped in the identity annotation ll.vec is a schedule's nominated vector loop; affine.lower_for strips the wrapper and re-attaches it to the loop's exit condition for the backend (when the ll plugin is loaded — otherwise it is dropped).

Affine Index Operations

Index

The affine index type is an opaque type that can only be manipulated through the following operations.

axm Index: *;

lit

Constructs an affine index from a constant natural number.

axm lit: Nat → Index;

op(add,sub)

Performs addition and subtraction on affine indices.

op(neg)

Negates an affine index.

axm semiop.(mul,ceildiv,floordiv,rem): [Index, Nat] → Index;

semiop(mul,ceildiv,floordiv,rem)

Performs semi-linear operations on affine indices. The second argument is a natural number that must be constant with regard to the indices.

• affine.op.mul Performs multiplication of an affine index with a constant natural number. The second argument must be a closed constant - a stronger requirement than affine.semiop imposes - which is enforced via refly.check.

pub mod op {
    axm (add, sub): [Index, Index] → Index;
    axm neg: Index → Index;
    anx lam mul (d : Index, c : Nat): Index =
        semiop.mul(d, refly.check (core.pe.is_closed(c), c, "The second argument of affine.op.mul must be a constant."));
}

map

Maps a function over a list of affine indices. Any captured variables of the function must be closed over the entire list of indices, i.e. they must not depend on the iteration index. The moduli are per-element: each input index i carries its own modulus sin#i and each output index j its own sout#j.

axm map: {m n: Nat}
       → {sin: «n; Nat», sout: «m; Nat»}
       → ([«n; Index»] → «m; Index»)
       → «i: n; Idx (sin#i)»
       → {a: Nat}
       → (mem.M a)
       → [(mem.M a), «i: m; Idx (sout#i)»];

id

The identity map on affine indices.

anx lam id {r: Nat} (o: «r; Index»): «r; Index» = o;

linearize

Row-major flattening of a multi-index: linearize(idxs, s) = Σ_k idxs#k · (∏_{j>k} s#j). s is the (row-major) shape; only its rank n need be a literal.

axm linearize: {n: Nat} → [idxs: «n; Index», s: «n; Nat»] → Index;

delinearize

Row-major unflattening: delinearize(lin, s)#d = (lin floordiv (∏_{j>d} s#j)) mod s#d. The inverse of affine.linearize; only the rank m of s need be a literal.

axm delinearize: {m: Nat} → [lin: Index, s: «m; Nat»] → «m; Index»;

Phases

lower_for

Lowers the affine.For operation to recursive function calls.

axm lower_for: compile.Phase;

lower_index

Lowers the affine index algebra (affine.Index, affine.lit, affine.op, affine.semiop, affine.map) to core arithmetic.

axm lower_index: compile.Phase;