[starkAhmed43]

Closes #4435

Related: #4418

This PR speeds up bilinear upsampling on Apple Silicon. It adds a fused Metal kernel for the forward pass and one for the backward pass.

What changed

• A new C++ primitive, exposed as mx.fast.upsample_bilinear(x, scale_factor, align_corners).
• The forward kernel uses one thread per output value. It reads 4 input values and blends them in float32.
• The backward kernel uses one thread per input value. It gathers without atomics, so the result is deterministic.
• nn.Upsample uses the primitive in linear mode for 4-D input without antialiasing.

Results

I measured on a base M1 with 16 GB. The baseline is stock nn.Upsample in MLX 0.32.4. Each number is the median of 10 fresh-process runs.

• All 4 shapes from the issue (upsample and downscale): 6.4–14.8× faster than stock.
• 53 real-model shapes, including RIFE: the kernel is faster on all 53.
• Geometric mean speedup: 16.6× forward, 13.9× forward + backward.
• float16 gradient error: 6.2e-4, against 2.7e-2 for stock.

Upsample shapes from the issue, in ms (F = forward, F+B = forward + backward):

Shape	Scale	Batch, dtype	Pass	Stock	#4418	This PR	vs stock
192×320	×4	1, fp16	F	5.67	14.90	0.48	11.9×
192×320	×4	1, fp16	F+B	11.58	25.21	1.82	6.4×
192×320	×4	8, fp32	F	41.73	94.36	2.91	14.4×
192×320	×4	8, fp32	F+B	88.89	170.07	10.20	8.7×
384×640	×2	1, fp16	F	6.89	14.80	0.49	14.1×
384×640	×2	1, fp16	F+B	13.93	25.43	1.48	9.4×
384×640	×2	8, fp32	F	49.61	92.70	3.36	14.8×
384×640	×2	8, fp32	F+B	111.03	174.07	7.69	14.4×

#4418 is slower than stock on these shapes. It did not trigger on the issue's shapes.

The issue also lists two downscale shapes: 768×1280 ×0.5 and ×0.25. Speedups there are 7.1–11.1× over stock. The full tables below include them.

Full results for all 53 shapes

All times are in ms. Each number is the median of 10 fresh-process runs.

Forward:

