`nn.Upsample(mode="linear")` is 2-4x slower than a matmul-based bilinear resize when upsampling

Description

mx.nn.Upsample with mode="linear"/align_corners=False is significantly slower than it needs to be specifically when upsampling (scale_factor > 1). I found this while porting a video frame-interpolation model (RIFE) to MLX, where repeated bilinear resizes are a meaningful fraction of the forward pass.

A fixed-ratio bilinear resize is separable into two 1D linear resizes, each of which is just a matrix multiply by a small, fixed (2-nonzeros-per-row) weight matrix — turning "resize" into "matmul" plays to MLX's own strength here. This is bit-exact with nn.Upsample's output (see repro) and 2-4x faster for upsampling shapes; for downsampling nn.Upsample is actually a bit faster than the matmul version, so this looks like a real, specific gap in the upsampling path rather than a case for just always avoiding nn.Upsample.

Repro

"""Repro: mx.nn.Upsample(mode="linear") is 4-8x slower than a
fixed-ratio bilinear resize implemented as two small matmuls, at
ordinary image-pipeline resize shapes. No torch involved - MLX only.
"""
import time
import numpy as np
import mlx.core as mx
import mlx.nn as nn


def matmul_resize_weight(in_size, out_size):
    scale = in_size / out_size
    src = (np.arange(out_size) + 0.5) * scale - 0.5
    src = np.clip(src, 0, in_size - 1)
    x0 = np.floor(src).astype(np.int64)
    x1 = np.minimum(x0 + 1, in_size - 1)
    w1 = (src - x0).astype(np.float32)
    w0 = 1.0 - w1
    w = np.zeros((out_size, in_size), dtype=np.float32)
    w[np.arange(out_size), x0] += w0
    w[np.arange(out_size), x1] += w1
    return mx.array(w)


def matmul_resize(x, scale_factor):
    n, h, w, c = x.shape
    out_h, out_w = round(h * scale_factor), round(w * scale_factor)
    wh, ww = matmul_resize_weight(h, out_h), matmul_resize_weight(w, out_w)
    x = (wh @ x.transpose(1, 0, 2, 3).reshape(h, n * w * c)).reshape(out_h, n, w, c).transpose(1, 0, 2, 3)
    x = (ww @ x.transpose(2, 0, 1, 3).reshape(w, n * out_h * c)).reshape(out_w, n, out_h, c).transpose(1, 2, 0, 3)
    return x


def timed(fn, n=30, warmup=8):
    for _ in range(warmup):
        mx.eval(fn())
    t0 = time.perf_counter()
    for _ in range(n):
        mx.eval(fn())
    return (time.perf_counter() - t0) / n * 1000


for in_h, in_w, scale in [(192, 320, 4.0), (384, 640, 2.0), (768, 1280, 0.25), (768, 1280, 0.5)]:
    x = mx.random.normal((1, in_h, in_w, 4))
    up = nn.Upsample(scale_factor=scale, mode="linear", align_corners=False)
    up_c = mx.compile(lambda a: up(a))
    resize_c = mx.compile(lambda a: matmul_resize(a, scale))

    up_ms = timed(lambda: up_c(x))
    resize_ms = timed(lambda: resize_c(x))
    diff = float(mx.abs(up_c(x) - resize_c(x)).max())
    out_h, out_w = round(in_h * scale), round(in_w * scale)
    print(f"{in_h}x{in_w} -> {out_h}x{out_w} (scale={scale}): "
          f"nn.Upsample={up_ms:.3f}ms  matmul_resize={resize_ms:.3f}ms  "
          f"({up_ms/resize_ms:.1f}x slower)  max_diff={diff:.2e}")

Output (M1 Pro, mlx 0.31.2, macOS 26.4.1)

192x320 -> 768x1280 (scale=4.0): nn.Upsample=4.008ms  matmul_resize=0.979ms  (4.1x slower)  max_diff=4.77e-07
384x640 -> 768x1280 (scale=2.0): nn.Upsample=3.937ms  matmul_resize=1.886ms  (2.1x slower)  max_diff=4.77e-07
768x1280 -> 192x320 (scale=0.25): nn.Upsample=0.509ms  matmul_resize=0.863ms  (0.6x slower)
768x1280 -> 384x640 (scale=0.5): nn.Upsample=1.180ms  matmul_resize=1.659ms  (0.7x slower)

(Both compiled with mx.compile; same pattern holds eager, just uniformly slower on both sides.)

Context

This showed up as a real end-to-end cost: in a RIFE (video frame interpolation) port, replacing nn.Upsample with this matmul-based resize cut ~18ms off a ~354ms forward pass at 768x1280 purely from the upsample calls (each block does one upsample of its flow/mask output back to full resolution). Happy to open a PR with the matmul-based implementation (resize()/_weight_matrix() in the repro) if that's a fix you'd want, or this may point to something fixable at a lower level in however upsample_linear currently walks the output for scale_factor > 1.

For context on where MLX is already doing well on this same machine (so this issue is a targeted gap, not a general performance complaint): a native mx.nn/mx.core GPT (minGPT's gpt-mini config, 2.71M params), compiled with mx.compile's stateful pattern (mx.compile(step, inputs=state, outputs=state)), trains 1.76x faster than the same architecture on PyTorch+MPS (88.7ms/step vs. 154.9ms/step)

• consistent with transformer_lm's existing example using the same compile pattern. The Upsample gap above is the one place in two separate model ports (this GPT, and the RIFE port) where MLX measured slower than PyTorch on identical work.

Environment

• mlx 0.31.2
• macOS 26.4.1, Apple M1 Pro