[Dynamo] Add MPSDevice interface #143891
Closed
[Dynamo] Add MPSDevice interface #143891
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/143891
Note: Links to docs will display an error until the docs builds have been completed. ⏳ No Failures, 25 PendingAs of commit d4c6d4e with merge base 809106a ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
This was referenced Dec 27, 2024
malfet
added
topic: not user facing
jansel
approved these changes
Dec 27, 2024
|
@pytorchbot merge -f "All Mac tests are green" |
Merge startedYour change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes). Please use Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
pytorchmergebot
pushed a commit
that referenced
this pull request
Dec 28, 2024
Mostly dummy interface as MPS backend currently limited to a single device Pull Request resolved: #143892 Approved by: https://github.com/jansel ghstack dependencies: #143891
pytorchmergebot
pushed a commit
that referenced
this pull request
Dec 28, 2024
Still work in progress, only works for element wise operations. Current implementation could be used to turn something like
```python
def f(x):
return x[:,::2].sin() + x[:, 1::2].cos()
```
into the following shader
```python
# Topologically Sorted Source Nodes: [sin, cos, add], Original ATen: [aten.sin, aten.cos, aten.add]
# Source node to ATen node mapping:
# add => add
# cos => cos
# sin => sin
# Graph fragment:
# %sin : [num_users=1] = call_function[target=torch.ops.aten.sin.default](args = (%slice_2,), kwargs = {})
# %cos : [num_users=1] = call_function[target=torch.ops.aten.cos.default](args = (%slice_4,), kwargs = {})
# %add : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%sin, %cos), kwargs = {})
mps_lib = torch.mps._compile_shader("""
kernel void kernel_0(
device float* out_ptr0,
constant float* in_ptr0,
uint xindex [[thread_position_in_grid]]
) {
int x0 = xindex;
auto tmp0 = in_ptr0[2*x0];
auto tmp1 = metal::precise::sin(tmp0);
auto tmp2 = in_ptr0[2*x0 + 1];
auto tmp3 = metal::precise::cos(tmp2);
auto tmp4 = tmp1 + tmp3;
out_ptr0[x0] = static_cast<float>(tmp4);
}
""")
```
Pull Request resolved: #143893
Approved by: https://github.com/jansel
ghstack dependencies: #143891, #143892
galexite
pushed a commit
to graphcore/pytorch-fork
that referenced
this pull request
Dec 30, 2024
That simply checks if device is available and whether or not it supports bf16 Pull Request resolved: pytorch#143891 Approved by: https://github.com/jansel
galexite
pushed a commit
to graphcore/pytorch-fork
that referenced
this pull request
Dec 30, 2024
Mostly dummy interface as MPS backend currently limited to a single device Pull Request resolved: pytorch#143892 Approved by: https://github.com/jansel ghstack dependencies: pytorch#143891
galexite
pushed a commit
to graphcore/pytorch-fork
that referenced
this pull request
Dec 30, 2024
Still work in progress, only works for element wise operations. Current implementation could be used to turn something like
```python
def f(x):
return x[:,::2].sin() + x[:, 1::2].cos()
```
into the following shader
```python
# Topologically Sorted Source Nodes: [sin, cos, add], Original ATen: [aten.sin, aten.cos, aten.add]
# Source node to ATen node mapping:
# add => add
# cos => cos
# sin => sin
# Graph fragment:
# %sin : [num_users=1] = call_function[target=torch.ops.aten.sin.default](args = (%slice_2,), kwargs = {})
# %cos : [num_users=1] = call_function[target=torch.ops.aten.cos.default](args = (%slice_4,), kwargs = {})
# %add : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%sin, %cos), kwargs = {})
mps_lib = torch.mps._compile_shader("""
kernel void kernel_0(
device float* out_ptr0,
constant float* in_ptr0,
uint xindex [[thread_position_in_grid]]
) {
int x0 = xindex;
auto tmp0 = in_ptr0[2*x0];
auto tmp1 = metal::precise::sin(tmp0);
auto tmp2 = in_ptr0[2*x0 + 1];
auto tmp3 = metal::precise::cos(tmp2);
auto tmp4 = tmp1 + tmp3;
out_ptr0[x0] = static_cast<float>(tmp4);
}
""")
```
Pull Request resolved: pytorch#143893
Approved by: https://github.com/jansel
ghstack dependencies: pytorch#143891, pytorch#143892
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Stack from ghstack (oldest at bottom):
That simply checks if device is available and whether or not it supports bf16
cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @chauhang @amjames