[MPS] Implement hardshrink metal kernel #155304

manuelcandales · 2025-06-06T02:53:31Z

Stack from ghstack (oldest at bottom):

Implements the forward and backward hardshrink operators as Metal kernels.
In order to support the lambda parameter, we extend the exec_unary_kernel and exec_binary_kernel methods. Now they take an optional Scalar and an optional ScalarType argument. When the optional ScalarType is provided, it overrides the type of the Scalar.
We add a new REGISTER_UNARY_ALPHA_OP macro, and modify the existing REGISTER_BINARY_ALPHA_OP to support the new feature.

[ghstack-poisoned]

pytorch-bot · 2025-06-06T02:53:34Z

🔗 Helpful Links

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

📄 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 Pending, 9 Unrelated Failures

As of commit 525ae01 with merge base d4d0ede ():

FLAKY - The following job failed but was likely due to flakiness present on trunk:

pull / linux-jammy-cuda12.6-py3.10-gcc11-sm89 / test (default, 3, 5, lf.linux.g6.4xlarge.experimental.nvidia.gpu) (gh) (disabled by #153395)
test_foreach.py::TestForeachCUDA::test_parity__foreach_add_fastpath_inplace_cuda_float64

BROKEN TRUNK - The following jobs failed but was present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

pull / cuda12.8-py3.10-gcc9-sm75 / test (pr_time_benchmarks, 1, 1, linux.g4dn.metal.nvidia.gpu) (gh) (trunk failure)
MISSING REGRESSION TEST
pull / linux-jammy-cuda12.6-py3.10-gcc11 / test (default, 5, 5, lf.linux.4xlarge.nvidia.gpu) (gh) (trunk failure)
dynamo/test_graph_deduplication.py::GraphDededuplicationTests::test_autocast_ordering
pull / linux-jammy-cuda12.6-py3.10-gcc11-sm89 / test (default, 5, 5, lf.linux.g6.4xlarge.experimental.nvidia.gpu) (gh) (trunk failure)
dynamo/test_graph_deduplication.py::GraphDededuplicationTests::test_autocast_ordering
pull / linux-jammy-py3.10-clang15-asan / test (default, 5, 6, lf.linux.4xlarge) (gh) (trunk failure)
dynamo/test_graph_deduplication.py::GraphDededuplicationTests::test_autocast_ordering
pull / linux-jammy-py3.13-clang12 / test (default, 2, 5, lf.linux.4xlarge) (gh) (trunk failure)
dynamo/test_graph_deduplication.py::GraphDededuplicationTests::test_autocast_ordering
pull / linux-jammy-py3.9-clang12 / test (default, 5, 5, lf.linux.4xlarge) (gh) (trunk failure)
dynamo/test_graph_deduplication.py::GraphDededuplicationTests::test_autocast_ordering
pull / linux-jammy-py3.9-gcc11 / test (default, 5, 5, lf.linux.2xlarge) (gh) (trunk failure)
dynamo/test_graph_deduplication.py::GraphDededuplicationTests::test_autocast_ordering

UNSTABLE - The following job is marked as unstable, possibly due to flakiness on trunk:

pull / linux-jammy-py3-clang12-executorch / test (executorch, 1, 1, lf.linux.2xlarge, unstable) (gh)
exir/backend/test/test_to_backend_multi_method.py::TestToBackendMultiMethod::test_multi_method_end_to_end

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

ghstack-source-id: 6f45291 Pull-Request: #155304

github-actions · 2025-06-06T02:57:33Z

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

[ghstack-poisoned]

ghstack-source-id: 0aed9db Pull-Request: #155304

malfet

type API looks a bit fragile, consider using std::pair and encode extra dtype in shader name

aten/src/ATen/native/mps/OperationUtils.mm

aten/src/ATen/native/mps/MetalShaderLibrary.h

[ghstack-poisoned]

malfet · 2025-06-09T15:58:42Z

aten/src/ATen/native/mps/MetalShaderLibrary.h

+      const std::optional<c10::Scalar> alpha = std::nullopt,
+      const std::optional<c10::ScalarType> scalar_arg_type = std::nullopt);


This API looks wrong, it should be something like

Suggested change

