Skip to content

[MPS] Add cholesky_solve support#176703

Closed
Kingwl wants to merge 1 commit intopytorch:mainfrom
Kingwl:feat/support-cholesky-solve-for-mps
Closed

[MPS] Add cholesky_solve support#176703
Kingwl wants to merge 1 commit intopytorch:mainfrom
Kingwl:feat/support-cholesky-solve-for-mps

Conversation

@Kingwl
Copy link
Copy Markdown
Contributor

@Kingwl Kingwl commented Mar 6, 2026 •
edited by malfet
Loading

As decomposition via two triangular solves

Frequently requested op in #154052

@Kingwl Kingwl requested review from malfet and mruberry as code owners March 6, 2026 09:05
@pytorch-bot
Copy link
Copy Markdown

pytorch-bot bot commented Mar 6, 2026 •
edited
Loading

🔗 Helpful Links

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

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

❌ 1 New Failure

As of commit 2738fc0 with merge base 08b6f48 (image):

NEW FAILURE - The following job has failed:

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

@pytorch-bot pytorch-bot bot added the release notes: mps Release notes category label Mar 6, 2026
@linux-foundation-easycla
Copy link
Copy Markdown

linux-foundation-easycla bot commented Mar 6, 2026 •
edited
Loading

CLA Signed

The committers listed above are authorized under a signed CLA.

  • ✅ login: Kingwl / name: Wenlu Wang (2738fc0)

@github-actions
Copy link
Copy Markdown
Contributor

github-actions bot commented Mar 6, 2026

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:

@malfet malfet added the ciflow/mps Run MPS tests (subset of trunk) label Mar 6, 2026
@pytorch-bot pytorch-bot bot removed the ciflow/mps Run MPS tests (subset of trunk) label Mar 6, 2026
@Kingwl Kingwl force-pushed the feat/support-cholesky-solve-for-mps branch 2 times, most recently from 458a6b7 to e5e17f7 Compare March 7, 2026 09:40
kurtamohler
kurtamohler reviewed Mar 11, 2026
View reviewed changes
kurtamohler
kurtamohler reviewed Mar 11, 2026
View reviewed changes
kurtamohler
kurtamohler reviewed Mar 11, 2026
View reviewed changes
kurtamohler
kurtamohler reviewed Mar 11, 2026
View reviewed changes
@kurtamohler kurtamohler added the ciflow/mps Run MPS tests (subset of trunk) label Mar 11, 2026
@pytorch-bot
Copy link
Copy Markdown

pytorch-bot bot commented Mar 11, 2026

To add the ciflow label ciflow/mps please first approve the workflows that are awaiting approval (scroll to the bottom of this page).

This helps ensure we don't trigger CI on this PR until it is actually authorized to do so. Please ping one of the reviewers if you do not have access to approve and run workflows.

@pytorch-bot pytorch-bot bot removed the ciflow/mps Run MPS tests (subset of trunk) label Mar 11, 2026
@kurtamohler kurtamohler added the ciflow/mps Run MPS tests (subset of trunk) label Mar 11, 2026
kurtamohler
kurtamohler reviewed Mar 11, 2026
View reviewed changes
kurtamohler
kurtamohler reviewed Mar 11, 2026
View reviewed changes
Copy link
Copy Markdown
Collaborator

@kurtamohler kurtamohler left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for the PR! Left a few suggestions

@Kingwl Kingwl force-pushed the feat/support-cholesky-solve-for-mps branch from e5e17f7 to f8ce1f2 Compare March 11, 2026 04:30
@pytorch-bot pytorch-bot bot removed the ciflow/mps Run MPS tests (subset of trunk) label Mar 11, 2026
@kurtamohler kurtamohler added the ciflow/mps Run MPS tests (subset of trunk) label Mar 11, 2026
@kurtamohler
Copy link
Copy Markdown
Collaborator

kurtamohler commented Mar 12, 2026

@pytorchbot rebase

@pytorchmergebot
Copy link
Copy Markdown
Collaborator

pytorchmergebot commented Mar 12, 2026

@pytorchbot started a rebase job onto refs/remotes/origin/viable/strict. Check the current status here

