grid_sampler_3d for MPS #159421

mlaves · 2025-07-29T21:24:36Z

Metal implementation of grid_sampler_3d, including backward pass.

Batch size:  1, CPU time: 7.164317 ms, MPS time: 0.817350 ms, speedup: 8.77x
Batch size:  2, CPU time: 14.793888 ms, MPS time: 1.111496 ms, speedup: 13.31x
Batch size:  4, CPU time: 30.779025 ms, MPS time: 1.712583 ms, speedup: 17.97x
Batch size:  8, CPU time: 68.820854 ms, MPS time: 3.391506 ms, speedup: 20.29x
Batch size: 16, CPU time: 142.618127 ms, MPS time: 5.831585 ms, speedup: 24.46x

Output accuracy test vs. CPU

Click to expand

Grid Sampler 3D Output Accuracy Test (MPS vs CPU)
input_shape           grid_shape         interp   padding   align_corners   max_diff      mean_diff     sum_diff
------------------------------------------------------------------------------------------------------------------------
(1, 3, 128, 128, 128) (1, 64, 64, 64, 3)   0        0         0   0.00000000   0.00000000   0.00000000
(1, 3, 128, 128, 128) (1, 64, 64, 64, 3)   0        0         1   0.00000000   0.00000000   0.00000000
(1, 3, 128, 128, 128) (1, 64, 64, 64, 3)   0        1         0   0.00000000   0.00000000   0.00000000
(1, 3, 128, 128, 128) (1, 64, 64, 64, 3)   0        1         1   0.00000000   0.00000000   0.00000000
(1, 3, 128, 128, 128) (1, 64, 64, 64, 3)   0        2         0   0.00000000   0.00000000   0.00000000
(1, 3, 128, 128, 128) (1, 64, 64, 64, 3)   0        2         1   0.00000000   0.00000000   0.00000000
(1, 3, 128, 128, 128) (1, 64, 64, 64, 3)   1        0         0   0.00000000   0.00000000   0.00000000
(1, 3, 128, 128, 128) (1, 64, 64, 64, 3)   1        0         1   0.00000000   0.00000000   0.00000000
(1, 3, 128, 128, 128) (1, 64, 64, 64, 3)   1        1         0   0.00000000   0.00000000   0.00000000
(1, 3, 128, 128, 128) (1, 64, 64, 64, 3)   1        1         1   0.00000000   0.00000000   0.00000000
(1, 3, 128, 128, 128) (1, 64, 64, 64, 3)   1        2         0   0.00000000   0.00000000   0.00000000
(1, 3, 128, 128, 128) (1, 64, 64, 64, 3)   1        2         1   0.00000000   0.00000000   0.00000000
(1, 3, 128, 128, 128) (1, 1, 1, 1, 3)      0        0         0   0.00000000   0.00000000   0.00000000
(1, 3, 128, 128, 128) (1, 1, 1, 1, 3)      0        0         1   0.00000000   0.00000000   0.00000000
(1, 3, 128, 128, 128) (1, 1, 1, 1, 3)      0        1         0   0.00000000   0.00000000   0.00000000
(1, 3, 128, 128, 128) (1, 1, 1, 1, 3)      0        1         1   0.00000000   0.00000000   0.00000000
(1, 3, 128, 128, 128) (1, 1, 1, 1, 3)      0        2         0   0.00000000   0.00000000   0.00000000
(1, 3, 128, 128, 128) (1, 1, 1, 1, 3)      0        2         1   0.00000000   0.00000000   0.00000000
(1, 3, 128, 128, 128) (1, 1, 1, 1, 3)      1        0         0   0.00000000   0.00000000   0.00000000
(1, 3, 128, 128, 128) (1, 1, 1, 1, 3)      1        0         1   0.00000000   0.00000000   0.00000000
(1, 3, 128, 128, 128) (1, 1, 1, 1, 3)      1        1         0   0.00000000   0.00000000   0.00000000
(1, 3, 128, 128, 128) (1, 1, 1, 1, 3)      1        1         1   0.00000000   0.00000000   0.00000000
(1, 3, 128, 128, 128) (1, 1, 1, 1, 3)      1        2         0   0.00000000   0.00000000   0.00000000
(1, 3, 128, 128, 128) (1, 1, 1, 1, 3)      1        2         1   0.00000000   0.00000000   0.00000000
(1, 3, 128, 128, 128) (1, 8, 8, 8, 3)      0        0         0   0.00000000   0.00000000   0.00000000
(1, 3, 128, 128, 128) (1, 8, 8, 8, 3)      0        0         1   0.00000000   0.00000000   0.00000000
(1, 3, 128, 128, 128) (1, 8, 8, 8, 3)      0        1         0   0.00000000   0.00000000   0.00000000
(1, 3, 128, 128, 128) (1, 8, 8, 8, 3)      0        1         1   0.00000000   0.00000000   0.00000000
(1, 3, 128, 128, 128) (1, 8, 8, 8, 3)      0        2         0   0.00000000   0.00000000   0.00000000
(1, 3, 128, 128, 128) (1, 8, 8, 8, 3)      0        2         1   0.00000000   0.00000000   0.00000000
(1, 3, 128, 128, 128) (1, 8, 8, 8, 3)      1        0         0   0.00000000   0.00000000   0.00000000
(1, 3, 128, 128, 128) (1, 8, 8, 8, 3)      1        0         1   0.00000000   0.00000000   0.00000000
(1, 3, 128, 128, 128) (1, 8, 8, 8, 3)      1        1         0   0.00000000   0.00000000   0.00000000
(1, 3, 128, 128, 128) (1, 8, 8, 8, 3)      1        1         1   0.00000000   0.00000000   0.00000000
(1, 3, 128, 128, 128) (1, 8, 8, 8, 3)      1        2         0   0.00000000   0.00000000   0.00000000
(1, 3, 128, 128, 128) (1, 8, 8, 8, 3)      1        2         1   0.00000000   0.00000000   0.00000000
(2, 1, 32, 32, 32)    (2, 16, 16, 16, 3)   0        0         0   0.00000000   0.00000000   0.00000000
(2, 1, 32, 32, 32)    (2, 16, 16, 16, 3)   0        0         1   0.00000000   0.00000000   0.00000000
(2, 1, 32, 32, 32)    (2, 16, 16, 16, 3)   0        1         0   0.00000000   0.00000000   0.00000000
(2, 1, 32, 32, 32)    (2, 16, 16, 16, 3)   0        1         1   0.00000000   0.00000000   0.00000000
(2, 1, 32, 32, 32)    (2, 16, 16, 16, 3)   0        2         0   0.00000000   0.00000000   0.00000000
(2, 1, 32, 32, 32)    (2, 16, 16, 16, 3)   0        2         1   0.00000000   0.00000000   0.00000000
(2, 1, 32, 32, 32)    (2, 16, 16, 16, 3)   1        0         0   0.00000000   0.00000000   0.00000000
(2, 1, 32, 32, 32)    (2, 16, 16, 16, 3)   1        0         1   0.00000000   0.00000000   0.00000000
(2, 1, 32, 32, 32)    (2, 16, 16, 16, 3)   1        1         0   0.00000000   0.00000000   0.00000000
(2, 1, 32, 32, 32)    (2, 16, 16, 16, 3)   1        1         1   0.00000000   0.00000000   0.00000000
(2, 1, 32, 32, 32)    (2, 16, 16, 16, 3)   1        2         0   0.00000000   0.00000000   0.00000000
(2, 1, 32, 32, 32)    (2, 16, 16, 16, 3)   1        2         1   0.00000000   0.00000000   0.00000000
(1, 1, 1, 1, 1)       (1, 64, 64, 64, 3)   0        0         0   0.00000000   0.00000000   0.00000000
(1, 1, 1, 1, 1)       (1, 64, 64, 64, 3)   0        0         1   0.00000000   0.00000000   0.00000000
(1, 1, 1, 1, 1)       (1, 64, 64, 64, 3)   0        1         0   0.00000000   0.00000000   0.00000000
(1, 1, 1, 1, 1)       (1, 64, 64, 64, 3)   0        1         1   0.00000000   0.00000000   0.00000000
(1, 1, 1, 1, 1)       (1, 64, 64, 64, 3)   0        2         0   0.00000000   0.00000000   0.00000000
(1, 1, 1, 1, 1)       (1, 64, 64, 64, 3)   0        2         1   0.00000000   0.00000000   0.00000000
(1, 1, 1, 1, 1)       (1, 64, 64, 64, 3)   1        0         0   0.00000000   0.00000000   0.00000000
(1, 1, 1, 1, 1)       (1, 64, 64, 64, 3)   1        0         1   0.00000000   0.00000000   0.00000000
(1, 1, 1, 1, 1)       (1, 64, 64, 64, 3)   1        1         0   0.00000000   0.00000000   0.00000000
(1, 1, 1, 1, 1)       (1, 64, 64, 64, 3)   1        1         1   0.00000000   0.00000000   0.00000000
(1, 1, 1, 1, 1)       (1, 64, 64, 64, 3)   1        2         0   0.00000000   0.00000000   0.00000000
(1, 1, 1, 1, 1)       (1, 64, 64, 64, 3)   1        2         1   0.00000000   0.00000000   0.00000000
(1, 1, 1, 1, 1)       (1, 1, 1, 1, 3)      0        0         0   0.00000000   0.00000000   0.00000000
(1, 1, 1, 1, 1)       (1, 1, 1, 1, 3)      0        0         1   0.00000000   0.00000000   0.00000000
(1, 1, 1, 1, 1)       (1, 1, 1, 1, 3)      0        1         0   0.00000000   0.00000000   0.00000000
(1, 1, 1, 1, 1)       (1, 1, 1, 1, 3)      0        1         1   0.00000000   0.00000000   0.00000000
(1, 1, 1, 1, 1)       (1, 1, 1, 1, 3)      0        2         0   0.00000000   0.00000000   0.00000000
(1, 1, 1, 1, 1)       (1, 1, 1, 1, 3)      0        2         1   0.00000000   0.00000000   0.00000000
(1, 1, 1, 1, 1)       (1, 1, 1, 1, 3)      1        0         0   0.00000000   0.00000000   0.00000000
(1, 1, 1, 1, 1)       (1, 1, 1, 1, 3)      1        0         1   0.00000000   0.00000000   0.00000000
(1, 1, 1, 1, 1)       (1, 1, 1, 1, 3)      1        1         0   0.00000000   0.00000000   0.00000000
(1, 1, 1, 1, 1)       (1, 1, 1, 1, 3)      1        1         1   0.00000000   0.00000000   0.00000000
(1, 1, 1, 1, 1)       (1, 1, 1, 1, 3)      1        2         0   0.00000000   0.00000000   0.00000000
(1, 1, 1, 1, 1)       (1, 1, 1, 1, 3)      1        2         1   0.00000000   0.00000000   0.00000000
(1, 1, 1, 1, 1)       (1, 8, 8, 8, 3)      0        0         0   0.00000000   0.00000000   0.00000000
(1, 1, 1, 1, 1)       (1, 8, 8, 8, 3)      0        0         1   0.00000000   0.00000000   0.00000000
(1, 1, 1, 1, 1)       (1, 8, 8, 8, 3)      0        1         0   0.00000000   0.00000000   0.00000000
(1, 1, 1, 1, 1)       (1, 8, 8, 8, 3)      0        1         1   0.00000000   0.00000000   0.00000000
(1, 1, 1, 1, 1)       (1, 8, 8, 8, 3)      0        2         0   0.00000000   0.00000000   0.00000000
(1, 1, 1, 1, 1)       (1, 8, 8, 8, 3)      0        2         1   0.00000000   0.00000000   0.00000000
(1, 1, 1, 1, 1)       (1, 8, 8, 8, 3)      1        0         0   0.00000000   0.00000000   0.00000000
(1, 1, 1, 1, 1)       (1, 8, 8, 8, 3)      1        0         1   0.00000000   0.00000000   0.00000000
(1, 1, 1, 1, 1)       (1, 8, 8, 8, 3)      1        1         0   0.00000000   0.00000000   0.00000000
(1, 1, 1, 1, 1)       (1, 8, 8, 8, 3)      1        1         1   0.00000000   0.00000000   0.00000000
(1, 1, 1, 1, 1)       (1, 8, 8, 8, 3)      1        2         0   0.00000000   0.00000000   0.00000000
(1, 1, 1, 1, 1)       (1, 8, 8, 8, 3)      1        2         1   0.00000000   0.00000000   0.00000000
(1, 4, 16, 16, 16)    (1, 64, 64, 64, 3)   0        0         0   0.00000000   0.00000000   0.00000000
(1, 4, 16, 16, 16)    (1, 64, 64, 64, 3)   0        0         1   0.00000000   0.00000000   0.00000000
(1, 4, 16, 16, 16)    (1, 64, 64, 64, 3)   0        1         0   0.00000000   0.00000000   0.00000000
(1, 4, 16, 16, 16)    (1, 64, 64, 64, 3)   0        1         1   0.00000000   0.00000000   0.00000000
(1, 4, 16, 16, 16)    (1, 64, 64, 64, 3)   0        2         0   0.00000000   0.00000000   0.00000000
(1, 4, 16, 16, 16)    (1, 64, 64, 64, 3)   0        2         1   0.00000000   0.00000000   0.00000000
(1, 4, 16, 16, 16)    (1, 64, 64, 64, 3)   1        0         0   0.00000000   0.00000000   0.00000000
(1, 4, 16, 16, 16)    (1, 64, 64, 64, 3)   1        0         1   0.00000000   0.00000000   0.00000000
(1, 4, 16, 16, 16)    (1, 64, 64, 64, 3)   1        1         0   0.00000000   0.00000000   0.00000000
(1, 4, 16, 16, 16)    (1, 64, 64, 64, 3)   1        1         1   0.00000000   0.00000000   0.00000000
(1, 4, 16, 16, 16)    (1, 64, 64, 64, 3)   1        2         0   0.00000000   0.00000000   0.00000000
(1, 4, 16, 16, 16)    (1, 64, 64, 64, 3)   1        2         1   0.00000000   0.00000000   0.00000000
(1, 4, 16, 16, 16)    (1, 1, 1, 1, 3)      0        0         0   0.00000000   0.00000000   0.00000000
(1, 4, 16, 16, 16)    (1, 1, 1, 1, 3)      0        0         1   0.00000000   0.00000000   0.00000000
(1, 4, 16, 16, 16)    (1, 1, 1, 1, 3)      0        1         0   0.00000000   0.00000000   0.00000000
(1, 4, 16, 16, 16)    (1, 1, 1, 1, 3)      0        1         1   0.00000000   0.00000000   0.00000000
(1, 4, 16, 16, 16)    (1, 1, 1, 1, 3)      0        2         0   0.00000000   0.00000000   0.00000000
(1, 4, 16, 16, 16)    (1, 1, 1, 1, 3)      0        2         1   0.00000000   0.00000000   0.00000000
(1, 4, 16, 16, 16)    (1, 1, 1, 1, 3)      1        0         0   0.00000000   0.00000000   0.00000000
(1, 4, 16, 16, 16)    (1, 1, 1, 1, 3)      1        0         1   0.00000000   0.00000000   0.00000000
(1, 4, 16, 16, 16)    (1, 1, 1, 1, 3)      1        1         0   0.00000000   0.00000000   0.00000000
(1, 4, 16, 16, 16)    (1, 1, 1, 1, 3)      1        1         1   0.00000000   0.00000000   0.00000000
(1, 4, 16, 16, 16)    (1, 1, 1, 1, 3)      1        2         0   0.00000000   0.00000000   0.00000000
(1, 4, 16, 16, 16)    (1, 1, 1, 1, 3)      1        2         1   0.00000000   0.00000000   0.00000000
(1, 4, 16, 16, 16)    (1, 8, 8, 8, 3)      0        0         0   0.00000000   0.00000000   0.00000000
(1, 4, 16, 16, 16)    (1, 8, 8, 8, 3)      0        0         1   0.00000000   0.00000000   0.00000000
(1, 4, 16, 16, 16)    (1, 8, 8, 8, 3)      0        1         0   0.00000000   0.00000000   0.00000000
(1, 4, 16, 16, 16)    (1, 8, 8, 8, 3)      0        1         1   0.00000000   0.00000000   0.00000000
(1, 4, 16, 16, 16)    (1, 8, 8, 8, 3)      0        2         0   0.00000000   0.00000000   0.00000000
(1, 4, 16, 16, 16)    (1, 8, 8, 8, 3)      0        2         1   0.00000000   0.00000000   0.00000000
(1, 4, 16, 16, 16)    (1, 8, 8, 8, 3)      1        0         0   0.00000000   0.00000000   0.00000000
(1, 4, 16, 16, 16)    (1, 8, 8, 8, 3)      1        0         1   0.00000000   0.00000000   0.00000000
(1, 4, 16, 16, 16)    (1, 8, 8, 8, 3)      1        1         0   0.00000000   0.00000000   0.00000000
(1, 4, 16, 16, 16)    (1, 8, 8, 8, 3)      1        1         1   0.00000000   0.00000000   0.00000000
(1, 4, 16, 16, 16)    (1, 8, 8, 8, 3)      1        2         0   0.00000000   0.00000000   0.00000000
(1, 4, 16, 16, 16)    (1, 8, 8, 8, 3)      1        2         1   0.00000000   0.00000000   0.00000000
------------------------------------------------------------------------------------------------------------------------
Summary:
Total tests: 120
Failed tests: 0
Max difference overall: 0.00000000
Mean difference overall: 0.00000000
Sum difference overall: 0.00000000