const std::optional<c10::Scalar> alpha = std::nullopt,

const std::optional<c10::ScalarType> scalar_arg_type = std::nullopt);

const std::optional<std::pair<c10::Scalar, c10::ScalarType>> alpha = std::nullopt,

c10/metal/indexing.h

aten/src/ATen/native/mps/kernels/ActivationKernel.metal

[ghstack-poisoned]

malfet

Thank you for working on the change

malfet · 2025-06-10T18:18:39Z

@pytorchbot merge -f "Lint + MPS are green, will submit some cleanups later on"

pytorchmergebot · 2025-06-10T18:20:14Z

Merge started

Your change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes). Please use -f as last resort and instead consider -i/--ignore-current to continue the merge ignoring current failures. This will allow currently pending tests to finish and report signal before the merge.

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

…155316) Pull Request resolved: #155316 Approved by: https://github.com/Skylion007, https://github.com/malfet ghstack dependencies: #155304

…5462) Pull Request resolved: #155462 Approved by: https://github.com/malfet ghstack dependencies: #155304, #155316

Pull Request resolved: #155479 Approved by: https://github.com/kulinseth, https://github.com/malfet ghstack dependencies: #155304, #155316, #155462

) Pull Request resolved: #155571 Approved by: https://github.com/malfet ghstack dependencies: #155304, #155316, #155462, #155479

) Pull Request resolved: #155586 Approved by: https://github.com/malfet ghstack dependencies: #155304, #155316, #155462, #155479, #155571

Implements the forward and backward hardshrink operators as Metal kernels. In order to support the lambda parameter, we extend the `exec_unary_kernel` and `exec_binary_kernel` methods. Now they take an optional Scalar and an optional ScalarType argument. When the optional ScalarType is provided, it overrides the type of the Scalar. We add a new `REGISTER_UNARY_ALPHA_OP` macro, and modify the existing `REGISTER_BINARY_ALPHA_OP` to support the new feature. Pull Request resolved: pytorch#155304 Approved by: https://github.com/malfet

…ytorch#155316) Pull Request resolved: pytorch#155316 Approved by: https://github.com/Skylion007, https://github.com/malfet ghstack dependencies: pytorch#155304

…orch#155462) Pull Request resolved: pytorch#155462 Approved by: https://github.com/malfet ghstack dependencies: pytorch#155304, pytorch#155316

…ch#155479) Pull Request resolved: pytorch#155479 Approved by: https://github.com/kulinseth, https://github.com/malfet ghstack dependencies: pytorch#155304, pytorch#155316, pytorch#155462

…rch#155571) Pull Request resolved: pytorch#155571 Approved by: https://github.com/malfet ghstack dependencies: pytorch#155304, pytorch#155316, pytorch#155462, pytorch#155479

…rch#155586) Pull Request resolved: pytorch#155586 Approved by: https://github.com/malfet ghstack dependencies: pytorch#155304, pytorch#155316, pytorch#155462, pytorch#155479, pytorch#155571

Pull Request resolved: #155735 Approved by: https://github.com/malfet ghstack dependencies: #155304, #155316, #155462, #155479, #155571, #155586

Update

f39e761

[ghstack-poisoned]

manuelcandales requested review from kulinseth and malfet as code owners June 6, 2025 02:53

pytorch-bot bot added ciflow/mps Run MPS tests (subset of trunk) release notes: mps Release notes category labels Jun 6, 2025

manuelcandales added a commit that referenced this pull request Jun 6, 2025

[MPS] Implement hardshrink metal kernel

3cda0e2

ghstack-source-id: 6f45291 Pull-Request: #155304

manuelcandales marked this pull request as draft June 6, 2025 04:53

Update

f5676f4

[ghstack-poisoned]

manuelcandales added a commit that referenced this pull request Jun 6, 2025

[MPS] Implement hardshrink metal kernel

f3e9849

ghstack-source-id: 0aed9db Pull-Request: #155304

manuelcandales mentioned this pull request Jun 6, 2025

[MPS][BE] Refactor round_decimals shader code to leverage new macro #155316

Closed

manuelcandales marked this pull request as ready for review June 6, 2025 05:51

malfet requested changes Jun 6, 2025

View reviewed changes

