JIT compilation prevents optimization of GEMM with bias

[zcbenz]

Consider following code, which is partial code of model training:

from functools import partial

import mlx.core as mx
import mlx.nn as nn
from mlx.utils import tree_map

model = nn.Linear(1024, 3072, bias=True)
params = model.trainable_parameters()
mx.eval(params)

state = [mx.random.state]

@partial(mx.compile, inputs=state, outputs=state)
def step(x, params):
    model.update(tree_map(lambda x: x.astype(mx.bfloat16), params))
    return model(x.astype(mx.float16))

x = mx.random.uniform(shape=(4,512,1024), dtype=mx.bfloat16)
y = step(x, params)
mx.eval(y)

Which converts the Linear module's weights and bias from f32 to bf16, broadcasts the bias to (4,512,3072), and then calls addmm. And after JIT compilation the astype and broadcast would be fused together.

The problem is that with JIT compilation we would lose some optimization opportunities: if we remove the mx.compile line, the AddMM::eval_gpu would receive the bias input with data size of 3072, from which we know it is a bias vector to be added to the last dim of output, and can use a faster kernel (#2569). However with JIT compilation the bias input's data size becomes 6291456 (4 *512 *3072) because of op fusion, and we can no longer take the same optimization.

I can think of 2 solutions:

1. Add bias epilogue as part of Matmul primitive and make addmm redirect to it. The downside is it would be relatively large change and we would also need to apply the same optimization to Metal kernels.
2. Move the broardcasting and reshaping of bias input from addmm op to AddMM::eval_gpu, the change would be minimal and could also be a bit ugly.

And with both solutions I think the vjp/jvp would be a bit tricky to get right.

[awni]

That's a very interesting observation. It seems suboptimal to me that we fuse broadcast when it's the final op.

For example, if you have a chain of elementwise ops that end with a broadcast:

import time
import mlx.core as mx

@mx.compile
def fun(x):
    return mx.broadcast_to(mx.sin(mx.exp(mx.log(mx.abs(x)))), (16384, 2048, 16))

x = mx.array([1.0])
for _ in range(10):
    x = fun(x)
mx.eval(x)

tic = time.time()
x = mx.array([1.0])
for _ in range(10):
    x = fun(x)
mx.eval(x)
toc = time.time()
print(toc - tic)

Time the above with and without compile. Compile makes it way way slower:

> python test.py
0.16287493705749512

> MLX_DISABLE_COMPILE=1 python tt.py 
0.0005729198455810547

[awni]

I think one way to solve this problem is to allow compile to produce non-row-contiguous outputs.

Maybe something like: if all the inputs are contiguous with the same strides the output can inherit the strides of the input and we can use a contiguous kernel to compute the output.

[angeloskath]

I am closing this based on #2622. I think having non-row-contiguous inputs generate non-row-contiguous outputs but run the contiguous kernel is a nice optimization that is independent of fusing the last broadcast.