Skip to content

Loopy representation of tensor and linear algebra operations

This is in response to our call last week (with @dham, @wence-, and myself). So the conclusion we came to before the meeting ended was that, in order to fully transition Firedrake's code-generation layout into loopy, we need to have loopy call foreign functions in some linear algebra backend (Eigen for example).

Since the call ended, we continued to discuss this further and we have arrived at another potential solution. Loopy already can sufficiently express most of the linear algebra operations we require: addition, multiplication (both matrix/vector and matrix/matrix can be expressed as sum(i, a[i, j]*c[j ...] ...), and transposes are straightforward (a[i, j] => a[j, i]). Even tensor operations like outer/inner/dot, cross (a bit ugly) , trace/diag, we think, can all be represented in loopy. The point is that most of the algebra we do on arrays/finite element tensors in Firedrake can be represented in loopy with the exception of two operations:

  • inverse
  • determinant

These are array operations which return arrays, but require some form of linear algebra backend to evaluate as loopy is designed to only perform scalar operations. And so we're proposing rather than introducing a big general backend system for all linear algebra operations, we just let loopy introduce a linear algebra backend for inverses/determinants when prompted. So expressions like: x = A.inv * b may be turned into: (after setting {[i, n, j] : 0<=i<nrow and 0<=k<ncol and 0<=j<=...})

out[j] = sum(k, AINV[i, k] * b[k]))

with

AINV = inv(i, j, A[i, j])

getting turned into something which generates code for evaluating this using a specified library. inv here is something that acts as an 'escape hatch' for foreign function calls relating to matrix inversion. There is still is a question about what is the best way to do this and how loopy should be "translated" to specific libraries. This is just a quick sketch of the talk. Is this at all reasonable/within scope? I am still review the loopy documentation to make sure I'm getting things right. I'll update this as we progress on this topic. One thing I am concerned about is if this sort of "reinvents" what is already widely available in existing linear algebra packages. I'm definitely open to considering other options, this is just something we discussed last week.