Is it possible to write a SpMV kernel for TPU using jax Pallas?

[chaoming0625]

For many cases of spiking neural network modeling, we are using sparse matrix, and heavily using SpMV operations. I am wondering how to customize a SpMV kernel for TPU devices using the jax Pallas? If it is possible, how can I achieve that?

So many thanks!

[jakevdp]

It's definitely possible to write such a kernel, but unless you have some kind of highly structured sparsity, I doubt you'd be able to write a very performant kernel given the computational characteristics of a TPU. I know @sharadmv has thought a bit about this – he may be able to say more.

[chaoming0625]

Dear @jakevdp , thanks for your reply.

[sharadmv]

As Jake mentioned, TPUs are best at exploiting certain kinds of structured sparsity (block sparsity most easily). If you have arbitrary sparsity, it will be more difficult to utilize the hardware to speed it up. It is possible to express blocksparse matrix multiplies with Pallas. We have examples/documentation on how to do so on the way (the key is to use the PrefetchScalarGridSpec).

[chaoming0625]

Dear @sharadmv , thank you so much for your help! I am really looking forward to the examples and documentation. I am eager to try them out as soon as possible. Could you please give me an estimate of how long I should wait? I am really excited about these examples.

[chaoming0625]

Dear @sharadmv , I have tried the Pallas tutorial given in the JAX official documentation (https://jax.readthedocs.io/en/latest/pallas/index.html). But I found I cannot import pallas from the latest jax. How long will the kernel customization support in JAX be online?

[chaoming0625]

Hello, I am wondering if there are any updates? It will be very helpful if there are some replies.

[jakevdp]

Have you tried with the more recent JAX release? It was not possible to use pallas with jax versions older than 0.4.15, but I believe it should be possible now.