Skip to content

[MPS] Add native implementation for shift ops #131813

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Closed
wants to merge 2 commits into from
Closed

[MPS] Add native implementation for shift ops #131813

wants to merge 2 commits into from

Conversation

malfet
Copy link
Contributor

@malfet malfet commented Jul 25, 2024

Similar to how AND/OR/XOR ops are implemented

TODO: Consider using MPS method calls rather than metal kernels

@malfet
[MPS] Add native implementation for shift ops
5642af3 
Similar to how AND/OR/XOR ops are implemented

TODO: Consider using MPS method calls rather than metal kernels
@malfet malfet requested a review from kulinseth as a code owner July 25, 2024 21:07
@pytorch-bot PyTorch Bot
Copy link

pytorch-bot bot commented Jul 25, 2024 •
edited
Loading

🔗 Helpful Links

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

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

❌ 1 New Failure

As of commit d1c7b2a with merge base bf6aae1 (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 ciflow/mps Run MPS tests (subset of trunk) release notes: mps Release notes category labels Jul 25, 2024
manuelcandales
manuelcandales reviewed Jul 25, 2024
View reviewed changes
aten/src/ATen/mps/MPSFallback.mm
@@ -85,8 +85,6 @@ static Tensor slow_conv2d_forward_mps(const Tensor& self,
// These ops are not supported via MPS backend currently, and we fallback to run on CPU.
// For the rest of unsupported ops the user needs to pass 'PYTORCH_ENABLE_MPS_FALLBACK=1'
// to fallback on CPU, otherwise we will error out.
m.impl("bitwise_left_shift.Tensor_out", torch::CppFunction::makeFromBoxedFunction<&mps_fallback>());
m.impl("bitwise_right_shift.Tensor_out", torch::CppFunction::makeFromBoxedFunction<&mps_fallback>());
Copy link
Contributor

Choose a reason for hiding this comment

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

If there was a fallback in place, why were things not working before?
Granted that the error was with __left__.Scalar. Not sure how that magic works

Copy link
Contributor

Choose a reason for hiding this comment

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

Oh, the fallback was for the Tensor variant, not for the Scalar ones

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Yeah, this only works if one calls torch.bitwise_left_shift, not sure what's the story with all those

manuelcandales
manuelcandales reviewed Jul 25, 2024
View reviewed changes
aten/src/ATen/native/native_functions.yaml
@@ -8457,29 +8457,29 @@
device_check: NoCheck # TensorIterator
variants: method, function
dispatch:
CPU, CUDA: __lshift__
CPU, CUDA, MPS: __lshift__
Copy link
Contributor

Choose a reason for hiding this comment

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

Just curious, where is the magic that ties __lshift__ to bitwise_left_shift_out

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Here:

pytorch/aten/src/ATen/native/BinaryOps.cpp

Line 1292 in 3d7c424

lshift_stub(iter.device_type(), iter);

@malfet
Copy link
Contributor Author

malfet commented Jul 26, 2024

@pytorchbot merge -f "MPS + Linr tests are green"

@pytorchmergebot
Copy link
Collaborator

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

@manuelcandales manuelcandales mentioned this pull request Jul 26, 2024
Merged
@henrylhtsang henrylhtsang mentioned this pull request Jul 31, 2024
Closed
@github-actions github-actions bot deleted the malfet/mps-add-shift-ops branch August 26, 2024 02:01
This was referenced Sep 6, 2024
Open
Closed
pytorchmergebot pushed a commit that referenced this pull request Sep 11, 2024
@malfet @pytorchmergebot
[MPS] Add missing dispatch to rshift.Tensor (#135607)
3bf6be4 
Missed it while working on #131813
Test plan: `python -c "import torch;print(torch.randint(100, 500, (64,), device='mps') >> torch.tensor([3,], device='mps'))"`

Pull Request resolved: #135607
Approved by: https://github.com/manuelcandales
yushangdi pushed a commit that referenced this pull request Sep 12, 2024
@malfet @yushangdi
[MPS] Add missing dispatch to rshift.Tensor (#135607)
1557dca 
Missed it while working on #131813
Test plan: `python -c "import torch;print(torch.randint(100, 500, (64,), device='mps') >> torch.tensor([3,], device='mps'))"`

Pull Request resolved: #135607
Approved by: https://github.com/manuelcandales
Chao1Han pushed a commit to Chao1Han/pytorch that referenced this pull request Sep 20, 2024
@malfet @Chao1Han
[MPS] Add missing dispatch to rshift.Tensor (pytorch#135607)
a6fa5d1 
Missed it while working on pytorch#131813
Test plan: `python -c "import torch;print(torch.randint(100, 500, (64,), device='mps') >> torch.tensor([3,], device='mps'))"`

Pull Request resolved: pytorch#135607
Approved by: https://github.com/manuelcandales
pytorchbot pushed a commit that referenced this pull request Oct 2, 2024
@malfet @pytorchbot
[MPS] Add missing dispatch to rshift.Tensor (#135607)
04b6ed6 
Missed it while working on #131813
Test plan: `python -c "import torch;print(torch.randint(100, 500, (64,), device='mps') >> torch.tensor([3,], device='mps'))"`

Pull Request resolved: #135607
Approved by: https://github.com/manuelcandales

(cherry picked from commit 3bf6be4)
@pytorchbot pytorchbot mentioned this pull request Oct 2, 2024
Merged
kit1980 pushed a commit that referenced this pull request Oct 2, 2024
@pytorchbot @malfet
[MPS] Add missing dispatch to rshift.Tensor (#137212)
5375201 
[MPS] Add missing dispatch to rshift.Tensor (#135607)

Missed it while working on #131813
Test plan: `python -c "import torch;print(torch.randint(100, 500, (64,), device='mps') >> torch.tensor([3,], device='mps'))"`

Pull Request resolved: #135607
Approved by: https://github.com/manuelcandales

(cherry picked from commit 3bf6be4)

Co-authored-by: Nikita Shulga <2453524+malfet@users.noreply.github.com>
@malfet malfet mentioned this pull request Jun 2, 2025
Open
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@manuelcandales manuelcandales manuelcandales approved these changes

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

Assignees
No one assigned
Labels
ciflow/mps Run MPS tests (subset of trunk) Merged release notes: mps Release notes category
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
3 participants
@malfet @pytorchmergebot @manuelcandales