Pallas: work around slicing by shifting pointers

[gautierronan]

I am trying to implement a so-called Sparse DIA matrix multiplication kernel, which stores sparse square matrices using their diagonals. Here is a typical implementation in plain JAX:

import jax.numpy as jnp
from jax import Array

def dia_plain(x: Array, y: Array, ks: tuple[int]) -> Array:
    """Plain implementation of a diagonal matrix-vector product."""
    # x has shape (M, N) and represents the matrix diagonals
    # y has shape (N,) and is the vector to multiply against
    z = jnp.zeros_like(y)
    N = y.shape[0]
    for i, k in enumerate(ks):
        xi = x[i, :]
        m, n = max(0, -k), min(N, N - k)
        z = z.at[m:n].add((xi * y)[m + k : n + k])
    return z

which achieves the required matmul, see e.g. this example:

# input examples
x = jnp.array([[0, 1, 2, 3], [1, 2, 3, 0]])
y = jnp.arange(1, 5)
ks = (1, -1)

# corresponding dense matrix
X = jnp.diag(x[0, ks[0]:], k=ks[0]) + jnp.diag(x[1, :ks[1]], k=ks[1])

#     X    @ y  =            =  z
# [0 1 0 0] [1] = [ 2] + [0] = [ 2]
# [1 0 2 0] [2] = [ 6] + [1] = [ 7]
# [0 2 0 3] [3] = [12] + [4] = [16]
# [0 0 3 0] [4] = [ 0] + [9] = [ 9]

print(X @ y)  # [ 2  7 16  9]: correct

# x contains X offset diagonals
#     x    * y  =            =  z
#                 [ 0]
#   [0 1]   [1] = [ 2] +     = [ 2]
#   [1 2]   [2] = [ 6] + [1] = [ 7]
#   [2 3]   [3] = [12] + [4] = [16]
#   [3 0]   [4] =      + [9] = [ 9]
#                        [0]

print(dia_plain(x, y, (1, -1)))  # [2 7 16 9]: correct

When going to Pallas, there is no slicing, so I am trying to workaround it. The workaround I found was to shift pointers by k so that storage in the result array would happen at the right indices. Here is my attempt at a kernel:

import jax
from functools import partial
from jax.experimental import pallas as pl

def dia_kernel(x_ref, y_ref, z_ref, *, ks):
    """Pallas kernel for a diagonal matrix-vector product."""
    z_ref[:] = jnp.zeros_like(y_ref)
    y = y_ref[:]
    for i, k in enumerate(ks):
        xi = x_ref[i, :]
        idx = jnp.arange(y_ref.shape[0]) - k # shift pointers by k
        z_ref[idx] += xi * y


@partial(jax.jit, static_argnums=(2,))
def dia(x: jax.Array, y: jax.Array, ks: tuple[int]) -> jax.Array:
    return pl.pallas_call(
        partial(dia_kernel, ks=ks),
        out_shape=jax.ShapeDtypeStruct(y.shape, y.dtype),
    )(x, y)

However, this does not seem to work. It yields a result without erroring out, but I believe something is wrong with the initialization (the first line z_ref[:] = jnp.zeros_like(y_ref) is not properly taken into account?).

print(dia(x, y, (1, -1)))  # [3 1 4 9]: incorrect

Any idea of how I can achieve this? Any help on how to progress is greatly appreciated. Thanks!

cc @chr1sj0nes @superbobry