Implement unfold_backward on MPS #135411
Implement unfold_backward on MPS #135411
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/135411
Note: Links to docs will display an error until the docs builds have been completed. ❗ 1 Active SEVsThere are 1 currently active SEVs. If your PR is affected, please view them below: ❌ 1 New FailureAs of commit 15a2972 with merge base 7578a0b ( NEW FAILURE - The following job has failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
Attention! native_functions.yaml was changedIf 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: |
malfet
added
topic: improvements
|
Looks like this PR hasn't been updated in a while so we're going to go ahead and mark this as |
malfet
changed the title
malfet
force-pushed
the
malfet-patch-14
branch
from
89d8310 to
e19d8f5
Compare
Co-authored-by: Manuel Candales <42380156+manuelcandales@users.noreply.github.com>
|
@pytorchbot merge -f "This was mostly green in the past" |
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 |
This PR adds native implementation of unfold_backward as metal shader, mostly copy-n-paste of algorithms used in CUDA and CPU implementations, i.e. considering `out = in.unfold(dim, size, step)`, then following holds true: * `out.shape[dim] == (in.shape[dim] - size) / step + 1` * `out.shape[-1] == size` * `out.ndim == in.ndim + 1` `unfold_backward` Metal kernel receives `grad_in` and returns `grad_out` such that: * `grad_in.shape == out.shape` * `grad_out.shape == in.shape` For each index in `grad_out` find the elements contributing to it and sum them up. Such algorithm requires no synchronization between threads. That is `grad_out[...,out_dim_idx,...]` accumulates all values `grad_in[...,in_dim_idx,...,in_last_idx]`, where `in_dim_idx` is range [`(out_dim_idx - size) / step`, `out_dim_idx / step`] clamped to (0, `in_dim_size`) and `in_last_idx` are equal `out_dim_idx - in_dim_idx * step` . Accumulation step is skipped if `in_last_idx` is outside of [0, size] range. This operator has been requested 16 times on pytorch#77764 Pull Request resolved: pytorch#135411 Approved by: https://github.com/manuelcandales Co-authored-by: Manuel Candales <42380156+manuelcandales@users.noreply.github.com>
3 participants
This PR adds native implementation of unfold_backward as metal shader, mostly copy-n-paste of algorithms used in CUDA and CPU implementations, i.e. considering
out = in.unfold(dim, size, step), then following holds true:out.shape[dim] == (in.shape[dim] - size) / step + 1out.shape[-1] == sizeout.ndim == in.ndim + 1unfold_backwardMetal kernel receivesgrad_inand returnsgrad_outsuch that:grad_in.shape == out.shapegrad_out.shape == in.shapeFor each index in
grad_outfind the elements contributing to it and sum them up. Such algorithm requires no synchronization between threads.That is
grad_out[...,out_dim_idx,...]accumulates all valuesgrad_in[...,in_dim_idx,...,in_last_idx], wherein_dim_idxis range [(out_dim_idx - size) / step,out_dim_idx / step] clamped to (0,in_dim_size) andin_last_idxare equalout_dim_idx - in_dim_idx * step. Accumulation step is skipped ifin_last_idxis outside of [0, size] range.This operator has been requested 16 times on #77764