viable/strict/1788168892: [MPS] Add complex64 support to linalg.lu_solve (#195369)
PyTorch 的 MPS 后端为 linalg.lu_solve 添加了 complex64 支持,使 det、slogdet、logdet、solve 和 solve_ex 的 complex64 反向传播在 MPS 上可用。lu_solve 的调度不变,仅将 trsmDiagSolveLU 和 luApplyPivotsRHS 内核模板化以支持 float 和 float2。float 路径保持字节级一致,complex 分支直接读取共享 tile 并使用 c10::metal::div/mul,通过块局部列边界防止 0/0。Schur 更新复用现有的 complex-capable gemmTiledLU。伴随情况无需内核更改,主机端在复制时物化 LU.mH()。反向传播 w.r.t. LU 因子及 lu_factor/lu/cholesky 的反向仍需要 complex solve_triangular,在 MPS 上委托给仅支持 float 的 MPSMatrixSolveTriangular。
Development
- First Reportviable/strict/1788168892: [MPS] Add complex64 support to linalg.lu_solve (#195369)PyTorch Core
- Industry Responsetrunk/2053ec330a7e9bc250fb9dd22d7c203ddd109410: [MPS] Add complex64 support to linalg.lu_solve (#195369)PyTorch Core
- Current AssessmentPyTorch 持续扩展 MPS 后端的 complex64 支持,表明 Apple Silicon 在科学计算和深度学习中的重要性日益增加。此提交解锁了 MPS 上多个线性代数操作的 complex64 反向传播,可能吸引更多研究者和工程师在 Mac 上进行复杂数值计算。Agent Pulse · analysis
PyTorch 的 MPS 后端为 linalg.lu_solve 添加了 complex64 支持,使 det、slogdet、logdet、solve 和 solve_ex 的 complex64 反向传播在 MPS 上可用。lu_solve 的调度不变,仅将 trsmDiagSolveLU 和 luApplyPivotsRHS 内核模板化以支持 float 和 float2。float 路径保持字节级一致,complex 分支直接读取共享 tile 并使用 c10::metal::div/mul,通过块局部列边界防止 0/0。Schur 更新复用现有的 complex-capable gemmTiledLU。伴随情况无需内核更改,主机端在复制时物化 LU.mH()。反向传播 w.r.t. LU 因子及 lu_factor/lu/cholesky 的反向仍需要 complex solve_triangular,在 MPS 上委托给仅支持 float 的 MPSMatrixSolveTriangular。
该提交通过将内核模板化并实例化 float 和 float2,为 MPS 上的 lu_solve 添加了 complex64 支持。float 路径保持字节级一致,complex 分支使用 c10::metal::div/mul 并防止 0/0。Schur 更新复用 complex-capable gemmTiledLU,因为 simdgroup_matrix 和 MPP matmul2d GEMM 是仅 float 的硬件路径。伴随情况无需内核更改,因为主机端在复制时物化 LU.mH()。反向传播 w.r.t. LU 因子仍需要 complex solve_triangular,在 MPS 上委托给仅 float 的 MPSMatrixSolveTriangular。
PyTorch 持续扩展 MPS 后端的 complex64 支持,表明 Apple Silicon 在科学计算和深度学习中的重要性日益增加。此提交解锁了 MPS 上多个线性代数操作的 complex64 反向传播,可能吸引更多研究者和工程师在 Mac 上进行复杂数值计算。
此提交增强了 PyTorch 在 Apple Silicon 上的功能,使 Mac 用户能够进行更广泛的科学计算和深度学习。这可能提高 PyTorch 在 Mac 开发者中的采用率,并促进 Apple 硬件在 AI 工作负载中的使用。
未来,MPS 后端可能继续扩展 complex64 支持,例如为 lu_factor、lu、cholesky 的反向添加 complex solve_triangular。此外,可能优化 complex 内核性能,或为其他线性代数操作添加 complex64 支持。