Pallas Long Sum Reduction Kernel

[domluna]

def p7_kernel(x_ref, z_ref):  
    @pl.when(pl.program_id(axis=1) == 0)
    def _():
        z_ref[...] = jnp.zeros_like(z_ref)    
    z_ref[...] += jnp.sum(x_ref[...])


@functools.partial(jax.jit, static_argnames=("block_size"))
def p7(
    x: jax.Array,
    *,
    block_size: int = 128,
):
    b, m = x.shape
    grid=(b, m // block_size)
    print(grid)
    return pl.pallas_call(
        p7_kernel,
        in_specs=[
            pl.BlockSpec(lambda i, j: (i,j), (1, block_size,)),
        ],
        out_specs=pl.BlockSpec(lambda i, j: (i,), (1,)),
        grid=grid,
        out_shape=jax.ShapeDtypeStruct((b,), x.dtype),
        interpret=False,
    )(x)

x = jnp.ones((4, 512))
res = p7(x)
print(res)
print(jnp.sum(x, axis=1))

This should work but it's giving this result

[128. 128. 128. 128.]

when interpret=True, it correctly gives

[512. 512. 512. 512.]

am I doing something wrong?

[sharadmv]

Is this on TPU or GPU?

[domluna]

This is on my GPU: GTX 3090.