Edge Cases Test
input_shape           grid_shape         interp   padding   align_corners   max_diff      mean_diff     sum_diff
------------------------------------------------------------------------------------------------------------------------
(1, 1, 2, 2, 2)       (1, 1, 1, 1, 3)      0        0         0   0.00000000   0.00000000   0.00000000
(1, 1, 2, 2, 2)       (1, 1, 1, 1, 3)      0        0         1   0.00000000   0.00000000   0.00000000
(1, 1, 2, 2, 2)       (1, 1, 1, 1, 3)      0        1         0   0.00000000   0.00000000   0.00000000
(1, 1, 2, 2, 2)       (1, 1, 1, 1, 3)      0        1         1   0.00000000   0.00000000   0.00000000
(1, 1, 2, 2, 2)       (1, 1, 1, 1, 3)      1        0         0   0.00000000   0.00000000   0.00000000
(1, 1, 2, 2, 2)       (1, 1, 1, 1, 3)      1        0         1   0.00000000   0.00000000   0.00000000
(1, 1, 2, 2, 2)       (1, 1, 1, 1, 3)      1        1         0   0.00000000   0.00000000   0.00000000
(1, 1, 2, 2, 2)       (1, 1, 1, 1, 3)      1        1         1   0.00000000   0.00000000   0.00000000
(1, 1, 1, 10, 10)     (1, 1, 5, 5, 3)      0        0         0   0.00000000   0.00000000   0.00000000
(1, 1, 1, 10, 10)     (1, 1, 5, 5, 3)      0        0         1   0.00000000   0.00000000   0.00000000
(1, 1, 1, 10, 10)     (1, 1, 5, 5, 3)      0        1         0   0.00000000   0.00000000   0.00000000
(1, 1, 1, 10, 10)     (1, 1, 5, 5, 3)      0        1         1   0.00000000   0.00000000   0.00000000
(1, 1, 1, 10, 10)     (1, 1, 5, 5, 3)      1        0         0   0.00000000   0.00000000   0.00000000
(1, 1, 1, 10, 10)     (1, 1, 5, 5, 3)      1        0         1   0.00000000   0.00000000   0.00000000
(1, 1, 1, 10, 10)     (1, 1, 5, 5, 3)      1        1         0   0.00000000   0.00000000   0.00000000
(1, 1, 1, 10, 10)     (1, 1, 5, 5, 3)      1        1         1   0.00000000   0.00000000   0.00000000
(1, 16, 8, 8, 8)      (1, 4, 4, 4, 3)      0        0         0   0.00000000   0.00000000   0.00000000
(1, 16, 8, 8, 8)      (1, 4, 4, 4, 3)      0        0         1   0.00000000   0.00000000   0.00000000
(1, 16, 8, 8, 8)      (1, 4, 4, 4, 3)      0        1         0   0.00000000   0.00000000   0.00000000
(1, 16, 8, 8, 8)      (1, 4, 4, 4, 3)      0        1         1   0.00000000   0.00000000   0.00000000
(1, 16, 8, 8, 8)      (1, 4, 4, 4, 3)      1        0         0   0.00000000   0.00000000   0.00000000
(1, 16, 8, 8, 8)      (1, 4, 4, 4, 3)      1        0         1   0.00000000   0.00000000   0.00000000
(1, 16, 8, 8, 8)      (1, 4, 4, 4, 3)      1        1         0   0.00000000   0.00000000   0.00000000
(1, 16, 8, 8, 8)      (1, 4, 4, 4, 3)      1        1         1   0.00000000   0.00000000   0.00000000

