Compiler passes for OpenXLA for finite element methods

[connorjward]

Hello,

I am interested in using OpenXLA (via JAX) for performing finite element computations. As far as I can tell the necessary interface already exists, but one critical compiler transformation appears to be missing.

To be more specific, I would like to be able to write code like the following:

# 'input' and 'output' are arrays with arbitrary (large) shapes (a,) and (b,) respectively
# 'input_map' has shape (n, 3) and 'output_map' has shape (n, 6) where 'n' is very large

# collect 'local' data into an n x 3 array
input_packed = input[input_map]

# perform row-wise computations on the input data, resulting in an n x 6 array
output_packed = ...

# scatter result back out with atomic adds
output.at[output_map].add(output_packed)

The crucial transformation that doesn't appear to be happening is a "tile and fuse" transformation. Since the axis n appears in all steps of the computation it should be possible to tile the iteration to compute, say, 10 rows at once. As things stand this doesn't happen and input_packed is fully materialised as an enormous temporary which is not performant.

I have found this previous unsolved issue which I believe describes a very similar problem.

I would love to hear your feedback regarding whether you think it is reasonable to expect OpenXLA to perform this transformation and what would be necessary to enable it. I would potentially be able to have a crack myself some time in the future.

Some additional points to emphasise include:

• Apart from the final atomic scatter, each row may be computed independently from the others
• The per row computation can get very complex. It involves tensor contractions with some very small matrices and pointwise operations.