@pytorchmergebot
Copy link
Copy Markdown
Collaborator

pytorchmergebot commented Mar 12, 2026

Successfully rebased feat/support-cholesky-solve-for-mps onto refs/remotes/origin/viable/strict, please pull locally before adding more changes (for example, via git checkout feat/support-cholesky-solve-for-mps && git pull --rebase)

@pytorchmergebot pytorchmergebot force-pushed the feat/support-cholesky-solve-for-mps branch from f8ce1f2 to 2738fc0 Compare March 12, 2026 01:03
@pytorch-bot pytorch-bot bot removed the ciflow/mps Run MPS tests (subset of trunk) label Mar 12, 2026
@kurtamohler
Copy link
Copy Markdown
Collaborator

kurtamohler commented Mar 12, 2026

@pytorchbot label "ciflow/mps"

@pytorch-bot pytorch-bot bot added the ciflow/mps Run MPS tests (subset of trunk) label Mar 12, 2026
@malfet malfet added the topic: improvements topic category label Mar 12, 2026
@malfet
Copy link
Copy Markdown
Contributor

malfet commented Mar 12, 2026

@pytorchbot merge

@pytorch-bot pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Mar 12, 2026
@pytorchmergebot
Copy link
Copy Markdown
Collaborator

pytorchmergebot commented Mar 12, 2026

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging
Check the merge workflow status
here

@Kingwl
Copy link
Copy Markdown
Contributor Author

Kingwl commented Mar 12, 2026

BTW: the solve triangular extremely slow on MPS. And It have no Idea how to be it faster after do some investigation. Any suggestion？😿

@pytorchmergebot
Copy link
Copy Markdown
Collaborator

pytorchmergebot commented Mar 12, 2026

Merge failed

Reason: 1 jobs have failed, first few of them are: trunk / macos-py3-arm64 / build

Details for Dev Infra team Raised by workflow job

@kurtamohler
Copy link
Copy Markdown
Collaborator

kurtamohler commented Mar 12, 2026 •
edited
Loading

BTW: the solve triangular extremely slow on MPS. And It have no Idea how to be it faster after do some investigation. Any suggestion？😿

Is it slow just for batched inputs with more than 2 dims, or is it also slow for non-batched inputs?

It seems that if the input is batched, each matrix is computed serially:

pytorch/aten/src/ATen/native/mps/operations/LinearAlgebra.mm

Line 1246 in f72a552

for (const auto i : c10::irange(batchSize)) {

So a major improvement would be to parallelize the batches. It looks like MPSMatrixSolveTriangular does not support batched matrices, so it would have to be implemented as a custom metal kernel.

@malfet
Copy link
Copy Markdown
Contributor

malfet commented Mar 12, 2026

@pytorchbot merge -i

@malfet malfet changed the title Feat: add cholesky_solve support for mps [MPS] Add cholesky_solve support Mar 12, 2026
@pytorchmergebot
Copy link
Copy Markdown
Collaborator

pytorchmergebot commented Mar 12, 2026

Merge started

Your change will be merged while ignoring the following 1 checks: pull / linux-jammy-py3.14t-clang15 / test (crossref, 2, 2, lf.linux.2xlarge)

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging
Check the merge workflow status
here

EmanueleCoradin pushed a commit to EmanueleCoradin/pytorch that referenced this pull request Mar 30, 2026
@Kingwl @EmanueleCoradin
[MPS] Add cholesky_solve support (pytorch#176703)
9c2c12f 
As decomposition via two triangular solves

Frequently requested op in  pytorch#154052
Pull Request resolved: pytorch#176703
Approved by: https://github.com/kurtamohler, https://github.com/malfet
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@malfet malfet malfet approved these changes

@kurtamohler kurtamohler kurtamohler approved these changes

@mruberry mruberry Awaiting requested review from mruberry mruberry is a code owner

Assignees

No one assigned

Labels

ciflow/mps Run MPS tests (subset of trunk) ciflow/trunk Trigger trunk jobs on your pull request Merged open source release notes: mps Release notes category topic: improvements topic category

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

5 participants

@Kingwl @kurtamohler @pytorchmergebot @malfet @pytorchbot