XLA custom call with Eigen

[nickmcgreivy]

Hi team,

I'm trying to implement a custom operation (a sparse solver) in Jax using Eigen. The next step in my journey is to use jit (on CPU, for now) which requires a custom call.

Here is a repository where I am trying to implement a jitable matmul primitive which uses Eigen. Unfortunately, my matmul_xla_translation function gives a std::bad_alloc error when jiting matmul.

My questions are:

1. Am I doing something wrong in my xla_client.ops.CustomCallWithLayout call? If that looks good, any insight into what is causing it not to work?
2. Does Eigen work with XLA's custom_call? Or is there some sort of incompatibility I don't know about?

Thanks,
Nick

[PhilipVinc]

It might help you going through https://github.com/PhilipVinc/numba4jax or https://github.com/mpi4jax/mpi4jax where we do exactly this kind of things.

By quickly looking at your code I see those possible problems:

• You cast the input pointers in_ptr[0] and in_ptr[1] to MatrixXd, but are you sure that is legal? MatrixXd should contain a pointer to the actual memory, and the shape of the matrix. I doubt that the data that XLA is passing you is aligned the same way that Eigen Expects? Or maybe you do the unpacking correctly, but its quite tricky.
• You cast the out_ptr[0] to a matrixXd, which for the same reason above, I'm quite sure is wrong. out_ptr[0] should be a pointer to the bare memory stored by the eigen matrix, without other stuff in it.

I'd try to create the eigen matrices manually with Eigen::Map, and to pass to your C++ function not only the pointers but also the shapes, by using CustomCall and not CustomCallWithShape

[nickmcgreivy]

This is a helpful answer! Thanks. I'm fairly new to C++, and all of this pointer conversion stuff has been making my head spin. You're absolutely right. I managed to connect Eigen to pybind11 to XLA to JAX by doing:

1. Casting in_ptr[0] and in_ptr[1] to float* or double*
2. Passing sizes in in_ptr[2], in_ptr[3], and in_ptr[4] and casting them to std::int64_t.
3. Converting input pointers to Eigen matrices manually with Eigen::Map, i.e. MatrixXf x = Eigen::Map<const MatrixXf>(x_ptr,size1,size2)
4. Computing output with Map<MatrixXf>(z_ptr, x.rows(), y.cols() ) = x * y

and then using pybind11 and the custom primitive as in the 'extending JAX with custom C++' tutorial.

I spent some time looking through numba4jax, and thinking about how I might replicate that with Eigen. My suspicion is that the stumbling block would be the xla_cpu_custom_call_target where I'd have to convert pointers to numpy arrays, but I ended up getting things working with pybind11 so not needing to implement something similar.

One question: why do you suggest using CustomCall rather than CustomCallWithLayout? I see you use CustomCallWithLayout in numba4jax and CustomCall in mpi4jax, but I can't find an explanation of the differences between the two.

[PhilipVinc]

I don't really recall anymore what's the difference...
When I wrote mpi4py I was really reverse engineering those things, so it might be that I had not found/understood how to use CustomCallWithLayout.

[nickmcgreivy]

I'll just point out, for anyone looking to implement custom operations in Eigen, that matrix multiplication and sparse solver primitives with Eigen backends have now been implemented in this repository.

Hopefully this code is enough to demonstrate how to link Eigen with JAX. I'm happy to write up a tutorial if that would be helpful.

[PhilipVinc]

Out of curiosity, what is your objective?
Most (not all, indeed) functionality offered by Eigen is already available in Jax. Even sparse matrices.

[nickmcgreivy]

A sparse lu solve in JAX, i.e. solving Mx=b where M is sparse, i.e. Poisson's equation.

JAX has an lu solve where M is dense, but I haven't found a way to lu_solve in JAX where M is sparse.

[PhilipVinc]

If you use jax.experimental.sparse.sparsify on dense iterative solvers such as cg or gmres I'd expect them to work.
A sparse LU solve, I don't think so.
But it might be worth for you to open a Feature Request.
They have been pushing a lot on sparse features recently and maybe they'll implement it.