Support Apple MPS acceleration #1129

ClaudiaComito · 2023-03-29T10:21:42Z

LAST EDITED DEC 12 2025
[Note from human: the most important changes in this PR are:

Apple MPS are now a valid GPU device in device.py. Both ht.array(..., device="gpu") and ht.array(..., device="mps") are allowed.
device attribute introduced for ht.random.permutation
added an item to the PR template checklist to test with MPS (manually for now until the CI is expanded, see Expand CI to macos-m1 #1747 )
codecov is unhappy but the introduced changes cannot be tested with our current setup.

Still to do:

~~update README~~ actually, it's probably best to update the README just before the next release

Below a copilot summary]

This pull request includes several changes to improve compatibility with Apple's Metal Performance Shaders (MPS) and correct some minor issues. The most important changes include modifications to handle unsupported data types on MPS, updates to unit tests, and minor corrections in documentation.

MPS Compatibility Improvements:

heat/core/_operations.py: Added checks to handle unsupported float64 data type on MPS and cast to float32 with appropriate warnings. [1] [2] [3]
heat/core/arithmetics.py: Updated hypot and hypot_ functions to raise errors for unsupported int64 data type on MPS. [1] [2]
heat/core/dndarray.py: Modified the size method to avoid using float64 on MPS.
heat/core/linalg/basics.py: Added a check to raise a RuntimeError if matrix inversion fails on MPS.

Unit Test Updates:

heat/cluster/tests/test_batchparallelclustering.py, heat/cluster/tests/test_kmeans.py, heat/cluster/tests/test_kmedians.py, heat/cluster/tests/test_kmedoids.py: Updated tests to handle unsupported float64 data type on MPS. [1] [2] [3] [4]
heat/cluster/tests/test_spectral.py: Added a condition to skip tests on MPS due to unsupported ComplexFloat operations.
heat/core/linalg/tests/test_basics.py: Updated tests to avoid using float64 on MPS. [1] [2] [3]

Minor Corrections:

heat/core/devices.py: Corrected documentation to use consistent naming for Heat. [1] [2]
heat/core/linalg/solver.py: Changed tensor creation to use int64 instead of int32 for cumulative sum operations. [1] [2]

Reference

Current status of PyTorch's MPS operations coverage

Issue/s resolved: #1053

Changes proposed:

Type of change

Memory requirements

Performance

Due Diligence

~~All split configurations tested~~ does not apply
Multiple dtypes tested in relevant functions
Documentation updated (if needed)
Title of PR is suitable for corresponding CHANGELOG entry

Does this change modify the behaviour of other functions? If so, which?

no

ghost · 2023-03-29T10:22:52Z

👇 Click on the image for a new way to code review

Legend

github-actions · 2023-03-29T10:26:32Z

Thank you for the PR!

codecov · 2023-03-29T10:41:02Z

Codecov Report

Attention: Patch coverage is 69.67213% with 37 lines in your changes missing coverage. Please review.

Project coverage is 91.99%. Comparing base (443afe3) to head (80a867e).

Files with missing lines	Patch %	Lines
heat/core/arithmetics.py	72.41%	8 Missing ⚠️
heat/core/tests/test_suites/basic_test.py	52.94%	8 Missing ⚠️
heat/core/devices.py	0.00%	7 Missing ⚠️
heat/core/_operations.py	42.85%	4 Missing ⚠️
heat/core/manipulations.py	86.95%	3 Missing ⚠️
heat/core/statistics.py	87.50%	2 Missing ⚠️
heat/core/dndarray.py	75.00%	1 Missing ⚠️
heat/core/linalg/basics.py	83.33%	1 Missing ⚠️
heat/core/relational.py	75.00%	1 Missing ⚠️
heat/core/signal.py	66.66%	1 Missing ⚠️
... and 1 more

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1129      +/-   ##
==========================================
- Coverage   92.26%   91.99%   -0.28%     
==========================================
  Files          84       84              
  Lines       12445    12535      +90     
==========================================
+ Hits        11482    11531      +49     
- Misses        963     1004      +41

Flag	Coverage Δ
unit	`91.99% <69.67%> (-0.28%)`	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

…Apple-silicon-GPUs

github-actions · 2023-03-29T10:51:15Z

Thank you for the PR!

github-actions · 2023-04-17T08:49:40Z

Thank you for the PR!

github-actions · 2023-04-18T08:54:44Z

Thank you for the PR!

github-actions · 2023-04-24T16:52:42Z

Thank you for the PR!

github-actions · 2023-04-27T09:03:30Z

Thank you for the PR!

…Apple-silicon-GPUs

github-actions · 2023-05-18T08:33:39Z

Thank you for the PR!

…:helmholtz-analytics/heat into features/1053-support-Apple-silicon-GPUs

ClaudiaComito · 2024-12-03T09:15:26Z

Tests failed; I rerun them to check whether thats just a HW problem.

Thanks, sadly it looks like an actual problem, conveniently without error message. I'll debug it, will probably get to it next week.

github-actions · 2024-12-10T04:55:16Z

Thank you for the PR!

github-actions · 2024-12-10T11:55:31Z

Thank you for the PR!

github-actions · 2024-12-11T04:47:29Z

Thank you for the PR!

github-actions · 2024-12-11T05:20:49Z

Thank you for the PR!

github-actions · 2024-12-12T04:50:13Z

Thank you for the PR!

github-actions · 2024-12-12T05:04:05Z

Thank you for the PR!

github-actions · 2024-12-12T09:37:53Z

Thank you for the PR!

github-actions · 2024-12-12T11:15:23Z

Thank you for the PR!

JuanPedroGHM · 2024-12-17T15:48:28Z

heat/cluster/tests/test_batchparallelclustering.py

+        if self.is_mps:
+            dtypes = [ht.float32]
+        else:
+            dtypes = [ht.float32, ht.float64]
+


This (and all subsequent tests that have to filter by system) would be a great target for parametrization (now that we talked about introducing hypothesis and parametrized tests).

A good example on how to skip certain possible parameters based on the os is here

JuanPedroGHM · 2024-12-17T15:53:46Z

heat/core/linalg/solver.py

@@ -339,7 +339,7 @@ def solve_triangular(A: DNDarray, b: DNDarray) -> DNDarray:
        else:  # A not split, b.split == -2
            b_lshapes_cum = torch.hstack(
                [
-                    torch.zeros(1, dtype=torch.int32, device=tdev),
+                    torch.zeros(1, dtype=torch.int64, device=tdev),


Is there a reason for change? Why not use the default dtype?

JuanPedroGHM · 2024-12-17T15:57:10Z

heat/core/linalg/tests/test_basics.py

@@ -2154,19 +2168,20 @@ def test_triu(self):
            self.assertTrue(result.larray[0, -1] == 1)

    def test_vdot(self):
-        a = ht.array([[1 + 1j, 2 + 2j], [3 + 3j, 4 + 4j]], split=0)
-        b = ht.array([[1 + 2j, 3 + 4j], [5 + 6j, 7 + 8j]], split=0)
+        if not self.is_mps:


Test should be skipped using unittest.skipIf or pytest.mark.skipif

JuanPedroGHM · 2024-12-17T15:57:54Z

heat/core/linalg/tests/test_qr.py

-            ht.allclose(q.transpose([0, 1, 3, 2]) @ q, batched_id, atol=1e-6, rtol=1e-6)
-        )
-        self.assertTrue(ht.allclose(q @ r, x, atol=1e-6, rtol=1e-6))
+        # skip float64 tests on MPS


Test should be skipped using unittest.skipIf or pytest.mark.skipif

JuanPedroGHM · 2024-12-17T15:58:37Z

heat/core/linalg/tests/test_svdtools.py

-        ]
-        rtols = [1e-1, 1e-2, 1e-3]
-        ranks = [5, 10, 15]
+        # not testing on MPS for now as torch.norm() is unstable


Test should be skipped using unittest.skipIf or pytest.mark.skipif

JuanPedroGHM · 2024-12-17T16:03:08Z

heat/core/relational.py

+    is_mps = x.larray.is_mps or y.larray.is_mps
+    if is_mps and result_type is types.float64:
+        result_type = types.float32


Instead of checking every time after calling types.result_type, the check could be done inside types.result_type(). This would save a lot of extra if statements and less chance of possibly forgetting to add that.

JuanPedroGHM · 2024-12-17T16:03:50Z

heat/core/signal.py

+    if a.larray.is_mps and promoted_type == float64:
+        # cannot cast to float64 on MPS
+        promoted_type = float32
+


Same with promote_types.

ClaudiaComito and others added 4 commits November 26, 2022 08:00

Support Apple MPS, first pass

4e9ed3a

Merge branch 'main' into features/1053-support-Apple-silicon-GPUs

50297d7

Include torch 2.0 in device check

febdcfc

Merge branch 'main' into features/1053-support-Apple-silicon-GPUs

a5642e9

reinstate quick_start.md

8445476

ClaudiaComito changed the title ~~Features/1053 support apple silicon gp us~~ Support Apple MPS acceleration Mar 29, 2023

ClaudiaComito added 2 commits March 29, 2023 12:46

Merge branch 'docs/reinstate-quick-start' into features/1053-support-…

8763146

…Apple-silicon-GPUs

[skip ci] edits

92306e1

Merge branch 'main' into features/1053-support-Apple-silicon-GPUs

b0d7f0f

ClaudiaComito added this to the 1.3.0 milestone Apr 17, 2023

ClaudiaComito self-assigned this Apr 17, 2023

Merge branch 'main' into features/1053-support-Apple-silicon-GPUs

ffa014e

Merge branch 'main' into features/1053-support-Apple-silicon-GPUs

6015ebf

Merge branch 'main' into features/1053-support-Apple-silicon-GPUs

f3a5ad8

ClaudiaComito added 5 commits April 27, 2023 12:09

fix tolerance for torch 2

a96441b

implement __array__ method

c7b70c6

test __array__ method

ff8af94

test __array__ method

fc059c2

Merge branch 'features/1153_array_method' into features/1053-support-…

ebf4c51

…Apple-silicon-GPUs

ClaudiaComito added 2 commits May 22, 2023 06:13

Merge branch 'main' into features/1053-support-Apple-silicon-GPUs

5c471e2

Merge branch 'features/1053-support-Apple-silicon-GPUs' of github.com…

34c9347

…:helmholtz-analytics/heat into features/1053-support-Apple-silicon-GPUs

adapt test_permutation

0efad5e

ClaudiaComito added 2 commits December 10, 2024 12:26

Merge branch 'main' into features/1053-support-Apple-silicon-GPUs

0a6b6ed

return inv as DNDarray, not Tensor

164b7b7

ClaudiaComito added 2 commits December 11, 2024 05:22

skip DMD tests on MPS

c7cdf34

increase allclose tolerance for test_inv

dc2ab84

ClaudiaComito added 3 commits December 11, 2024 06:02

skip line formatting

6907c71

skip line formatting

58cc44f

skip line formatting

1e902d1

debugging test_sort on AMD

71e83c5

update test_iris

fa3e900

indices sorting workaround for CUDA

103dc7e

ClaudiaComito added PR talk HW:MPS labels Dec 12, 2024

update PR template, docs

80a867e

ClaudiaComito mentioned this pull request Dec 12, 2024

Expand CI to macos-m1 #1747

Open

ClaudiaComito requested a review from mtar December 12, 2024 11:12

JuanPedroGHM self-requested a review December 16, 2024 08:58

ClaudiaComito requested review from JuanPedroGHM and removed request for JuanPedroGHM December 16, 2024 08:58

ClaudiaComito mentioned this pull request Dec 16, 2024

Update documentation for usage with MPS #1752

Open

12 tasks

JuanPedroGHM reviewed Dec 17, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Support Apple MPS acceleration #1129

Support Apple MPS acceleration #1129

ClaudiaComito commented Mar 29, 2023 •

edited

Loading

ghost commented Mar 29, 2023 •

edited by ghost

Loading

Legend

github-actions bot commented Mar 29, 2023

codecov bot commented Mar 29, 2023 •

edited

Loading

github-actions bot commented Mar 29, 2023

github-actions bot commented Apr 17, 2023

github-actions bot commented Apr 18, 2023

github-actions bot commented Apr 24, 2023

github-actions bot commented Apr 27, 2023

github-actions bot commented May 18, 2023

ClaudiaComito commented Dec 3, 2024

github-actions bot commented Dec 10, 2024

github-actions bot commented Dec 10, 2024

github-actions bot commented Dec 11, 2024

github-actions bot commented Dec 11, 2024

github-actions bot commented Dec 12, 2024

github-actions bot commented Dec 12, 2024

github-actions bot commented Dec 12, 2024

github-actions bot commented Dec 12, 2024

JuanPedroGHM Dec 17, 2024

JuanPedroGHM Dec 17, 2024

JuanPedroGHM Dec 17, 2024

JuanPedroGHM Dec 17, 2024

JuanPedroGHM Dec 17, 2024

JuanPedroGHM Dec 17, 2024

JuanPedroGHM Dec 17, 2024

Support Apple MPS acceleration #1129

Are you sure you want to change the base?

Support Apple MPS acceleration #1129

Conversation

ClaudiaComito commented Mar 29, 2023 • edited Loading

MPS Compatibility Improvements:

Unit Test Updates:

Minor Corrections:

Reference

Changes proposed:

Type of change

Memory requirements

Performance

Due Diligence

Does this change modify the behaviour of other functions? If so, which?

ghost commented Mar 29, 2023 • edited by ghost Loading

Legend

github-actions bot commented Mar 29, 2023

codecov bot commented Mar 29, 2023 • edited Loading

Codecov Report

github-actions bot commented Mar 29, 2023

github-actions bot commented Apr 17, 2023

github-actions bot commented Apr 18, 2023

github-actions bot commented Apr 24, 2023

github-actions bot commented Apr 27, 2023

github-actions bot commented May 18, 2023

ClaudiaComito commented Dec 3, 2024

github-actions bot commented Dec 10, 2024

github-actions bot commented Dec 10, 2024

github-actions bot commented Dec 11, 2024

github-actions bot commented Dec 11, 2024

github-actions bot commented Dec 12, 2024

github-actions bot commented Dec 12, 2024

github-actions bot commented Dec 12, 2024

github-actions bot commented Dec 12, 2024

JuanPedroGHM Dec 17, 2024

Choose a reason for hiding this comment

JuanPedroGHM Dec 17, 2024

Choose a reason for hiding this comment

JuanPedroGHM Dec 17, 2024

Choose a reason for hiding this comment

JuanPedroGHM Dec 17, 2024

Choose a reason for hiding this comment

JuanPedroGHM Dec 17, 2024

Choose a reason for hiding this comment

JuanPedroGHM Dec 17, 2024

Choose a reason for hiding this comment

JuanPedroGHM Dec 17, 2024

Choose a reason for hiding this comment

ClaudiaComito commented Mar 29, 2023 •

edited

Loading

ghost commented Mar 29, 2023 •

edited by ghost

Loading

codecov bot commented Mar 29, 2023 •

edited

Loading