Skip to content

viable/strict/1788168892

@malfet malfet tagged this 31 Aug 05:15
Extends the native Metal `lu_solve` to complex64, unblocking complex64 backward for `det`, `slogdet`, `logdet`, `solve` and `solve_ex` on MPS

The lu_solve schedule is unchanged; only the per-step kernels are made dtype-generic:
- `trsmDiagSolveLU` and `luApplyPivotsRHS` are templated on the scalar type and instantiated for `float` and `float2`. The float path is kept byte-identical (dcol register staging + `fma`, identity-padded diagonal); the complex branch reads the shared tile directly and uses `c10::metal::div`/`mul`, guarding the divide on the block-local column bound so the zero padding never yields a 0/0 (there is no complex "one" to pad the diagonal with).
- The Schur update reuses the existing complex-capable `gemmTiledLU`, since the `simdgroup_matrix` and MPP `matmul2d` GEMMs are float-only hardware paths.
- The adjoint case needs no kernel change: the host already materializes `LU.mH()` into  the augmented matrix, so conjugation happens at copy time and the kernels only do complex arithmetic on stored values.

Backward w.r.t. the LU factor itself (`linalg_lu_solve_LU`) and the backward of `lu_factor`/`lu`/`cholesky` still need a complex `solve_triangular`, which on MPS delegates to the float-only `MPSMatrixSolveTriangular`; that is left as a follow-up. `lu_solve`'s OpInfo therefore keeps `backward_dtypesIfMPS=floating` (the B-gradient works, the LU-gradient does not).

> [!NOTE]
> This PR was authored with the assistance of Claude Code.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/195369
Approved by: https://github.com/Skylion007

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Assets 2
Loading