[MPS] Speed up QR with blocked householder updates (#198732)
First reported here https://github.com/kornia/kornia/issues/4522 QR decomposition runs reaaallly slow now on main. Embarassingly slow, batch size 1 **512 x 512 QR at 6.976042 SECONDS**... Perf, comparisons against before on MPS are deliberately excluded for obvious reason: | Shape | CPU (ms) | MPS (ms) | CPU / MPS | |---|---:|---:|---:| | 128 x 128 | 0.219 | 0.298 | 0.73x | | 384 x 384 | 1.740 | 0.993 | 1.75x | | 512 x 512 | 4.083 | 1.465 | 2.78x | | 1024 x 1024 | 18.796 | 4.954 | 3.79x | | 2048 x 2048 | 97.228 | 19.503 | 4.98x | | 1152 x 128 | 1.959 | 0.897 | 2.18x | | 2048 x 512 | 15.830 | 6.025 | 2.63x | | 16384 x 512 | 106.555 | 32.437 | 3.29x | | 64 x 9 x 8 | 0.067 | 0.031 | 2.13x | | 256 x 3 x 3 | 0.066 | 0.020 | 3.35x | | 2048 x 9 x 8 | 1.504 | 0.219 | 6.87x | | 4096 x 9 x 8 | 2.957 | 0.413 | 7.17x | | 16384 x 9 x 8 | 11.555 | 1.566 | 7.38x | | 65536 x 9 x 8 | 46.528 | 6.225 | 7.47x | | 2048 x 3 x 3 | 0.403 | 0.040 | 9.98x | | 4096 x 3 x 3 | 0.781 | 0.067 | 11.67x | | 16384 x 3 x 3 | 3.066 | 0.222 | 13.83x | | 65536 x 3 x 3 | 11.598 | 0.827 | 14.02x | In terms of numerics, they are better than the ones currently present on main | Case | Metric | CPU (LAPACK) | MPS main | MPS this PR | |---|---|---:|---:|---:| | 512 x 512 | backward error | 8.0e-07 | 4.1e-06 | 7.7e-07 | | 512 x 512 | orthogonality | 1.9e-05 | 1.8e-04 | 1.9e-05 | | 1024 x 64 | backward error | 2.2e-07 | 2.2e-06 | 3.2e-07 | | 64 x 32 x 32 | backward error | 2.4e-07 | 3.5e-07 | 3.7e-07 | | 128 x 128, scaled by 1e-10 | backward error | 4.0e-07 | 0.70 | 5.0e-07 | | 128 x 128, scaled by 1e20 | backward error | 4.0e-07 | NaN | 5.1e-07 | | 256 x 256, condition number 1e6 | backward error | 4.3e-07 | 3.9e-05 | 4.7e-07 | | R only, 2048 x 2048 | error in R | 1.7e-06 | 9.9e-07 | 1.4e-06 | Pull Request resolved: https://github.com/pytorch/pytorch/pull/198732 Approved by: https://github.com/malfet
I
Isalia20 committed
4d19930b18fadc326ffe2a9618b8b457f3794043
Parent: e4112a1
Committed by PyTorch MergeBot <pytorchmergebot@users.noreply.github.com>
on 10/1/2026, 5:47:32 AM