Cholesky mps implementation by Isalia20 · Pull Request #144193 · pytorch/pytorch

Isalia20 · 2025-01-04T16:36:37Z

Requested in #77764

PR is still in draft because it needs some cleanups and optimizations to get to cpu performance the least. Tasks:

Make upper=True work, only upper=False works now
Code cleanup
Optimizations(Though might need some help on this)(tried my best, maybe there is still some more to squeeze out)
Checks for positive definite input
Support for (*, N, N) input, currently only supports (B, N, N) input
Support other dtypes(float16, bfloat16)

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @chauhang @aakhundov

pytorch-bot · 2025-01-04T16:36:41Z

🔗 Helpful Links

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

📄 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.

❌ 1 New Failure, 26 Unrelated Failures

As of commit dc7617d with merge base b7bef1c ():

NEW FAILURE - The following job has failed:

pull / linux-jammy-py3-clang12-executorch / test (executorch, 1, 1, lf.linux.2xlarge) (gh)
Process completed with exit code 1.

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

inductor / cuda12.4-py3.10-gcc9-sm86 / test (inductor_timm, 1, 2, linux.g5.4xlarge.nvidia.gpu) (gh) (similar failure)
levit_128
inductor / linux-jammy-cpu-py3.9-gcc11-inductor / test (dynamic_cpu_inductor_timm, 1, 2, linux.8xlarge.amx) (gh) (similar failure)
levit_128
pull / linux-focal-cuda11.8-py3.10-gcc9 / test (distributed, 1, 3, lf.linux.g4dn.12xlarge.nvidia.gpu) (gh) (disabled by #137771 but the issue was closed recently and a rebase is needed to make it pass)
distributed/test_c10d_ucc.py::DistributedDataParallelTest::test_save_load_checkpoint
pull / linux-focal-cuda12.4-py3.10-gcc9 / test (default, 1, 5, lf.linux.4xlarge.nvidia.gpu) (gh) (disabled by #131082, #131083, #131088, #131101, #131119, #131136, #138675, #138884 but the issue was closed recently and a rebase is needed to make it pass)
export/test_export.py::TestExport::test_slice_with_floordiv
pull / linux-focal-cuda12.4-py3.10-gcc9 / test (default, 2, 5, lf.linux.4xlarge.nvidia.gpu) (gh) (disabled by #131119 but the issue was closed recently and a rebase is needed to make it pass)
export/test_serdes.py::SerDesExportTestExport::test_slice_with_floordiv_serdes
pull / linux-focal-cuda12.4-py3.10-gcc9 / test (default, 3, 5, lf.linux.4xlarge.nvidia.gpu) (gh) (disabled by #131082 but the issue was closed recently and a rebase is needed to make it pass)
export/test_export_training_ir_to_run_decomp.py::TrainingIRToRunDecompExportTestExport::test_slice_with_floordiv_training_ir_to_decomp
pull / linux-focal-cuda12.4-py3.10-gcc9 / test (default, 4, 5, lf.linux.4xlarge.nvidia.gpu) (gh) (disabled by #131083 but the issue was closed recently and a rebase is needed to make it pass)
export/test_retraceability.py::RetraceExportTestExport::test_slice_with_floordiv_retraceability
pull / linux-focal-cuda12.4-py3.10-gcc9-sm89 / test (default, 3, 5, lf.linux.g6.4xlarge.experimental.nvidia.gpu) (gh) (disabled by #131083 but the issue was closed recently and a rebase is needed to make it pass)
export/test_retraceability.py::RetraceExportTestExport::test_slice_with_floordiv_retraceability
pull / linux-focal-cuda12.4-py3.10-gcc9-sm89 / test (default, 4, 5, lf.linux.g6.4xlarge.experimental.nvidia.gpu) (gh) (disabled by #131082 but the issue was closed recently and a rebase is needed to make it pass)
export/test_export_training_ir_to_run_decomp.py::TrainingIRToRunDecompExportTestExport::test_slice_with_floordiv_training_ir_to_decomp
pull / linux-focal-cuda12.4-py3.10-gcc9-sm89 / test (default, 5, 5, lf.linux.g6.4xlarge.experimental.nvidia.gpu) (gh) (disabled by #131119 but the issue was closed recently and a rebase is needed to make it pass)
export/test_serdes.py::SerDesExportTestExport::test_slice_with_floordiv_serdes
pull / linux-focal-py3.13-clang10 / test (crossref, 2, 2, lf.linux.2xlarge) (gh) (disabled by #131082, #131083, #131088, #131101, #131119, #131136, #138675, #138884 but the issue was closed recently and a rebase is needed to make it pass)
export/test_export.py::TestExport::test_slice_with_floordiv
pull / linux-focal-py3.13-clang10 / test (default, 1, 5, lf.linux.4xlarge) (gh) (disabled by #131082, #131083, #131088, #131101, #131119, #131136, #138675, #138884 but the issue was closed recently and a rebase is needed to make it pass)
export/test_export.py::TestExport::test_slice_with_floordiv
pull / linux-focal-py3.13-clang10 / test (default, 2, 5, lf.linux.4xlarge) (gh) (disabled by #131082 but the issue was closed recently and a rebase is needed to make it pass)
export/test_export_training_ir_to_run_decomp.py::TrainingIRToRunDecompExportTestExport::test_slice_with_floordiv_training_ir_to_decomp
pull / linux-focal-py3.13-clang10 / test (default, 4, 5, lf.linux.4xlarge) (gh) (disabled by #131119 but the issue was closed recently and a rebase is needed to make it pass)
export/test_serdes.py::SerDesExportTestExport::test_slice_with_floordiv_serdes
pull / linux-focal-py3.9-clang10 / test (crossref, 2, 2, lf.linux.2xlarge) (gh) (disabled by #131082, #131083, #131088, #131101, #131119, #131136, #138675, #138884 but the issue was closed recently and a rebase is needed to make it pass)
export/test_export.py::TestExport::test_slice_with_floordiv
pull / linux-focal-py3.9-clang10 / test (default, 1, 5, lf.linux.4xlarge) (gh) (disabled by #131119 but the issue was closed recently and a rebase is needed to make it pass)
export/test_serdes.py::SerDesExportTestExport::test_slice_with_floordiv_serdes
pull / linux-focal-py3.9-clang10 / test (default, 2, 5, lf.linux.4xlarge) (gh) (disabled by #131083 but the issue was closed recently and a rebase is needed to make it pass)
export/test_retraceability.py::RetraceExportTestExport::test_slice_with_floordiv_retraceability
pull / linux-focal-py3.9-clang10 / test (default, 4, 5, lf.linux.4xlarge) (gh) (disabled by #131082, #131083, #131088, #131101, #131119, #131136, #138675, #138884 but the issue was closed recently and a rebase is needed to make it pass)
export/test_export.py::TestExport::test_slice_with_floordiv
pull / linux-focal-py3.9-clang10 / test (default, 5, 5, lf.linux.4xlarge) (gh) (disabled by #131082 but the issue was closed recently and a rebase is needed to make it pass)
export/test_export_training_ir_to_run_decomp.py::TrainingIRToRunDecompExportTestExport::test_slice_with_floordiv_training_ir_to_decomp
pull / linux-focal-py3.9-clang10 / test (dynamo_wrapped, 2, 3, lf.linux.2xlarge) (gh) (disabled by #116746 but the issue was closed recently and a rebase is needed to make it pass)
torch_np/numpy_tests/lib/test_function_base.py::TestGradient::test_second_order_accurate
pull / linux-jammy-py3.10-clang15-asan / test (default, 2, 6, lf.linux.4xlarge) (gh) (disabled by #131088 but the issue was closed recently and a rebase is needed to make it pass)
export/test_export_nonstrict.py::NonStrictExportTestExport::test_slice_with_floordiv_non_strict
pull / linux-jammy-py3.10-clang15-asan / test (default, 5, 6, lf.linux.4xlarge) (gh) (disabled by #131082, #131083, #131088, #131101, #131119, #131136, #138675, #138884 but the issue was closed recently and a rebase is needed to make it pass)
export/test_export.py::TestExport::test_slice_with_floordiv
pull / linux-jammy-py3.10-clang15-asan / test (default, 6, 6, lf.linux.4xlarge) (gh) (disabled by #131083 but the issue was closed recently and a rebase is needed to make it pass)
export/test_retraceability.py::RetraceExportTestExport::test_slice_with_floordiv_retraceability
pull / linux-jammy-py3.9-gcc11 / test (default, 5, 5, lf.linux.2xlarge) (gh) (disabled by #131082 but the issue was closed recently and a rebase is needed to make it pass)
export/test_export_training_ir_to_run_decomp.py::TrainingIRToRunDecompExportTestExport::test_slice_with_floordiv_training_ir_to_decomp
trunk / macos-py3-arm64 / test (default, 2, 3, macos-m1-stable) (gh) (disabled by #131119 but the issue was closed recently and a rebase is needed to make it pass)
export/test_serdes.py::SerDesExportTestExport::test_slice_with_floordiv_serdes
trunk / macos-py3-arm64 / test (default, 3, 3, macos-m1-stable) (gh) (disabled by #131082 but the issue was closed recently and a rebase is needed to make it pass)
export/test_export_training_ir_to_run_decomp.py::TrainingIRToRunDecompExportTestExport::test_slice_with_floordiv_training_ir_to_decomp

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

github-actions · 2025-01-04T16:40:41Z

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

Isalia20 · 2025-01-05T09:05:15Z

PR is not ready yet to merge, but I'll mark it as ready for review to get some comments from you before proceeding further. There is still some things to optimize like launching a kernel to deal with syrk and trsm in parallel instead of in loops which I will tackle in the following days

Isalia20 · 2025-01-05T16:36:48Z

Added more optimizations but still slower than CPU 😞 Comments about improving speed are very welcome

Isalia20 · 2025-01-05T16:42:45Z

I compare implementations like this:

import torch
import numpy as np
from torch.profiler import profile, record_function, ProfilerActivity
from torch.profiler import schedule

n = 512
batch_size = 64
num_runs = 20

torch.manual_seed(42)
A_cpu = torch.randn(batch_size, n, n, dtype=torch.float32)
A_cpu = A_cpu @ A_cpu.transpose(-2, -1) + n * torch.eye(n).expand(batch_size, -1, -1)

A_mps = A_cpu.to("mps")

def run_cholesky_cpu(A):
    with record_function("cholesky_cpu"):
        b = torch.linalg.cholesky(A, upper=False)
        torch.cpu.synchronize()
        return b

def run_cholesky_mps(A):
    with record_function("cholesky_mps"):
        b = torch.linalg.cholesky(A, upper=False)
        torch.mps.synchronize()
        return b

my_schedule = schedule(
    skip_first=5,
    wait=0,
    warmup=3,
    active=1,
    repeat=1
)

with profile(
    activities=[ProfilerActivity.CPU],
    schedule=my_schedule,
    record_shapes=True
) as prof:
    for _ in range(num_runs):
        run_cholesky_mps(A_mps)
        run_cholesky_cpu(A_cpu)

# Profile CPU
cpu_times = []
with profile(
    activities=[ProfilerActivity.CPU],
    record_shapes=True,
    with_stack=True
) as prof_cpu:
    for _ in range(num_runs):
        L_cpu = run_cholesky_cpu(A_cpu)

print("\nCPU Profile:")
print(prof_cpu.key_averages().table(sort_by="cpu_time_total"))

# Warmup run for MPS
_ = run_cholesky_mps(A_mps)

# Profile MPS
mps_times = []
with profile(
    activities=[ProfilerActivity.CPU],
    record_shapes=True,
    with_stack=True
) as prof_mps:
    for _ in range(num_runs):
        L_mps = run_cholesky_mps(A_mps)

print("\nMPS Profile:")
print(prof_mps.key_averages().table(sort_by="cpu_time_total"))



L_mps_cpu = L_mps.cpu()

# check
tolerance = 1e-5
assert torch.allclose(L_cpu.cpu(), L_mps_cpu, rtol=tolerance, atol=tolerance), \
    f"Maximum difference between CPU and MPS results: {torch.max(torch.abs(L_cpu - L_mps_cpu))}"

mps_events = [e for e in prof_mps.key_averages() if "cholesky_mps" in e.key]
cpu_events = [e for e in prof_cpu.key_averages() if "cholesky_cpu" in e.key]

mps_mean = np.mean([e.cpu_time * 1e-6 for e in mps_events])  # Convert to seconds
cpu_mean = np.mean([e.cpu_time * 1e-6 for e in cpu_events])

print("\nAll assertions passed! CPU and MPS results match within tolerance.")
print(f"\nCPU Cholesky mean time: {cpu_mean:.4f} seconds")
print(f"MPS Cholesky mean time: {mps_mean:.4f} seconds")
print(f"\nSpeed comparison: MPS is {cpu_mean/mps_mean:.2f}x faster than CPU on average")

pytorchmergebot · 2025-01-16T21:37:35Z

@Isalia20 your PR has been successfully reverted.

This reverts commit 727ae13. Reverted #144193 on behalf of https://github.com/malfet due to Alas, inductor changes broke inductor tests, see https://hud.pytorch.org/hud/pytorch/pytorch/aa4a1ff027316c6ad850180cfc20c3962eafb4e4/1?per_page=50&name_filter=inductor%20%2F%20linux-jammy-cpu ([comment](#144193 (comment)))

This PR was reopened (likely due to being reverted), so your approval was removed. Please request another review.

Isalia20 · 2025-01-16T22:26:35Z

removed the op from fallback(originally had there before annotating linalg cholesky with CompImplicitAutograd, wrote devices like CUDA, CPU, Meta and thats why I had to have it originally in the fallback)

Tests should pass now but I suggest running the workflow here before merging @malfet

Isalia20 · 2025-01-17T07:01:33Z

Hmm still failing. This is interesting. I'm not really sure why it's failing though. I looked into the test but don't see any cholesky ops in it. Any suggestions?

Isalia20 · 2025-01-18T10:51:31Z

Some of the PRs still had the issue with levit after reverting this so we can say that issues weren't caused by this. Maybe we can merge? @malfet

Isalia20 · 2025-01-22T17:14:02Z

bump

Isalia20 · 2025-01-23T19:24:07Z

@malfet Can we merge this? Want to submit another PR with speeding this up a lot

malfet · 2025-01-26T03:45:47Z

@pytorchbot merge

pytorchmergebot · 2025-01-26T03:47:36Z

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

pytorchmergebot · 2025-01-26T05:09:08Z

Merge failed

Reason: 1 mandatory check(s) failed. The first few are:

pull / linux-jammy-py3-clang12-executorch / test (executorch, 1, 1, lf.linux.2xlarge)

Dig deeper by viewing the failures on hud

Details for Dev Infra team

Raised by workflow job

Failing merge rule: Core Maintainers

malfet · 2025-01-26T05:28:22Z

@pytorchbot merge -r

pytorchmergebot · 2025-01-26T05:29:42Z

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

pytorchmergebot · 2025-01-26T05:29:44Z

Rebase failed due to Command git -C /home/runner/work/pytorch/pytorch rebase refs/remotes/origin/viable/strict pull/144193/head returned non-zero exit code 1

Rebasing (1/33)
Rebasing (2/33)
Rebasing (3/33)
Rebasing (4/33)
Rebasing (5/33)
Rebasing (6/33)
Rebasing (7/33)
Rebasing (8/33)
Rebasing (9/33)
Rebasing (10/33)
Rebasing (11/33)
Rebasing (12/33)
Rebasing (13/33)
Rebasing (14/33)
Rebasing (15/33)
Rebasing (16/33)
Rebasing (17/33)
Rebasing (18/33)
Rebasing (19/33)
Rebasing (20/33)
Rebasing (21/33)
Rebasing (22/33)
Rebasing (23/33)
Rebasing (24/33)
Rebasing (25/33)
Rebasing (26/33)
Auto-merging aten/src/ATen/native/native_functions.yaml
Auto-merging tools/autograd/derivatives.yaml
Auto-merging torch/_inductor/lowering.py
CONFLICT (content): Merge conflict in torch/_inductor/lowering.py
error: could not apply 372da90d281... fix failing tests
hint: Resolve all conflicts manually, mark them as resolved with
hint: "git add/rm <conflicted_files>", then run "git rebase --continue".
hint: You can instead skip this commit: run "git rebase --skip".
hint: To abort and get back to the state before "git rebase", run "git rebase --abort".
hint: Disable this message with "git config set advice.mergeConflict false"
Could not apply 372da90d281... fix failing tests

Raised by https://github.com/pytorch/pytorch/actions/runs/12971799472

malfet · 2025-01-26T05:32:30Z

@Isalia20 do you mind rebasing this PR, because as of right now I could not get a clear signal. (Though I can always try merging it and see what will happen...)

Isalia20 · 2025-01-26T11:26:17Z

Tried rebasing but there are lot of commits, so I just decided to create a new one checking out from main
#145701

Requested in #77764 Closed #144193 due to a lot of conflicts when rebasing Pull Request resolved: #145701 Approved by: https://github.com/malfet

Requested in pytorch#77764 Closed pytorch#144193 due to a lot of conflicts when rebasing Pull Request resolved: pytorch#145701 Approved by: https://github.com/malfet

Isalia20 added 7 commits January 3, 2025 13:11

minimum working example with bs=1

d1ac18b

batch size support and check for non positive definite matrices

917394a

add cholesky to linalg module

10bd5bb

working implementation for bs=1 and N<64

6fd7f75

updated metal kernel and caller to support bs>1

1675fb2

linter

a22df6b

parallel processing of batches

b673a2d

pytorch-bot bot added the release notes: mps Release notes category label Jan 4, 2025

pytorchbot added the open source label Jan 4, 2025

Isalia20 added 2 commits January 4, 2025 21:47

bugfix for N>64 matrices

3343fab

some performance optimizations

c50264e

Isalia20 marked this pull request as ready for review January 5, 2025 09:05

Isalia20 requested review from kulinseth and malfet as code owners January 5, 2025 09:05

Isalia20 added 2 commits January 5, 2025 15:05

clone with contiguous format

f4e5be6

more optimizations

e199d6e

Isalia20 added 4 commits January 6, 2025 13:38

faster syrk

8a1372e

a bit faster syrk kernel

5aafb8a

coalesced access for right matrix

7eb7b4d

correct upper support

eb606a9

cpuhrsch requested a review from albanD January 7, 2025 06:36

cpuhrsch added the triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module label Jan 7, 2025

Isalia20 added 3 commits January 7, 2025 13:00

support other dtypes(float16, bfloat16)

83b0c66

check for positive definite matrix

6987e24

more performance optimizations

06b4419

pytorchmergebot added Reverted ci-no-td Do not run TD on this PR labels Jan 16, 2025

pytorchmergebot reopened this Jan 16, 2025

remove linalg_cholesky from lowering fallback

dc7617d

malfet mentioned this pull request Jan 26, 2025

General MPS op coverage tracking issue #77764

Open

malfet approved these changes Jan 26, 2025

View reviewed changes

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Jan 26, 2025

pytorchmergebot added the merging label Jan 26, 2025

pytorchmergebot removed the merging label Jan 26, 2025

Isalia20 mentioned this pull request Jan 26, 2025

[MPS] cholesky implementation #145701

Closed

Isalia20 closed this Jan 26, 2025

pytorchmergebot pushed a commit that referenced this pull request Jan 27, 2025

[MPS] cholesky implementation (#145701)

b75afa2

Requested in #77764 Closed #144193 due to a lot of conflicts when rebasing Pull Request resolved: #145701 Approved by: https://github.com/malfet

Conversation

Isalia20 commented Jan 4, 2025 • edited by pytorch-bot bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Jan 4, 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/144193

❌ 1 New Failure, 26 Unrelated Failures

Uh oh!

github-actions bot commented Jan 4, 2025

Attention! native_functions.yaml was changed

Uh oh!

Isalia20 commented Jan 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Isalia20 commented Jan 5, 2025

Uh oh!

Isalia20 commented Jan 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorchmergebot commented Jan 16, 2025

Uh oh!

Isalia20 commented Jan 16, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Isalia20 commented Jan 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Isalia20 commented Jan 18, 2025

Uh oh!

Isalia20 commented Jan 22, 2025

Uh oh!

Isalia20 commented Jan 23, 2025

Uh oh!

malfet commented Jan 26, 2025

Uh oh!

pytorchmergebot commented Jan 26, 2025

Merge started

Uh oh!

pytorchmergebot commented Jan 26, 2025

Merge failed

Uh oh!

malfet commented Jan 26, 2025

Uh oh!

pytorchmergebot commented Jan 26, 2025

Uh oh!

pytorchmergebot commented Jan 26, 2025

Uh oh!

malfet commented Jan 26, 2025

Uh oh!

Isalia20 commented Jan 26, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

Isalia20 commented Jan 4, 2025 •

edited by pytorch-bot bot

Loading

pytorch-bot bot commented Jan 4, 2025 •

edited

Loading

Isalia20 commented Jan 5, 2025 •

edited

Loading

Isalia20 commented Jan 5, 2025 •

edited

Loading

Isalia20 commented Jan 16, 2025 •

edited

Loading

Isalia20 commented Jan 17, 2025 •

edited

Loading