All tests passed!

Gradient accuracy test vs. CPU

Click to expand

grad wrt. input
input_shape           grid_shape         interp   padding   align_corners   diff
(1, 3, 128, 128, 128) (1, 64, 64, 64, 3)   0        0         0    6.324674117763607e-10
(1, 3, 128, 128, 128) (1, 64, 64, 64, 3)   0        0         1    7.693850001544433e-10
(1, 3, 128, 128, 128) (1, 64, 64, 64, 3)   0        1         0    6.303384481043395e-10
(1, 3, 128, 128, 128) (1, 64, 64, 64, 3)   0        1         1    7.976201921167103e-10
(1, 3, 128, 128, 128) (1, 64, 64, 64, 3)   1        0         0                      0.0
(1, 3, 128, 128, 128) (1, 64, 64, 64, 3)   1        0         1                      0.0
(1, 3, 128, 128, 128) (1, 64, 64, 64, 3)   1        1         0                      0.0
(1, 3, 128, 128, 128) (1, 64, 64, 64, 3)   1        1         1                      0.0
(1, 3, 128, 128, 128) (1, 1, 1, 1, 3)      0        0         0                      0.0
(1, 3, 128, 128, 128) (1, 1, 1, 1, 3)      0        0         1                      0.0
(1, 3, 128, 128, 128) (1, 1, 1, 1, 3)      0        1         0                      0.0
(1, 3, 128, 128, 128) (1, 1, 1, 1, 3)      0        1         1                      0.0
(1, 3, 128, 128, 128) (1, 1, 1, 1, 3)      1        0         0                      0.0
(1, 3, 128, 128, 128) (1, 1, 1, 1, 3)      1        0         1                      0.0
(1, 3, 128, 128, 128) (1, 1, 1, 1, 3)      1        1         0                      0.0
(1, 3, 128, 128, 128) (1, 1, 1, 1, 3)      1        1         1                      0.0
(1, 1, 1, 1, 1)       (1, 64, 64, 64, 3)   0        0         0   1.1622905731201172e-06
(1, 1, 1, 1, 1)       (1, 64, 64, 64, 3)   0        0         1                      0.0
(1, 1, 1, 1, 1)       (1, 64, 64, 64, 3)   0        1         0                      0.0
(1, 1, 1, 1, 1)       (1, 64, 64, 64, 3)   0        1         1                      0.0
(1, 1, 1, 1, 1)       (1, 64, 64, 64, 3)   1        0         0                      0.0
(1, 1, 1, 1, 1)       (1, 64, 64, 64, 3)   1        0         1                      0.0
(1, 1, 1, 1, 1)       (1, 64, 64, 64, 3)   1        1         0                      0.0
(1, 1, 1, 1, 1)       (1, 64, 64, 64, 3)   1        1         1                      0.0
(1, 1, 1, 1, 1)       (1, 1, 1, 1, 3)      0        0         0                      0.0
(1, 1, 1, 1, 1)       (1, 1, 1, 1, 3)      0        0         1                      0.0
(1, 1, 1, 1, 1)       (1, 1, 1, 1, 3)      0        1         0                      0.0
(1, 1, 1, 1, 1)       (1, 1, 1, 1, 3)      0        1         1                      0.0
(1, 1, 1, 1, 1)       (1, 1, 1, 1, 3)      1        0         0                      0.0
(1, 1, 1, 1, 1)       (1, 1, 1, 1, 3)      1        0         1                      0.0
(1, 1, 1, 1, 1)       (1, 1, 1, 1, 3)      1        1         0                      0.0
(1, 1, 1, 1, 1)       (1, 1, 1, 1, 3)      1        1         1                      0.0
max diff for input gradients:  1.1622905731201172e-06
grad wrt. grid
input_shape           grid_shape         interp   padding   align_corners   diff
(1, 3, 128, 128, 128) (1, 64, 64, 64, 3)   0        0         0                      0.0
(1, 3, 128, 128, 128) (1, 64, 64, 64, 3)   0        0         1                      0.0
(1, 3, 128, 128, 128) (1, 64, 64, 64, 3)   0        1         0                      0.0
(1, 3, 128, 128, 128) (1, 64, 64, 64, 3)   0        1         1                      0.0
(1, 3, 128, 128, 128) (1, 64, 64, 64, 3)   1        0         0                      0.0
(1, 3, 128, 128, 128) (1, 64, 64, 64, 3)   1        0         1                      0.0
(1, 3, 128, 128, 128) (1, 64, 64, 64, 3)   1        1         0                      0.0
(1, 3, 128, 128, 128) (1, 64, 64, 64, 3)   1        1         1                      0.0
(1, 3, 128, 128, 128) (1, 1, 1, 1, 3)      0        0         0                      0.0
(1, 3, 128, 128, 128) (1, 1, 1, 1, 3)      0        0         1                      0.0
(1, 3, 128, 128, 128) (1, 1, 1, 1, 3)      0        1         0                      0.0
(1, 3, 128, 128, 128) (1, 1, 1, 1, 3)      0        1         1                      0.0
(1, 3, 128, 128, 128) (1, 1, 1, 1, 3)      1        0         0                      0.0
(1, 3, 128, 128, 128) (1, 1, 1, 1, 3)      1        0         1                      0.0
(1, 3, 128, 128, 128) (1, 1, 1, 1, 3)      1        1         0                      0.0
(1, 3, 128, 128, 128) (1, 1, 1, 1, 3)      1        1         1                      0.0
(1, 1, 1, 1, 1)       (1, 64, 64, 64, 3)   0        0         0                      0.0
(1, 1, 1, 1, 1)       (1, 64, 64, 64, 3)   0        0         1                      0.0
(1, 1, 1, 1, 1)       (1, 64, 64, 64, 3)   0        1         0                      0.0
(1, 1, 1, 1, 1)       (1, 64, 64, 64, 3)   0        1         1                      0.0
(1, 1, 1, 1, 1)       (1, 64, 64, 64, 3)   1        0         0                      0.0
(1, 1, 1, 1, 1)       (1, 64, 64, 64, 3)   1        0         1                      0.0
(1, 1, 1, 1, 1)       (1, 64, 64, 64, 3)   1        1         0                      0.0
(1, 1, 1, 1, 1)       (1, 64, 64, 64, 3)   1        1         1                      0.0
(1, 1, 1, 1, 1)       (1, 1, 1, 1, 3)      0        0         0                      0.0
(1, 1, 1, 1, 1)       (1, 1, 1, 1, 3)      0        0         1                      0.0
(1, 1, 1, 1, 1)       (1, 1, 1, 1, 3)      0        1         0                      0.0
(1, 1, 1, 1, 1)       (1, 1, 1, 1, 3)      0        1         1                      0.0
(1, 1, 1, 1, 1)       (1, 1, 1, 1, 3)      1        0         0                      0.0
(1, 1, 1, 1, 1)       (1, 1, 1, 1, 3)      1        0         1                      0.0
(1, 1, 1, 1, 1)       (1, 1, 1, 1, 3)      1        1         0                      0.0
(1, 1, 1, 1, 1)       (1, 1, 1, 1, 3)      1        1         1                      0.0
max diff for grid gradients:  0.0