Workload	B	Input H×W×C	Scale	align_corners	dtype	Stock ms	#4418 ms	This PR ms	vs stock	vs #4418
deeplabv3p.aspp	1	33×33×256	3.90909	True	fp16	15.6	1.94	0.491	31.86×	3.96×
deeplabv3p.aspp	2	33×33×256	3.90909	True	fp32	26.1	3.05	0.863	30.29×	3.53×
deeplabv3p.logits	1	129×129×21	3.97674	True	fp16	20	20	0.572	35.04×	34.92×
deeplabv3p.logits	2	129×129×21	3.97674	True	fp32	33.6	33.8	1.12	30.07×	30.23×
depth_anything_l.fuse148	1	148×148×256	2	True	fp16	75.6	75.8	1.6	47.10×	47.24×
depth_anything_l.fuse37	1	37×37×256	2	True	fp16	6.04	1.4	0.369	16.38×	3.81×
depth_anything_l.fuse37	4	37×37×256	2	True	fp32	17.5	3.01	0.747	23.43×	4.03×
depth_anything_l.fuse74	1	74×74×256	2	True	fp16	20.1	3.12	0.554	36.17×	5.63×
depth_anything_l.fuse74	4	74×74×256	2	True	fp32	66.6	9.36	2.41	27.66×	3.89×
depth_anything_s.fuse148	1	148×148×64	2	True	fp16	20.1	20.1	0.552	36.43×	36.36×
depth_anything_s.fuse148	4	148×148×64	2	True	fp32	67.1	67.2	2.46	27.33×	27.36×
depth_anything_s.fuse37	1	37×37×64	2	True	fp16	2.28	0.88	0.274	8.33×	3.21×
depth_anything_s.fuse37	4	37×37×64	2	True	fp32	5.79	1.61	0.381	15.19×	4.23×
depth_anything_s.fuse74	1	74×74×64	2	True	fp16	6.12	1.58	0.4	15.31×	3.96×
depth_anything_s.fuse74	4	74×74×64	2	True	fp32	17.5	3.15	0.752	23.32×	4.19×
pspnet.bin1	1	1×1×512	60	True	fp16	6.7	0.529	0.406	16.50×	1.30×
pspnet.bin1	2	1×1×512	60	True	fp32	12.2	1.14	0.479	25.39×	2.38×
pspnet.bin2	1	2×2×512	30	True	fp16	6.71	0.584	0.411	16.34×	1.42×
pspnet.bin2	2	2×2×512	30	True	fp32	11.9	0.965	0.499	23.80×	1.93×
pspnet.bin3	1	3×3×512	20	True	fp16	6.85	0.598	0.405	16.93×	1.48×
pspnet.bin3	2	3×3×512	20	True	fp32	11.9	1.03	0.487	24.54×	2.12×
pspnet.bin6	1	6×6×512	10	True	fp16	7.01	0.674	0.405	17.32×	1.67×
pspnet.bin6	2	6×6×512	10	True	fp32	11.8	1.1	0.481	24.54×	2.29×
rife.down2	1	768×1280×4	0.5	False	fp16	4.15	4.56	0.409	10.13×	11.14×
rife.down2	8	768×1280×4	0.5	False	fp32	25.7	23.9	3.06	8.38×	7.82×
rife.down4	1	768×1280×4	0.25	False	fp16	2.39	2.09	0.258	9.26×	8.12×
rife.down4	8	768×1280×4	0.25	False	fp32	11.7	6.92	1.62	7.23×	4.28×
rife.up2	1	384×640×4	2	False	fp16	6.89	14.8	0.488	14.10×	30.30×
rife.up2	8	384×640×4	2	False	fp32	49.6	92.7	3.36	14.75×	27.56×
rife.up4	1	192×320×4	4	False	fp16	5.67	14.9	0.478	11.85×	31.16×
rife.up4	8	192×320×4	4	False	fp32	41.7	94.4	2.91	14.36×	32.47×
sam.masks	1	256×256×3	4	False	fp16	4.73	12.2	0.456	10.37×	26.79×
sam.masks	4	256×256×3	4	False	fp32	17.3	37.8	1.32	13.11×	28.64×
segformer.logits	1	128×128×150	4	False	fp16	49.6	10.7	2.19	22.65×	4.91×
segformer_b0.mlp16	1	16×16×256	8	False	fp16	5.54	1.22	0.469	11.82×	2.60×
segformer_b0.mlp16	2	16×16×256	8	False	fp32	10.3	1.95	0.817	12.58×	2.39×
segformer_b0.mlp32	1	32×32×256	4	False	fp16	6.22	1.68	0.48	12.96×	3.50×
segformer_b0.mlp32	2	32×32×256	4	False	fp32	11.3	2.41	0.851	13.27×	2.83×
segformer_b0.mlp64	1	64×64×256	2	False	fp16	7.3	2.13	0.489	14.91×	4.35×
segformer_b0.mlp64	2	64×64×256	2	False	fp32	13.4	3.38	1.05	12.78×	3.23×
segformer_b2.mlp16	1	16×16×768	8	False	fp16	14.6	2.5	0.813	17.96×	3.07×
segformer_b2.mlp16	2	16×16×768	8	False	fp32	28.7	4.01	2.08	13.76×	1.93×
segformer_b2.mlp32	1	32×32×768	4	False	fp16	16.3	3.03	0.814	20.04×	3.72×
segformer_b2.mlp32	2	32×32×768	4	False	fp32	31.7	5.13	2.22	14.26×	2.30×
segformer_b2.mlp64	1	64×64×768	2	False	fp16	20	4.63	0.884	22.66×	5.23×
segformer_b2.mlp64	2	64×64×768	2	False	fp32	38.6	8.8	2.58	15.00×	3.42×
upernet.fpn16	1	16×16×512	2	False	fp16	1.33	0.675	0.293	4.54×	2.30×
upernet.fpn16	2	16×16×512	2	False	fp32	2.27	0.85	0.318	7.13×	2.68×
upernet.fpn32	1	32×32×512	2	False	fp16	4.55	1.4	0.405	11.23×	3.45×
upernet.fpn32	2	32×32×512	2	False	fp32	7.18	2.1	0.59	12.17×	3.56×
upernet.fpn64	1	64×64×512	2	False	fp16	13.8	3.44	0.646	21.35×	5.33×
upernet.fpn64	2	64×64×512	2	False	fp32	26.3	5.91	1.86	14.18×	3.19×
upernet.logits	1	128×128×150	4	False	fp16	49.2	10.7	2.2	22.33×	4.87×

