Parametrize block size according to input shape in Pallas

[gautierronan]

In Pallas, what is the proper way to parametrize the size of blocks depending on the input matrix dimensions?

For instance, say I have a _matmul function which calls pl.pallas_call(matmul_kernel, ...)(x, y) with inputs x and y of shape (N, N). I want to change my BlockSpec depending on the size N, to get optimal performance. What is the proper way to achieve this?

def matmul_kernel(x_ref, y_ref, z_ref):
    z_ref[...] = x_ref[...] @ y_ref[...]

def _matmul(x: Array, y: Array, *, nblocks: int) -> Array:
    pl.pallas_call(
        matmul_kernel,
        out_shape=jax.ShapeDtypeStruct((x.shape[0], y.shape[1]), x.dtype),
        grid=(nblocks, nblocks),
        in_specs=[
            pl.BlockSpec(lambda i, j: (i, 0), (x.shape[0] // nblocks, x.shape[1])),
            pl.BlockSpec(lambda i, j: (0, j), (y.shape[0], y.shape[1] // nblocks))
        ],
        out_specs=pl.BlockSpec(
            lambda i, j: (i, j), (x.shape[0] // nblocks, y.shape[1] // nblocks)
        )
    )(x, y)

def matmul(x: Array, y: Array) -> Array:
    # What is the right way to write this function?
    N = x.shape[1]
    if N <= 32:
        return _matmul(x, y, nblocks=2)
    elif N <= 64:
        return _matmul(x, y, nblocks=4)
    elif N <= 128:
        return _matmul(x, y, nblocks=8)
    elif N <= 256:
        return _matmul(x, y, nblocks=16)
    else:
        return _matmul(x, y, nblocks=32)

Sorry if this is a newbie question!

[sharadmv]

For an (m, k, n) matmul, I usually parameterize the kernel with (bm: int, bk: int, bn: int). This will make block shapes (bm, bk), (bk, bn), (bm, bn) and a grid of (m // bm, n // bn, k // bn).

I then usually sweep over combinations of bm, bk, bn and pick the best one.

[gautierronan]

Yes if the shape of matrices is fixed in advance it's fine to sweep block sizes, but what if shapes can be changed by the user? I would like to abstract away this "autotuning" complexity. Do you have recommendations in this regard?

[sharadmv]

We don't have any APIs for this at the moment but you have a couple options:

1. write a function wrapper that does autotuning at trace time (i.e. uses input/output shape/types, generates inputs, sweeps over block sizes, times the function, and picks the best one)
2. write a function wrapper that uses a heuristic to choose block sizes given the input/output shapes/types

[gautierronan]

Thanks for the pointers!

Would you have a rough idea of what option 1 would look like? Is this something we want to run at compilation time, or really a wrapper (e.g. an @autotune decorator)? I'm not sure what the overall implementation would look like.

[ASKabalan]

Hello,

I am very interested in this and I started to make a toy autotuner.
For the moment I only have this

def matmul_kernel(x_ref, y_ref, z_ref, *, bk: int):
  m, n = x_ref.shape[0], y_ref.shape[1]
  print(f"m: {m}, n: {n}")
  k = y_ref.shape[0]
  acc = jnp.zeros((m, n), dtype=jnp.float32)
  def body(i, acc):
    x_k = pl.load(x_ref, (slice(None), pl.ds(i * bk, bk)))
    y_k = pl.load(y_ref, (pl.ds(i * bk, bk), slice(None)))
    return acc + x_k @ y_k
  z_ref[...] = lax.fori_loop(0, k // bk, body, acc).astype(z_ref.dtype)


@partial(jax.jit, static_argnames=('debug', 'interpret','bn', 'bk', 'bm'))
def matmul(x: jax.Array, y: jax.Array, *, bm: int = 32, bn: int = 32, bk: int = 32,
           debug: bool = False, interpret: bool = False):
  m, n, k = x.shape[0], y.shape[1], y.shape[0]
  grid = (m // bm, n // bn)
  return pl.pallas_call(
    partial(matmul_kernel, bk=bk),
    out_shape=jax.ShapeDtypeStruct((m, n), x.dtype),
    grid=grid,
    in_specs=[
      pl.BlockSpec(lambda i, j: (i, 0), (bm, k)),
      pl.BlockSpec(lambda i, j: (0, j), (k, bn))
    ],
    out_specs=pl.BlockSpec(lambda i, j: (i, j), (bm, bn)),
    debug=debug, interpret=interpret,
  )(x, y)


def autotune(x , y , c, precision, dtype, output):

  times = {}
  grid_spec = product([16 , 32 , 64, 128, 256], repeat=3)
  for bm , bn , bk in grid_spec:

    print(f"Trying bm={bm}, bn={bn}, bk={bk}")
    try:
      z = matmul(x, y, bm=bm, bn=bn, bk=bk).block_until_ready()
      start = time.perf_counter()
      z = matmul(x, y, bm=bm, bn=bn, bk=bk).block_until_ready()
      end = time.perf_counter()
      elapsed = (end - start)*1000
      times[elapsed] = (bm, bn, bk)
    except Exception as e:
      print(f"Failed with bm={bm}, bn={bn}, bk={bk}")
      print(f"Error: {e}")
      continue

  best_time = min(times.keys())

  for key in times:
    if key == best_time:
      print(f"Best time: {key:.2f} ms, bm={times[key][0]}, bn={times[key][1]}, bk={times[key][2]}")
    else:
      print(f"Time: {key:.2f} ms, bm={times[key][0]}, bn={times[key][1]}, bk={times[key][2]}")

  bm, bn, bk = times[best_time]
  return bm, bn, bk

I did this also in with a wrapped and I believe that the autotune happends in compile/lower time but not trace time as @sharadmv suggested, which means that the autotuning is redone everytime which is suboptimal.
To use the very convinient caching that JAX use for jitted functions, I will try to do what @sharadmv said.
Just to be sure, trace time is abstract evaluation for JAX users right?
In that case we don't have concrete values and we would have to generate random matrices with the tracer shape during the trace time which in my case can be a lot of memory (having a jno.ones matrices does use memory but a lot less than PRNG generated matrices ... is the multiplication of static or jnp.ones matrices enough to autotune?).
Is the JAX team interested in making a triton-like autotuner / do you accept PR in that regard?.

Thank you