Pallas autotuning

[sharadmv]

Enables autotuning in Pallas

Example:

@functools.partial(
  pl.pallas_call, out_shape=jax.ShapeDtypeStruct((8,), jnp.float32),
  autotuning_configs=[
    pl.KernelConfig(meta=dict(block_size=block_size),
                    in_specs=[
                      pl.BlockSpec(lambda i: i, (block_size,))
                    ],
                    out_specs=[
                      pl.BlockSpec(lambda i: i, (block_size,))
                    ],
                    grid=8 // block_size)
  for block_size in [1, 2, 4, 8]
])
def add_one(x_ref, o_ref, *, block_size):
  del block_size
  o_ref[...] = x_ref[...] + 1.