|
3 | 3 | # This source code is licensed under the BSD-style license found in the |
4 | 4 | # LICENSE file in the root directory of this source tree. |
5 | 5 |
|
6 | | -from collections.abc import Callable |
| 6 | +from collections.abc import Callable, Sequence |
7 | 7 | from typing import Tuple |
8 | 8 |
|
9 | 9 | import pytest |
@@ -720,6 +720,109 @@ def test_down_pass_moves_matching_input_permutations_after_binary_op() -> None: |
720 | 720 | assert targets.index(ADD) < targets.index(PERMUTE) |
721 | 721 |
|
722 | 722 |
|
| 723 | +def test_down_pass_moves_permute_through_single_exit_pointwise_region() -> None: |
| 724 | + """Move a dimension reorder after related per-value operations. |
| 725 | +
|
| 726 | + Here ``P`` means permute, Both ``floor`` and ``sub`` process values independently: |
| 727 | +
|
| 728 | + Before: P(x) --+-----------------> sub -> output |
| 729 | + | ^ |
| 730 | + +-> floor -----------+ |
| 731 | +
|
| 732 | + After: x --+-----------------> sub -> P -> output |
| 733 | + | ^ |
| 734 | + +-> floor -----------+ |
| 735 | +
|
| 736 | + The first check confirms that the permute moved after both operations. The second |
| 737 | + check confirms that the rewritten graph still calculates the same values. |
| 738 | +
|
| 739 | + """ |
| 740 | + graph = torch.fx.Graph() |
| 741 | + x = graph.placeholder("x") |
| 742 | + x.meta["val"] = torch.empty((1, 2, 3, 4)) |
| 743 | + permute = graph.call_function(PERMUTE, args=(x, [0, 2, 3, 1])) |
| 744 | + permute.meta["val"] = torch.empty((1, 3, 4, 2)) |
| 745 | + floor = graph.call_function(FLOOR, args=(permute,)) |
| 746 | + floor.meta["val"] = torch.empty((1, 3, 4, 2)) |
| 747 | + sub = graph.call_function(SUB, args=(permute, floor)) |
| 748 | + sub.meta["val"] = torch.empty((1, 3, 4, 2)) |
| 749 | + graph.output(sub) |
| 750 | + |
| 751 | + x_data = torch.randn(1, 2, 3, 4) |
| 752 | + expected = x_data.permute(0, 2, 3, 1) |
| 753 | + expected = expected - expected.floor() |
| 754 | + graph_module = _run_pass_on_graph_module(graph, PropagateViewCopyPermuteDownPass) |
| 755 | + call_nodes = [ |
| 756 | + node for node in graph_module.graph.nodes if node.op == "call_function" |
| 757 | + ] |
| 758 | + targets = [node.target for node in call_nodes] |
| 759 | + |
| 760 | + assert targets.index(FLOOR) < targets.index(SUB) < targets.index(PERMUTE) |
| 761 | + torch.testing.assert_close(graph_module(x_data), expected) |
| 762 | + |
| 763 | + |
| 764 | +class _BlockFloorToSubPass(PropagateViewCopyPermuteDownPass): |
| 765 | + def blocks_moving( |
| 766 | + self, |
| 767 | + moving_node: torch.fx.Node, |
| 768 | + frontier: torch.fx.Node, |
| 769 | + next_nodes: Sequence[torch.fx.Node], |
| 770 | + ) -> bool: |
| 771 | + return frontier.target == FLOOR and any( |
| 772 | + next_node.target == SUB for next_node in next_nodes |
| 773 | + ) |
| 774 | + |
| 775 | + |
| 776 | +def test_down_pass_checks_each_path_into_shared_pointwise_node() -> None: |
| 777 | + """Keep a permute when a backend rejects one path into a shared node.""" |
| 778 | + graph = torch.fx.Graph() |
| 779 | + x = graph.placeholder("x") |
| 780 | + x.meta["val"] = torch.empty((1, 2, 3, 4)) |
| 781 | + permute = graph.call_function(PERMUTE, args=(x, [0, 2, 3, 1])) |
| 782 | + permute.meta["val"] = torch.empty((1, 3, 4, 2)) |
| 783 | + floor = graph.call_function(FLOOR, args=(permute,)) |
| 784 | + floor.meta["val"] = torch.empty((1, 3, 4, 2)) |
| 785 | + sub = graph.call_function(SUB, args=(permute, floor)) |
| 786 | + sub.meta["val"] = torch.empty((1, 3, 4, 2)) |
| 787 | + graph.output(sub) |
| 788 | + |
| 789 | + targets = _run_pass_on_graph(graph, _BlockFloorToSubPass) |
| 790 | + |
| 791 | + assert targets.index(PERMUTE) < targets.index(FLOOR) < targets.index(SUB) |
| 792 | + |
| 793 | + |
| 794 | +def test_down_pass_keeps_permute_before_multiple_exit_pointwise_region() -> None: |
| 795 | + """Keep a dimension reorder in place when two results are returned. |
| 796 | +
|
| 797 | + Here ``P`` means permute. Both ``floor`` and ``sub`` are returned from the graph: |
| 798 | +
|
| 799 | + P(x) --+-----------------> sub -> output 1 |
| 800 | + | ^ |
| 801 | + +-> floor -----------+ |
| 802 | + | |
| 803 | + +---------------> output 0 |
| 804 | +
|
| 805 | + Moving the permute would require adding a separate reorder before each output. |
| 806 | + This rewrite deliberately handles only one outgoing result, so the graph |
| 807 | + must remain unchanged. |
| 808 | +
|
| 809 | + """ |
| 810 | + graph = torch.fx.Graph() |
| 811 | + x = graph.placeholder("x") |
| 812 | + x.meta["val"] = torch.empty((1, 2, 3, 4)) |
| 813 | + permute = graph.call_function(PERMUTE, args=(x, [0, 2, 3, 1])) |
| 814 | + permute.meta["val"] = torch.empty((1, 3, 4, 2)) |
| 815 | + floor = graph.call_function(FLOOR, args=(permute,)) |
| 816 | + floor.meta["val"] = torch.empty((1, 3, 4, 2)) |
| 817 | + sub = graph.call_function(SUB, args=(permute, floor)) |
| 818 | + sub.meta["val"] = torch.empty((1, 3, 4, 2)) |
| 819 | + graph.output((floor, sub)) |
| 820 | + |
| 821 | + targets = _run_pass_on_graph(graph, PropagateViewCopyPermuteDownPass) |
| 822 | + |
| 823 | + assert targets.index(PERMUTE) < targets.index(FLOOR) < targets.index(SUB) |
| 824 | + |
| 825 | + |
723 | 826 | def test_down_pass_keeps_lower_rank_singleton_input_before_binary_op() -> None: |
724 | 827 | graph = torch.fx.Graph() |
725 | 828 | x = graph.placeholder("x") |
|
0 commit comments