Timing difference between mx.eval and mx.async_eval

[sck-at-ucy]

Ok, continuing from the X discussion. I know I must be doing something wrong here, but I just don't see it. Here is the code that produces the large timing difference:

import mlx.core as mx
import time
from functools import partial

# Set the default device to GPU for all operations
mx.set_default_device(mx.gpu)

# Define dimensions
nxR, nyR, nzR = 500, 500, 100
dx = 1

# Create random arrays using MLX
mx.random.seed(1)
u = mx.random.uniform(shape=(2 * nxR + 5, 2 * nyR + 5, 2 * nzR + 5))
v = mx.random.uniform(shape=(2 * nxR + 5, 2 * nyR + 5, 2 * nzR + 5))
w = mx.random.uniform(shape=(2 * nxR + 5, 2 * nyR + 5, 2 * nzR + 5))

# Define the convolution function for MLX arrays with lazy evaluation
@partial(mx.compile, inputs=[u, v, w], outputs=["C_u", "C_v", "C_w"])
def conv_mlx(u, v, w, nxR, nyR, nzR, dx):
    # Implement the C_u calculation using MLX-compatible slicing
    C_u = (
        u[4:2*nxR+4:2, 3:2*nyR+3:2, 3:2*nzR+3:2] *
        (u[6:2*nxR+6:2, 3:2*nyR+3:2, 3:2*nzR+3:2] - u[2:2*nxR+2:2, 3:2*nyR+3:2, 3:2*nzR+3:2]) / (2 * dx) +
        (v[3:2*nxR+3:2, 4:2*nyR+4:2, 3:2*nzR+3:2] +
         v[5:2*nxR+5:2, 4:2*nyR+4:2, 3:2*nzR+3:2] +
         v[3:2*nxR+3:2, 2:2*nyR+2:2, 3:2*nzR+3:2] +
         v[5:2*nxR+5:2, 2:2*nyR+2:2, 3:2*nzR+3:2]) *
        (u[4:2*nxR+4:2, 5:2*nyR+5:2, 3:2*nzR+3:2] - u[4:2*nxR+4:2, 1:2*nyR+1:2, 3:2*nzR+3:2]) / (8 * dx) +
        (w[3:2*nxR+3:2, 3:2*nyR+3:2, 4:2*nzR+4:2] +
         w[3:2*nxR+3:2, 3:2*nyR+3:2, 2:2*nzR+2:2] +
         w[5:2*nxR+5:2, 3:2*nyR+3:2, 4:2*nzR+4:2] +
         w[5:2*nxR+5:2, 3:2*nyR+3:2, 2:2*nzR+2:2]) *
        (u[4:2*nxR+4:2, 3:2*nyR+3:2, 5:2*nzR+5:2] - u[4:2*nxR+4:2, 3:2*nyR+3:2, 1:2*nzR+1:2]) / (8 * dx)
    )

    # Implement the C_v calculation
    C_v = (
        (u[4:2*nxR+4:2, 3:2*nyR+3:2, 3:2*nzR+3:2] +
         u[4:2*nxR+4:2, 5:2*nyR+5:2, 3:2*nzR+3:2] +
         u[2:2*nxR+2:2, 3:2*nyR+3:2, 3:2*nzR+3:2] +
         u[2:2*nxR+2:2, 5:2*nyR+5:2, 3:2*nzR+3:2]) *
        (v[5:2*nxR+5:2, 4:2*nyR+4:2, 3:2*nzR+3:2] - v[1:2*nxR+1:2, 4:2*nyR+4:2, 3:2*nzR+3:2]) / (8 * dx) +
        v[3:2*nxR+3:2, 4:2*nyR+4:2, 3:2*nzR+3:2] *
        (v[3:2*nxR+3:2, 6:2*nyR+6:2, 3:2*nzR+3:2] - v[3:2*nxR+3:2, 2:2*nyR+2:2, 3:2*nzR+3:2]) / (2 * dx) +
        (w[3:2*nxR+3:2, 3:2*nyR+3:2, 4:2*nzR+4:2] +
         w[3:2*nxR+3:2, 5:2*nyR+5:2, 4:2*nzR+4:2] +
         w[3:2*nxR+3:2, 3:2*nyR+3:2, 2:2*nzR+2:2] +
         w[3:2*nxR+3:2, 5:2*nyR+5:2, 2:2*nzR+2:2]) *
        (v[3:2*nxR+3:2, 4:2*nyR+4:2, 5:2*nzR+5:2] - v[3:2*nxR+3:2, 4:2*nyR+4:2, 1:2*nzR+1:2]) / (8 * dx)
    )

    # Implement the C_w calculation
    C_w = (
        (u[4:2*nxR+4:2, 3:2*nyR+3:2, 3:2*nzR+3:2] +
         u[4:2*nxR+4:2, 3:2*nyR+3:2, 5:2*nzR+5:2] +
         u[2:2*nxR+2:2, 3:2*nyR+3:2, 3:2*nzR+3:2] +
         u[2:2*nxR+2:2, 3:2*nyR+3:2, 5:2*nzR+5:2]) *
        (w[5:2*nxR+5:2, 3:2*nyR+3:2, 4:2*nzR+4:2] - w[1:2*nxR+1:2, 3:2*nyR+3:2, 4:2*nzR+4:2]) / (8 * dx) +
        (v[3:2*nxR+3:2, 4:2*nyR+4:2, 3:2*nzR+3:2] +
         v[3:2*nxR+3:2, 4:2*nyR+4:2, 5:2*nzR+5:2] +
         v[3:2*nxR+3:2, 2:2*nyR+2:2, 3:2*nzR+3:2] +
         v[3:2*nxR+3:2, 2:2*nyR+2:2, 5:2*nzR+5:2]) *
        (w[3:2*nxR+3:2, 5:2*nyR+5:2, 4:2*nzR+4:2] - w[3:2*nxR+3:2, 1:2*nyR+1:2, 4:2*nzR+4:2]) / (8 * dx) +
        w[3:2*nxR+3:2, 3:2*nyR+3:2, 4:2*nzR+4:2] *
        (w[3:2*nxR+3:2, 3:2*nyR+3:2, 6:2*nzR+6:2] - w[3:2*nxR+3:2, 3:2*nyR+3:2, 2:2*nzR+2:2]) / (2 * dx)
    )

    return C_u, C_v, C_w

