forall tensor constructs in VADL

[benedikt-l-huber]

This is a summary of a discussion where @Jozott00 and @benedikt-l-huber tried to
find an agreement on the semantics of VADL's forall constructs.

There are three forall constructs in VADL.

forall ... tensor

The forall i in 3..0 tensor someExpression expression.
someExpression can be of any type (also a tensor).
The number of dimensions of the resulting tensor is the number of
indices in the index list ( e.g., i in 3..0, j in 4..1) plus the number of dimensions of someExpression
e.g.:

let e = forall i in 3..0, j in 2..0 tensor (forall k in 2..0 tensor (i, j, k) as UInt<3><32>) in

Has 5 dimensions. We can access a single Bit like this e(1)(0)(1)(2)(14).
Note that all operations in someExpression have to be defined in VADL.

let z = (...) as UInt<2><2><32> in
let p = (...) as UInt<2><2><32> in

let fl = forall j in 1..0, k in 1..0, l in 1..0 tensor z(j)(k)(l) + p(j)(k)(l) in  // OK, addition of Bits
let fk = forall j in 1..0, k in 1..0            tensor z(j)(k)    + p(j)(k)    in  // OK, addition of UInt<32>
let fj = forall j in 1..0                       tensor z(j)       + p(j)       in  // Currently not OK in VADL
let fn = z + p                                                                     // Currently not OK in VADL

The reason is the types on which + is defined. In the fj example z(j) and p(j)
are of type UInt<2><32> for which + is not defined.

Tensor construction with forall ... tensor is done from left to right. I.e., the "next" element is
appended to the right.
In the following example x = b and a is x in reversed order.

let x = (6, 5, 4) as UInt<3><32> in
let a = forall i in 0..2 tensor x(i) in
let b = forall i in 2..0 tensor x(i) in
let cb = (x = b) in

forall ... do

The forall i in 3..0 do someStatement statement which can read write registers.
VADL's semantics ensures that all reads occur before all writes.
The forall ... do statement is especially useful for strides writes.
e.g.:

forall i in 3..0 do
  X(i*2) := p(i)

In theory this could be rewritten using only the forall ... tensor expression as

let z = X in
let e = forall i in 7..0 tensor
  if (i % 2 = 1) then z(i)
  else p(i/2) in
X := e

But this is much less readable and it would necessitate further analysis
to find out that writes to the odd indices of X are the identity function.

forall ... fold

The forall i in 0..3 fold + with someExpression can be used to make a reduction of one dimension.
There can only be a single index in the list. It signifies the dimension across which the reduction
occurs. Again the operation (in this example the +) has to work with someExpression.
Right now VADL restricts the operation to a set of these binary operators + - * & | ^.
Note that this list contains also contains - which is not distributive, thus the
order of the indices has to be fixed. It is defined by the index list which can be
increasing or decreasing. Again the "next" operand is applied to the right of the operand.
So e = g and f = y in the following example:

let e = forall i in 0..2 fold - with k(i) in 
let g = (k(0) -  k(1)) - k(2)  in
let f = forall i in 2..0 fold - with k(i) in 
let y = (k(2) -  k(1)) - k(0)  in

This semantics leads to a composition problem, like in the following example:

let k = ((1,2,3), (4,5,6)) as UInt<2><3><32> in // UInt<outer dim><inner dim><bit width>
// fold on both dimensions, i.e., add all elements
let d = forall i in 1..0 fold + with (forall j in 2..0 fold + with k(i)(j)) in
let d_t = ( d = 21 ) in

This works in VADL since all uses of + are supported.
If we want to compose the same thing with let expressions for the individual dimensions
the problem becomes apparent. The expression e still has to list all dimensions, so
composition does not work well.

// Again fold both dimensions, but try to compose the individual dimensions.
let e = forall i in 1..0 tensor (forall j in 2..0 fold + with k(i)(j)) in
let f = forall i in 1..0 fold + with e(i) in
let f_t = ( f = d ) in

A possible solution would be to define + also for more dimensional tensors in VADL.
Then the following would be possible, which composes better.
Note that the + for expression m now has to be defined on for the type UInt<3><32>

let m = forall i in 1..0 fold + with k(i) in
let n = forall j in 2..0 fold + with m(j) in
let n_t = ( n = d ) in

As a comparison: Pytorch provides a number of reduction operators e.g., argmax that
do a similar thing than VADL's forall ... fold. These pytorch operators explicitly
take the dimension on which to perform the reduction as a parameter.

For the type UInt<32>, in VADL + is already defined on both the inner type (Bits) and
the outer type (UInt). So in the following both k and l are defined in VADL at the moment.

let m = { 3, 2 }  as UInt<2><32> in
let n = { 5, 4 }  as UInt<2><32> in
let k = m(0) + n(1) in
let l = m(0)(3)  + n(1)(4) in

So extending operations to even more dimensions seems to be quite natural for VADL.
This would make it possible to describe more complex tensor operations in a convenient way.

Nevertheless, we think that this proposed extension is not necessary at the moment.
Since most tensor operation that will be described in VADL are far less complex than
the ones that appear in e.g., pytorch.