KEMBAR78
[MPS] Add CompileShader method by malfet · Pull Request #141478 · pytorch/pytorch · GitHub
Skip to content

[MPS] Add CompileShader method #141478

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 19 commits into from
Closed

[MPS] Add CompileShader method #141478

wants to merge 19 commits into from

Conversation

@malfet
Copy link
Contributor

@malfet malfet commented Nov 25, 2024 •
edited
Loading

Stack from ghstack (oldest at bottom):

This allows one to do something like that

import torch
x = torch.ones(10, device="mps")
m = torch.mps._compile_shader("""
   kernel void foo(device float* x, uint idx [[thread_position_in_grid]]) {
     x[idx] += idx;
   }
")
m.foo(x)

And in general enables writing custom operators using Metal shaders purely in Python

@malfet malfet requested a review from kulinseth as a code owner November 25, 2024 06:11
This was referenced Nov 25, 2024
Closed
Closed
@pytorch-bot
Copy link

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

🔗 Helpful Links

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

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

✅ No Failures

As of commit 8b827fa with merge base 95e7fcf (image):
💚 Looks good so far! There are no failures yet. 💚

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

@malfet malfet mentioned this pull request Nov 25, 2024
Closed
@pytorch-bot pytorch-bot bot added ciflow/mps Run MPS tests (subset of trunk) release notes: mps Release notes category labels Nov 25, 2024
@malfet malfet mentioned this pull request Nov 25, 2024
Closed
malfet added a commit that referenced this pull request Nov 25, 2024
@malfet
[WIP] Add CompileShader method
20f7b09 
ghstack-source-id: 9622dee
Pull Request resolved: #141478
@malfet malfet marked this pull request as draft November 25, 2024 06:11
@malfet malfet mentioned this pull request Nov 25, 2024
Closed
malfet added a commit that referenced this pull request Nov 25, 2024
@malfet
[WIP] Add CompileShader method
e71d1e5 
ghstack-source-id: 98879ec
Pull Request resolved: #141478
malfet added a commit that referenced this pull request Nov 25, 2024
@malfet
[WIP] Add CompileShader method
ca2a771 
ghstack-source-id: 2709938
Pull Request resolved: #141478
@malfet malfet mentioned this pull request Nov 26, 2024
Closed
malfet added a commit that referenced this pull request Nov 26, 2024
@malfet
[WIP] Add CompileShader method
5e6b17f 
ghstack-source-id: fec0ed8
Pull Request resolved: #141478
@malfet
Update
c3b2962 
[ghstack-poisoned]
@qqaatw
Copy link
Collaborator

qqaatw commented Nov 28, 2024

I'm curious about the use of this addition since it accepts one argument only. I assume there is no cuda counterpart in PyTorch.

malfet added 3 commits December 2, 2024 12:00
@malfet
Update
71fef69 
[ghstack-poisoned]
@malfet
Update
c7c1186 
[ghstack-poisoned]
@malfet
Update
9e5dc13 
[ghstack-poisoned]
@malfet
Copy link
Contributor Author

malfet commented Dec 3, 2024

I'm curious about the use of this addition since it accepts one argument only. I assume there is no cuda counterpart in PyTorch.

@qqaatw it's still work in progress :) But final version will accept any number of arguments. As for usefulness, at the very least it unblocks quick prototyping and usecases like shader toys, but in the long run it can enable a pretty advanced codegen.

@malfet
Update
737ee7c 
[ghstack-poisoned]
This was referenced Dec 4, 2024
Closed
Closed
malfet added 2 commits December 4, 2024 09:34
@malfet
Update
0731d6e 
[ghstack-poisoned]
@malfet
Update
bd441cd 
[ghstack-poisoned]
@malfet
Copy link
Contributor Author

malfet commented Dec 11, 2024

@pytorchbot revert -m "Broke doctests, by trying to run MPS example on Linux" -c nosignal

@pytorchmergebot
Copy link
Collaborator

@pytorchbot successfully started a revert job. Check the current status here.
Questions? Feedback? Please reach out to the PyTorch DevX Team

pytorchmergebot added a commit that referenced this pull request Dec 11, 2024
@pytorchmergebot
Revert "[MPS] Add CompileShader method (#141478)"
393cf46 
This reverts commit 0478fee.

Reverted #141478 on behalf of https://github.com/malfet due to Broke doctests, by trying to run MPS example on Linux ([comment](#141478 (comment)))
@pytorchmergebot
Copy link
Collaborator

@malfet your PR has been successfully reverted.

@pytorchmergebot pytorchmergebot added Reverted ci-no-td Do not run TD on this PR labels Dec 11, 2024
@malfet malfet added the no-runner-experiments Bypass Meta/LF runner determinator label Dec 11, 2024
@malfet
Update
f391cba 
[ghstack-poisoned]
@malfet
Update
8b827fa 
[ghstack-poisoned]
@malfet malfet requested a review from a team as a code owner December 11, 2024 01:36
@malfet
Copy link
Contributor Author

malfet commented Dec 11, 2024

@pytorchbot merge -f "No pain no gain"

@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

Esquains pushed a commit to Esquains/study1 that referenced this pull request Dec 15, 2024
@malfet
[WIP] Add CompileShader method
cb2c11e 
ghstack-source-id: 0d9e621
Pull Request resolved: pytorch/pytorch#141478
@github-actions github-actions bot deleted the gh/malfet/71/head branch January 11, 2025 02:10
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

@albanD albanD Awaiting requested review from albanD

Assignees

No one assigned

Labels

ci-no-td Do not run TD on this PR ciflow/mps Run MPS tests (subset of trunk) Merged no-runner-experiments Bypass Meta/LF runner determinator release notes: mps Release notes category Reverted topic: improvements topic category

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@malfet @qqaatw @pytorchmergebot @manuelcandales