Forward + backward:

Workload	B	Input H×W×C	Scale	align_corners	dtype	Stock ms	#4418 ms	This PR ms	vs stock	vs #4418
deeplabv3p.aspp	1	33×33×256	3.90909	True	fp16	27.3	3.67	1.3	21.04×	2.83×
deeplabv3p.aspp	2	33×33×256	3.90909	True	fp32	46.8	6.05	1.76	26.55×	3.44×
deeplabv3p.logits	1	129×129×21	3.97674	True	fp16	36.4	36.7	1.31	27.74×	28.02×
deeplabv3p.logits	2	129×129×21	3.97674	True	fp32	60.9	60.5	2.36	25.79×	25.62×
depth_anything_l.fuse148	1	148×148×256	2	True	fp16	137	137	3.7	37.05×	37.01×
depth_anything_l.fuse37	1	37×37×256	2	True	fp16	10.2	2.36	0.779	13.11×	3.03×
depth_anything_l.fuse37	4	37×37×256	2	True	fp32	32.8	5.51	1.55	21.21×	3.56×
depth_anything_l.fuse74	1	74×74×256	2	True	fp16	35.7	6.32	1.51	23.67×	4.19×
depth_anything_l.fuse74	4	74×74×256	2	True	fp32	128	21.4	4.55	28.12×	4.71×
depth_anything_s.fuse148	1	148×148×64	2	True	fp16	35.7	35.8	1.42	25.10×	25.19×
depth_anything_s.fuse148	4	148×148×64	2	True	fp32	127	127	4.66	27.38×	27.39×
depth_anything_s.fuse37	1	37×37×64	2	True	fp16	4.27	1.38	0.369	11.57×	3.75×
depth_anything_s.fuse37	4	37×37×64	2	True	fp32	9.59	2.31	0.789	12.16×	2.93×
depth_anything_s.fuse74	1	74×74×64	2	True	fp16	10.1	2.49	0.794	12.76×	3.13×
depth_anything_s.fuse74	4	74×74×64	2	True	fp32	33	5.96	1.55	21.33×	3.86×
pspnet.bin1	1	1×1×512	60	True	fp16	14.4	1.64	1.65	8.73×	0.99×
pspnet.bin1	2	1×1×512	60	True	fp32	20.5	2.25	1.94	10.53×	1.16×
pspnet.bin2	1	2×2×512	30	True	fp16	20.9	1.65	1.64	12.78×	1.01×
pspnet.bin2	2	2×2×512	30	True	fp32	21.5	2.46	1.89	11.43×	1.30×
pspnet.bin3	1	3×3×512	20	True	fp16	15.8	1.76	1.5	10.47×	1.17×
pspnet.bin3	2	3×3×512	20	True	fp32	21.1	2.52	1.87	11.26×	1.35×
pspnet.bin6	1	6×6×512	10	True	fp16	13.6	1.89	0.807	16.83×	2.34×
pspnet.bin6	2	6×6×512	10	True	fp32	20.9	2.64	1.26	16.60×	2.10×
rife.down2	1	768×1280×4	0.5	False	fp16	8.48	9.12	1.04	8.12×	8.74×
rife.down2	8	768×1280×4	0.5	False	fp32	68.7	72.1	6.17	11.14×	11.69×
rife.down4	1	768×1280×4	0.25	False	fp16	5.26	5.19	0.74	7.10×	7.00×
rife.down4	8	768×1280×4	0.25	False	fp32	36.9	40.8	4.19	8.81×	9.73×
rife.up2	1	384×640×4	2	False	fp16	13.9	25.4	1.48	9.43×	17.21×
rife.up2	8	384×640×4	2	False	fp32	111	174	7.69	14.44×	22.64×
rife.up4	1	192×320×4	4	False	fp16	11.6	25.2	1.82	6.36×	13.84×
rife.up4	8	192×320×4	4	False	fp32	88.9	170	10.2	8.71×	16.67×
sam.masks	1	256×256×3	4	False	fp16	9.53	20.4	1.69	5.63×	12.08×
sam.masks	4	256×256×3	4	False	fp32	35.9	68.1	4.84	7.42×	14.08×
segformer.logits	1	128×128×150	4	False	fp16	105	25.3	4.99	20.98×	5.07×
segformer_b0.mlp16	1	16×16×256	8	False	fp16	11	2.87	1.26	8.74×	2.27×
segformer_b0.mlp16	2	16×16×256	8	False	fp32	21.2	4.51	1.73	12.27×	2.61×
segformer_b0.mlp32	1	32×32×256	4	False	fp16	12.3	3.27	1.12	10.94×	2.92×
segformer_b0.mlp32	2	32×32×256	4	False	fp32	23.8	5.21	1.74	13.67×	3.00×
segformer_b0.mlp64	1	64×64×256	2	False	fp16	14.6	4.21	1.31	11.18×	3.21×
segformer_b0.mlp64	2	64×64×256	2	False	fp32	29.5	6.87	1.96	15.04×	3.50×
segformer_b2.mlp16	1	16×16×768	8	False	fp16	30.4	6.1	1.82	16.76×	3.36×
segformer_b2.mlp16	2	16×16×768	8	False	fp32	60.9	11	4.03	15.11×	2.72×
segformer_b2.mlp32	1	32×32×768	4	False	fp16	34	7.35	1.91	17.80×	3.84×
segformer_b2.mlp32	2	32×32×768	4	False	fp32	69.2	13.4	4.3	16.08×	3.12×
segformer_b2.mlp64	1	64×64×768	2	False	fp16	41.4	10.6	2.15	19.30×	4.92×
segformer_b2.mlp64	2	64×64×768	2	False	fp32	86.6	19.4	5.06	17.11×	3.84×
upernet.fpn16	1	16×16×512	2	False	fp16	3.31	1.5	0.402	8.24×	3.73×
upernet.fpn16	2	16×16×512	2	False	fp32	5.78	2	0.518	11.17×	3.86×
upernet.fpn32	1	32×32×512	2	False	fp16	8.53	2.64	0.936	9.12×	2.82×
upernet.fpn32	2	32×32×512	2	False	fp32	15.3	4	1.42	10.80×	2.82×
upernet.fpn64	1	64×64×512	2	False	fp16	27.9	7.32	1.55	17.94×	4.71×
upernet.fpn64	2	64×64×512	2	False	fp32	57.3	12.9	3.53	16.21×	3.66×
upernet.logits	1	128×128×150	4	False	fp16	104	25.2	4.99	20.76×	5.05×

Limits

• The kernel runs only on Apple Silicon (Metal), for 4-D input without antialiasing.
• Other cases use a fallback built from ordinary MLX ops. These are CPU, CUDA, empty input, input or output with 2³¹ or more elements, and vmap/jvp. The fallback gives the same values and no speedup.
• Antialias, cubic and nearest modes, and 3-D and 5-D input, are unchanged.
• Each new shape costs about 6 ms once, to compile its pipeline.

Tests

I added 13 tests in test_fast.py, plus an export round trip.

Questions for maintainers

1. Should mx.fast.upsample_bilinear stay public? The other option is to make it internal and use it only through nn.Upsample.
2. For float16 input, my kernel returns float16. Stock returns float32 as a side effect of its float32 weights. Is float16 out acceptable?

☑️ I understand it is strictly prohibited to use AI to write PR description

AI usage disclosure:
I wrote the Metal kernels. An AI assistant reviewed them for errors and refactored them for readability and wrote the C++ primitive, the Python binding, the nn.Upsample change and the tests. It also wrote the benchmark and analysis scripts. I set the direction, made the design decisions and reviewed the changes and results.