aten/src/ATen/native/mps/OperationUtils.mm Outdated Show resolved Hide resolved

aten/src/ATen/native/mps/MetalShaderLibrary.h Outdated Show resolved Hide resolved

manuelcandales added 4 commits June 6, 2025 13:56

Update

e393dba

[ghstack-poisoned]

Update

0ac7f8e

[ghstack-poisoned]

Update

b474056

[ghstack-poisoned]

Update

f2e018b

[ghstack-poisoned]

manuelcandales requested a review from malfet June 7, 2025 02:12

manuelcandales mentioned this pull request Jun 9, 2025

[MPS] Migrate hardsigmoid (forward and backward) to Metal kernel #155462

Closed

malfet reviewed Jun 9, 2025

View reviewed changes

c10/metal/indexing.h Outdated Show resolved Hide resolved

malfet reviewed Jun 9, 2025

View reviewed changes

aten/src/ATen/native/mps/kernels/ActivationKernel.metal Outdated Show resolved Hide resolved

manuelcandales added 3 commits June 9, 2025 14:19

Update

2fdce93

[ghstack-poisoned]

Update

c635f81

[ghstack-poisoned]

Update

3a26396

[ghstack-poisoned]

manuelcandales mentioned this pull request Jun 9, 2025

[MPS] Migrate hardswish (forward and backward) to Metal kernel #155479

Closed

manuelcandales added 3 commits June 9, 2025 16:00

Update

d48911b

[ghstack-poisoned]

Update

fcef81e

[ghstack-poisoned]

Update

5314497

[ghstack-poisoned]

Update

525ae01

[ghstack-poisoned]

manuelcandales requested a review from malfet June 10, 2025 14:03

manuelcandales mentioned this pull request Jun 10, 2025

[MPS] Migrate leaky_relu (forward and backward) to Metal kernel #155571

Closed

malfet approved these changes Jun 10, 2025

View reviewed changes

malfet added the topic: improvements topic category label Jun 10, 2025

pytorchmergebot added the merging label Jun 10, 2025

pytorchmergebot closed this in 0f47e76 Jun 10, 2025

pytorchmergebot added Merged and removed merging labels Jun 10, 2025

manuelcandales mentioned this pull request Jun 5, 2025

Most requested ops for the MPS backend #154052

Open

github-actions bot deleted the gh/manuelcandales/2/head branch July 13, 2025 02:22

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[MPS] Implement hardshrink metal kernel #155304

[MPS] Implement hardshrink metal kernel #155304

Uh oh!

manuelcandales commented Jun 6, 2025 •

edited

Loading

Uh oh!

pytorch-bot bot commented Jun 6, 2025 •

edited

Loading

Uh oh!

github-actions bot commented Jun 6, 2025

Uh oh!

malfet left a comment

Uh oh!

Uh oh!

Uh oh!

malfet Jun 9, 2025

Uh oh!

Uh oh!

Uh oh!

malfet left a comment

Uh oh!

malfet commented Jun 10, 2025

Uh oh!

pytorchmergebot commented Jun 10, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

		const std::optional<c10::Scalar> alpha = std::nullopt,
		const std::optional<c10::ScalarType> scalar_arg_type = std::nullopt);

	const std::optional<c10::Scalar> alpha = std::nullopt,
	const std::optional<c10::ScalarType> scalar_arg_type = std::nullopt);
	const std::optional<std::pair<c10::Scalar, c10::ScalarType>> alpha = std::nullopt,

[MPS] Implement hardshrink metal kernel #155304

[MPS] Implement hardshrink metal kernel #155304

Uh oh!

Conversation

manuelcandales commented Jun 6, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Jun 6, 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/155304

⏳ 1 Pending, 9 Unrelated Failures

Uh oh!

github-actions bot commented Jun 6, 2025

Attention! native_functions.yaml was changed

Uh oh!

malfet left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

malfet Jun 9, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

malfet left a comment

Choose a reason for hiding this comment

Uh oh!

malfet commented Jun 10, 2025

Uh oh!

pytorchmergebot commented Jun 10, 2025

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

manuelcandales commented Jun 6, 2025 •

edited

Loading

pytorch-bot bot commented Jun 6, 2025 •

edited

Loading