Enable nvprims.transpose fusions for nvFuser #86967

IvanYashchuk · 2022-10-14T15:31:15Z

This PR allows transposes to be fused with other operations. If a fusion group is formed only from operations that just manipulate metadata in PyTorch (transpose, view, etc.) then this group is not sent to nvFuser.
On top of that if we have converted to nvprims but then decided to not form a fusion group we modify the graph use prim.impl_aten attribute instead of calling prim(*args, **kwargs) that has a higher overhead.

cc @kevinstephano @jjsjann123

…ityBasedPartitioner

pytorch-bot · 2022-10-14T15:31:19Z

🔗 Helpful Links

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

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit our office hours

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

✅ No Failures

As of commit b013d6f:
💚 Looks good so far! There are no failures yet. 💚

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

IvanYashchuk · 2022-10-14T15:31:58Z

@SherlockNoMad could you please review the torch/fx/passes/infra/partitioner.py changes?

``` commit 5863911 (ivan/nvprims-transpose-partitioner) Merge: 73669a7 fc3afc8 Author: Ivan Yashchuk <IvanYashchuk@users.noreply.github.com> Date: Fri Oct 14 22:58:55 2022 +0300 Merge branch 'master' into nvprims-transpose-partitioner ```

IvanYashchuk · 2022-10-17T09:05:12Z

@ngimel could you please take a look? It's purely nvFuser related change and doesn't modify anything else.

jjsjann123

LGTM.

jjsjann123 · 2022-10-17T17:14:16Z

torch/_prims/nvfuser_executor.py

@@ -268,6 +268,29 @@ def __call__(self, *args):
        )


+# A set of operators that are supported by nvFuser
+# but should not form a fusion group solely on their own
+# _non_compute_ops = {


nitpick: commented code should be removed.

It's an example, but it should have been marked so more clearly 😄

jjsjann123 · 2022-10-17T17:32:39Z

torch/_prims/nvfuser_executor.py

