RFC: Index Expression API

[fredrikbk]

This issue requests comments for a new index expression API proposal. It supercedes issues #72 and #111. The index expression module relates to other modules in the following way:

The API is in the taco namespace. It is designed to be convenient to use by us and by end users directly.

IndexVar

An IndexVar is an index variable in a tensor index expression.

IndexExpr

An IndexExpr is a tensor index expression. It has several subclasses including the standard expressions (e.g., +, - and *), Assign (e.g., a(i) = b(i)), and Access (e.g., A(i,j) where A is a TensorVar).

TensorVar

A TensorVar is either a variable in an index expression (e.g., A(i,j)) or the result of an index expression (e.g., A(i,j) = B(i,j) + C(i,j)).

A TensorVar is a source for tensor data, whether it comes from memory, disk, or a computed index expression. It is different from the Tensor class because it does not actually store data. When computing a compiled expression (using the function API) TensorVar instances must be bound to data.

TensorVar objects must be given a type and a format, or else they will be given defaults. The constructor is:

TensorVar::TensorVar(Type, Format);

Type consist of a component datatype (e.g. double) and the tensor dimensions. Tensor dimensions are given as a vector or initializer list and can either be hard-coded or parametric. E.g.,

Type fixed(Double, {256,256});

Dimension n, m;
Type variable(Double, {n,m});

Type blocked(Double, {n, m, 3, 3});

Lowering Index Expressions

Index expressions can be lowered using the API's in the lower module, which turns them into the imperative IR in the ir module.

We can turn a tensor index expression into a runnable function with the compile function. For example

IndexVar i, j;
TensorVar A(Type(Double, {256,256}), Sparse);
TensorVar B(Type(Double, {256,256}), Sparse);
TensorVar C(Type(Double, {256,256}), Sparse);
A(i,j) = B(i,j) + C(i,j);

Function function = compile(A, {B,C}, Target::C);

TensorStorage myA;
TensorStorage myB = read("B.tns");
TensorStorage myC = read("C.tns");

function.assemble(myA, myB, myC);
function.compute(myA, myB, myC);

, where the arguments to compile are the input and output parameters of the function (in order) as well as the target backend, e.g., C, .

In the future, we may want to simultaneously lower multiple index expressions to take advantage of cross-expression optimization opportunities. Multiple expressions can also be lowered with the lower function by declaring multiple results.

TensorVar D, E;
D(i,j) = A(i,j) * E(i,j);

Function function = compile({A,D}, {B,C,E}, Target::C);