pytorch-bot · 2025-07-29T21:24:40Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/159421

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

❌ 2 New Failures, 2 Unrelated Failures

As of commit 36d9498 with merge base c665594 ():

NEW FAILURES - The following jobs have failed:

pull / linux-jammy-py3.9-clang12 / test (default, 1, 5, lf.linux.4xlarge) (gh)
'Test'
pull / linux-jammy-py3.9-gcc11 / test (default, 1, 5, lf.linux.2xlarge) (gh)
'Test'

FLAKY - The following job failed but was likely due to flakiness present on trunk:

pull / linux-jammy-cuda12.8-py3.10-gcc11-test / test (distributed, 1, 3, lf.linux.g4dn.12xlarge.nvidia.gpu) (gh) (disabled by #147726 but the issue was closed recently and a rebase is needed to make it pass)
distributed/test_c10d_functional_native.py::CompileTest::test_inductor_all_reduce_coalesced

UNSTABLE - The following job is marked as unstable, possibly due to flakiness on trunk:

pull / linux-jammy-py3_9-clang9-xla / test (xla, 1, 1, lf.linux.12xlarge, unstable) (gh) (#158876)
sccache: error: couldn't connect to server

This comment was automatically generated by Dr. CI and updates every 15 minutes.

github-actions · 2025-07-29T21:28:42Z

Attention! native_functions.yaml was changed

If you are adding a new function or defaulted argument to native_functions.yaml, you cannot use it from pre-existing Python frontend code until our FC window passes (two weeks). Split your PR into two PRs, one which adds the new C++ functionality, and one that makes use of it from Python, and land them two weeks apart. See https://github.com/pytorch/pytorch/wiki/PyTorch's-Python-Frontend-Backward-and-Forward-Compatibility-Policy#forwards-compatibility-fc for more info.

Caused by:

aten/src/ATen/native/native_functions.yaml

aten/src/ATen/native/mps/kernels/GridSampler3D.metal

aten/src/ATen/native/mps/operations/GridSampler.mm

mlaves · 2025-08-22T07:41:40Z

Any updates on this? @malfet

mlaves · 2025-08-27T07:16:20Z

I see that there was another PR #160541 that was opened after this one, but already merged. It looks like that the merged PR does not implement the backward pass or nearest interpolation. Therefore, my PR might still be worth merging. @malfet @kurtamohler

kurtamohler · 2025-08-27T19:46:44Z

Hey @mlaves, I'm sorry I had not noticed that you opened this PR. I have started to work on the backward part as well, but I could stop if you'd like to update your PR to only add the backward function.

You'll need to enable tests for it by removing this line:

pytorch/torch/testing/_internal/common_mps.py

Line 665 in 06c7516

"grid_sampler_3d": None,

You can run the tests on your machine with python test/test_mps.py -k grid_sampler_3d

@malfet , does that sound good to you?

malfet · 2025-08-27T21:01:51Z

@mlaves sorry I've missed your change earlier (I usually go over all the changes that claim to Fix some of the MPS issues or has ciflow/mps in the title.

Please rebase your change and I'll have a look at it soon

mlaves · 2025-08-28T07:06:56Z

@kurtamohler No worries! I also missed your PR and saw it just now. It's rather remarkable that this has been missing for years, only to be implemented by two individuals within a couple of days lol.

@malfet I will rebase my changes on current main as soon as possible.

github-actions · 2025-10-27T07:35:01Z

Looks like this PR hasn't been updated in a while so we're going to go ahead and mark this as Stale.
Feel free to remove the Stale label if you feel this was a mistake.
If you are unable to remove the Stale label please contact a maintainer in order to do so.
If you want the bot to never mark this PR stale again, add the no-stale label.
Stale pull requests will automatically be closed after 30 days of inactivity.

mlaves added 6 commits July 29, 2025 23:21

initial version of grid_sampler_3d_mps forward pass

6a7b040

added vectorized version for speed improvements

816bda7

added boilerplate for backward pass

fc57913

added backward kernel for grid_sampler_3d wrt. input

5b79872

some clean-ups and refactoring

3e0cf6a

added CPU fallback for backward pass

8ea2c3f

mlaves requested review from kulinseth and malfet as code owners July 29, 2025 21:24

pytorch-bot bot added the release notes: mps Release notes category label Jul 29, 2025

pytorchbot added the open source label Jul 29, 2025

kulinseth approved these changes Jul 30, 2025

View reviewed changes

Skylion007 reviewed Aug 3, 2025

View reviewed changes

aten/src/ATen/native/mps/kernels/GridSampler3D.metal Outdated Show resolved Hide resolved

aten/src/ATen/native/mps/operations/GridSampler.mm Outdated Show resolved Hide resolved

aten/src/ATen/native/mps/operations/GridSampler.mm Outdated Show resolved Hide resolved

Skylion007 reviewed Aug 3, 2025

View reviewed changes

aten/src/ATen/native/mps/operations/GridSampler.mm Outdated Show resolved Hide resolved

resolved reviewer comments and lint issues

4d67278

mlaves requested a review from Skylion007 August 3, 2025 20:58

Skylion007 approved these changes Aug 5, 2025

View reviewed changes

aten/src/ATen/native/mps/operations/GridSampler.mm Outdated Show resolved Hide resolved

aten/src/ATen/native/mps/operations/GridSampler.mm Outdated Show resolved Hide resolved

aten/src/ATen/native/mps/operations/GridSampler.mm Outdated Show resolved Hide resolved

minor reviewer comments

36d9498

kurtamohler mentioned this pull request Aug 28, 2025

Most requested ops for the MPS backend #154052

Open

github-actions bot added the Stale label Oct 27, 2025

github-actions bot closed this Nov 26, 2025

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

grid_sampler_3d for MPS #159421

grid_sampler_3d for MPS #159421

Uh oh!

mlaves commented Jul 29, 2025 •

edited

Loading

Uh oh!

pytorch-bot bot commented Jul 29, 2025 •

edited

Loading

Uh oh!

github-actions bot commented Jul 29, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mlaves commented Aug 22, 2025

Uh oh!

mlaves commented Aug 27, 2025

Uh oh!

kurtamohler commented Aug 27, 2025 •

edited

Loading

Uh oh!

malfet commented Aug 27, 2025

Uh oh!

mlaves commented Aug 28, 2025

Uh oh!

github-actions bot commented Oct 27, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

grid_sampler_3d for MPS #159421

grid_sampler_3d for MPS #159421

Uh oh!

Conversation

mlaves commented Jul 29, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Jul 29, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/159421

❌ 2 New Failures, 2 Unrelated Failures

Uh oh!

github-actions bot commented Jul 29, 2025

Attention! native_functions.yaml was changed

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mlaves commented Aug 22, 2025

Uh oh!

mlaves commented Aug 27, 2025

Uh oh!

kurtamohler commented Aug 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

malfet commented Aug 27, 2025

Uh oh!

mlaves commented Aug 28, 2025

Uh oh!

github-actions bot commented Oct 27, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

mlaves commented Jul 29, 2025 •

edited

Loading

pytorch-bot bot commented Jul 29, 2025 •

edited

Loading

kurtamohler commented Aug 27, 2025 •

edited

Loading