+#     "torch.ops.nvprims.broadcast_in_dim.default",
+#     "torch.ops.nvprims.squeeze.default",
+# }
+_non_compute_ops = [


nitpick: peeking through op return type and categorizing view-like return to non-compute ops is counter-intuitive. (thinking about in-place update).

I think this is safe for our uses, but should we rename this to _view_like_ops and add a line in the comment that view-like ops are likely no-compute ops, since in-place update is handled by functionalization

I agree with you, I used the same name as was already used in the partitioner code.
We also don't have the in-place update (not merged yet).

IvanYashchuk · 2022-10-17T20:16:45Z

@SherlockNoMad could you please review the torch/fx/passes/infra/partitioner.py changes?

@SherlockNoMad

SherlockNoMad

partitioner change looks good to me.

…anspose-partitioner

IvanYashchuk · 2022-10-26T14:21:28Z

@pytorchbot merge -g

pytorchmergebot · 2022-10-26T14:23:28Z

Merge started

Your change will be merged once all checks on your PR pass since you used the green (-g) flag (ETA: 0-4 Hours).

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

@kevinstephano

This PR allows transposes to be fused with other operations. If a fusion group is formed only from operations that just manipulate metadata in PyTorch (transpose, view, etc.) then this group is not sent to nvFuser. On top of that if we have converted to `nvprims` but then decided to not form a fusion group we modify the graph use `prim.impl_aten` attribute instead of calling `prim(*args, **kwargs)` that has a higher overhead. cc @kevinstephano @jjsjann123 Pull Request resolved: pytorch#86967 Approved by: https://github.com/jjsjann123, https://github.com/SherlockNoMad

@kevinstephano

This PR allows transposes to be fused with other operations. If a fusion group is formed only from operations that just manipulate metadata in PyTorch (transpose, view, etc.) then this group is not sent to nvFuser. On top of that if we have converted to `nvprims` but then decided to not form a fusion group we modify the graph use `prim.impl_aten` attribute instead of calling `prim(*args, **kwargs)` that has a higher overhead. cc @kevinstephano @jjsjann123 Pull Request resolved: pytorch#86967 Approved by: https://github.com/jjsjann123, https://github.com/SherlockNoMad

@kevinstephano

This PR allows transposes to be fused with other operations. If a fusion group is formed only from operations that just manipulate metadata in PyTorch (transpose, view, etc.) then this group is not sent to nvFuser. On top of that if we have converted to `nvprims` but then decided to not form a fusion group we modify the graph use `prim.impl_aten` attribute instead of calling `prim(*args, **kwargs)` that has a higher overhead. cc @kevinstephano @jjsjann123 Pull Request resolved: pytorch#86967 Approved by: https://github.com/jjsjann123, https://github.com/SherlockNoMad

IvanYashchuk added 5 commits October 14, 2022 18:20

Add non_compute_ops and allowed_single_node_partition_ops for Capabil…

7117d4c

…ityBasedPartitioner

Remove transpose and unsqueeze from aten_ops_to_skip

b70b902

Add _non_compute_ops to nvfuser partitioner

ab27563

Replace unfused nvprims with torch calls

c56ff6b

Set default allow_single_op_fusion to False

748743d

IvanYashchuk added the module: nvfuser label Oct 14, 2022

IvanYashchuk requested a review from jjsjann123 October 14, 2022 15:31

pytorch-bot bot added the release notes: fx release notes category label Oct 14, 2022

pytorchbot added the open source label Oct 14, 2022

IvanYashchuk requested a review from SherlockNoMad October 14, 2022 15:31

IvanYashchuk changed the title ~~Enable transpose fusions for nvFuser~~ Enable nvprims.transpose fusions for nvFuser Oct 14, 2022

IvanYashchuk and others added 2 commits October 14, 2022 18:55

ufmt

73669a7

Merge branch 'master' into nvprims-transpose-partitioner

5863911

IvanYashchuk requested a review from ngimel October 14, 2022 19:59

jjsjann123 mentioned this pull request Oct 14, 2022

Mega issue tracking torchbenchPerf on benchmark runs csarofeen/pytorch#2065

Open

26 tasks

IvanYashchuk added 3 commits October 15, 2022 12:14

Add a test for nvprims.view partitioning and replacing with aten impl

a922236

Fix test

b486ad0

Fix nvprims.view forming fusion groups

7944194

jjsjann123 approved these changes Oct 17, 2022

View reviewed changes

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Oct 17, 2022

SherlockNoMad approved these changes Oct 18, 2022

View reviewed changes

IvanYashchuk added 3 commits October 18, 2022 12:55

Merge remote-tracking branch 'upstream/viable/strict' into nvprims-tr…

b815962

…anspose-partitioner

Remove example _non_compute_ops

fd1b4ba

Merge remote-tracking branch 'upstream/viable/strict' into nvprims-tr…

d1b4bd2

…anspose-partitioner

IvanYashchuk added the topic: not user facing topic category label Oct 26, 2022

Merge branch 'master' into nvprims-transpose-partitioner

b013d6f

pytorchmergebot added the Merged label Oct 26, 2022

pytorchmergebot closed this in ae4fbac Oct 26, 2022

IvanYashchuk deleted the nvprims-transpose-partitioner branch October 26, 2022 18:47

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Enable nvprims.transpose fusions for nvFuser #86967

Enable nvprims.transpose fusions for nvFuser #86967

IvanYashchuk commented Oct 14, 2022 •

edited by pytorch-bot bot

pytorch-bot bot commented Oct 14, 2022 •

edited

IvanYashchuk commented Oct 14, 2022

IvanYashchuk commented Oct 17, 2022

jjsjann123 left a comment

jjsjann123 Oct 17, 2022

IvanYashchuk Oct 17, 2022

jjsjann123 Oct 17, 2022

IvanYashchuk Oct 17, 2022

IvanYashchuk commented Oct 17, 2022

SherlockNoMad left a comment

IvanYashchuk commented Oct 26, 2022

pytorchmergebot commented Oct 26, 2022

Enable nvprims.transpose fusions for nvFuser #86967

Enable nvprims.transpose fusions for nvFuser #86967

Conversation

IvanYashchuk commented Oct 14, 2022 • edited by pytorch-bot bot

pytorch-bot bot commented Oct 14, 2022 • edited

🔗 Helpful Links

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

✅ No Failures

IvanYashchuk commented Oct 14, 2022

IvanYashchuk commented Oct 17, 2022

jjsjann123 left a comment

Choose a reason for hiding this comment

jjsjann123 Oct 17, 2022

Choose a reason for hiding this comment

IvanYashchuk Oct 17, 2022

Choose a reason for hiding this comment

jjsjann123 Oct 17, 2022

Choose a reason for hiding this comment

IvanYashchuk Oct 17, 2022

Choose a reason for hiding this comment

IvanYashchuk commented Oct 17, 2022

SherlockNoMad left a comment

Choose a reason for hiding this comment

IvanYashchuk commented Oct 26, 2022

pytorchmergebot commented Oct 26, 2022

Merge started

IvanYashchuk commented Oct 14, 2022 •

edited by pytorch-bot bot

pytorch-bot bot commented Oct 14, 2022 •

edited