Metal backend: Add operator implementations #15023
Conversation
|
Stack from ghstack (oldest at bottom): |
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/15023
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 New Failure, 1 Unrelated FailureAs of commit 4367977 with merge base 6e0c9f6 ( NEW FAILURE - The following job has failed:
FLAKY - The following job failed but was likely due to flakiness present on trunk:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
bot
added
the
CLA Signed
manuelcandales
requested review from
larryliu0820 and
mergennachin
and removed request for
cccclai and
shoumikhin
| * ExecutorTorch implementation of aoti_torch_mps_mm_out. | ||
| * Performs simple matrix multiplication: out = self @ mat2 | ||
| */ | ||
| AOTITorchError aoti_torch_mps_mm_out( |
There was a problem hiding this comment.
Does custom ops use caching mechanism like the ETMetalShaderLibrary?
There was a problem hiding this comment.
No, not yet. These fallback ops are implemented using MPSGraph, so, here we would be caching the graph. This is something I want to look into later when optimizing performance. But this deserves time. In particular, since I never understood why MPSGraph operations have a non-trivial CPU overhead in PyTorch, in spite of PyTorch having a caching mechanism for MPSGraphs.
| // For attention weights, zero-fill the GPU buffer (shared memory allows CPU memset) | ||
| std::memset(attn_contents_ptr, 0, attn_size_bytes); |
There was a problem hiding this comment.
do you need zero filling here
There was a problem hiding this comment.
Well, I though it was nicer to return 0, rather than some random stuff.
|
|
||
| // Set output tensor handles | ||
| *ret0 = out_tensor_handle; | ||
| *ret1 = attn_tensor_handle; |
There was a problem hiding this comment.
Is ret1 actually populated or just zerod
There was a problem hiding this comment.
This is just zeroed.
We are using MPSGraph's scaledDotProductAttention which only returns the output tensor.
We need to return an attention tensor because we need to match _scaled_dot_product_attention_math_for_mps signature. But we don't really need it, it gets thrown away here
manuelcandales
added
the
release notes: none
Adds bfloat16/float32 working implementations of the following AOTI shim ops: - aoti_torch_mps_mm_out - aoti_torch_mps_convolution - aoti_torch_mps__scaled_dot_product_attention_math_for_mps Adds a stub implementation of aoti_torch_mps_addmm_out ghstack-source-id: 61b8cc4 ghstack-comment-id: 3392300522 Pull-Request: pytorch#15023
2 participants
Adds bfloat16/float32 working implementations of the following AOTI shim ops:
Adds a stub implementation of aoti_torch_mps_addmm_out