# Run the function and measure GPU execution time
start_time = time.time()
C_u_mlx_S, C_v_mlx_S, C_w_mlx_S = conv_mlx(u, v, w, nxR, nyR, nzR, dx)
mx.eval(C_u_mlx_S)  # Trigger evaluation
elapsed_time_gpu = time.time() - start_time

print(f"Elapsed time on M2 GPU with mx.eval: {elapsed_time_gpu} seconds")


# Run the function and measure GPU execution time
start_time = time.time()
C_u_mlx_A, C_v_mlx_A, C_w_mlx_A = conv_mlx(u, v, w, nxR, nyR, nzR, dx)
mx.async_eval(C_u_mlx_A)  # Trigger evaluation
mx.synchronize()
elapsed_time_gpu = time.time() - start_time

print(f"Elapsed time on M2 GPU with MLX async_eval: {elapsed_time_gpu} seconds")


print(C_u_mlx_A-C_u_mlx_S)

When I run this I get:

Elapsed time on M2 GPU with mx.eval: 0.15218305587768555 seconds
Elapsed time on M2 GPU with MLX async_eval: 0.006497859954833984 seconds

[awni]

The time difference is mostly due to the fact that there is no warmup. The first time you call a compiled function it gets compiled, which will make it slower. After that, the function will be a lot faster.

Also I strongly recommend measuring more than one call to the function to reduce the variance.

[sck-at-ucy]

Yes, confirmed! This is the output with multiple mx.eval (no mx.async_eval). The first time around much slower, then fast. Curious observation, the second call is always a little slower than the rest as well, no matter how many times I run the code.

Elapsed time on M2 GPU with mx.eval #1: 0.14909100532531738 seconds
Elapsed time on M2 GPU with MLX mx.eval #2: 0.006442070007324219 seconds
Elapsed time on M2 GPU with MLX mx.eval #3: 0.0025229454040527344 seconds
Elapsed time on M2 GPU with MLX mx.eval #4: 0.002500772476196289 seconds
Elapsed time on M2 GPU with MLX mx.eval #5: 0.0024862289428710938 seconds