[Metal] Tune splitk gemm dispatch conditions and partition sizes - #3087

import time
import mlx.core as mx

mx.random.seed(0)


for M, N in [(128, 128), (512, 512), (512, 1024)]:
    for K in [128, 256, 512, 1024, 2048, 4096, 8192, 16384]:
        A = mx.random.normal(shape=(M, K)).astype(mx.bfloat16)
        B = mx.random.normal(shape=(K, N)).astype(mx.bfloat16)

        def fun(A):
            for i in range(60):
                C = A @ B
                # Hack to create a dependency chain so matmuls are in sequence
                A[0, 0] = C[0, 0]
            mx.eval(A)

        for _ in range(10):
            fun(A)

        tic = time.time()
        for i in range(10):
            fun(A)

        toc = time.time()
        s = toc - tic
        ms = 1e3 * (toc - tic)
        print(f"{M=}, {N=}, {K=}, {ms=:.3f}")