|
8 | 8 | # pyre-unsafe |
9 | 9 |
|
10 | 10 | import copy |
| 11 | +import itertools |
11 | 12 | import unittest |
12 | 13 | from typing import cast |
13 | 14 |
|
|
38 | 39 | from executorch.exir import EdgeCompileConfig, to_edge |
39 | 40 | from executorch.exir.dialects._ops import ops as exir_ops |
40 | 41 | from executorch.exir.pass_base import PassResult |
| 42 | +from parameterized import parameterized |
41 | 43 | from torch.utils import _pytree as pytree |
42 | 44 |
|
43 | 45 |
|
@@ -2734,6 +2736,276 @@ def test_lower_rank_constant_is_transposed_into_region_layout(self) -> None: |
2734 | 2736 | ) |
2735 | 2737 |
|
2736 | 2738 |
|
| 2739 | +_NHWC_SHAPE = [2, 3, 4, 5] |
| 2740 | +_AXIS_NAMES = "NHWC" |
| 2741 | +# x.permute(_NHWC_FROM_NCHW) turns an NCHW tensor into NHWC. |
| 2742 | +_NHWC_FROM_NCHW = [0, 2, 3, 1] |
| 2743 | + |
| 2744 | + |
| 2745 | +def _nchw_order(nhwc: list[int]) -> list[int]: |
| 2746 | + return [nhwc[0], nhwc[3], nhwc[1], nhwc[2]] |
| 2747 | + |
| 2748 | + |
| 2749 | +def _case_name(broadcast: tuple[int, ...], kept_states: dict[int, str]) -> str: |
| 2750 | + parts = ( |
| 2751 | + ["broadcast_" + "".join(_AXIS_NAMES[a] for a in broadcast)] if broadcast else [] |
| 2752 | + ) |
| 2753 | + parts += [f"{_AXIS_NAMES[a]}_{state}" for a, state in kept_states.items()] |
| 2754 | + return "_".join(parts) |
| 2755 | + |
| 2756 | + |
| 2757 | +def _region_expand_cases() -> list[tuple[str, tuple[int, ...], dict[int, str]]]: |
| 2758 | + cases = [] |
| 2759 | + for count in range(1, 5): |
| 2760 | + for broadcast in itertools.combinations(range(4), count): |
| 2761 | + kept = [a for a in range(4) if a not in broadcast] |
| 2762 | + for states in itertools.product( |
| 2763 | + ("size", "minus_one", "dynamic"), repeat=len(kept) |
| 2764 | + ): |
| 2765 | + kept_states = dict(zip(kept, states)) |
| 2766 | + cases.append( |
| 2767 | + (_case_name(broadcast, kept_states), broadcast, kept_states) |
| 2768 | + ) |
| 2769 | + return cases |
| 2770 | + |
| 2771 | + |
| 2772 | +def _constant_expand_cases() -> list[tuple[str, int, tuple[int, ...], dict[int, str]]]: |
| 2773 | + cases = [] |
| 2774 | + for rank in range(1, 4): |
| 2775 | + own_axes = list(range(4 - rank, 4)) |
| 2776 | + for count in range(rank + 1): |
| 2777 | + for broadcast in itertools.combinations(own_axes, count): |
| 2778 | + kept = [a for a in own_axes if a not in broadcast] |
| 2779 | + for states in itertools.product( |
| 2780 | + ("size", "minus_one"), repeat=len(kept) |
| 2781 | + ): |
| 2782 | + kept_states = dict(zip(kept, states)) |
| 2783 | + name = f"rank{rank}_" + _case_name(broadcast, kept_states) |
| 2784 | + cases.append((name.rstrip("_"), rank, broadcast, kept_states)) |
| 2785 | + return cases |
| 2786 | + |
| 2787 | + |
| 2788 | +def _export_to_edge( |
| 2789 | + module: torch.nn.Module, |
| 2790 | + example_inputs: tuple[torch.Tensor, ...], |
| 2791 | + dynamic_shapes: dict[str, dict[int, torch.export.Dim] | None] | None, |
| 2792 | +) -> torch.fx.GraphModule: |
| 2793 | + exported = torch.export.export( |
| 2794 | + module, example_inputs, dynamic_shapes=dynamic_shapes |
| 2795 | + ) |
| 2796 | + edge = to_edge(exported, compile_config=EdgeCompileConfig(_check_ir_validity=False)) |
| 2797 | + return edge.exported_program().graph_module |
| 2798 | + |
| 2799 | + |
| 2800 | +class RemovePermutesAroundExpandCopyTest(unittest.TestCase): |
| 2801 | + """An NHWC region around ``expand_copy``, for every set of broadcast axes, |
| 2802 | + with kept axes written as their size, as -1, or as -1 on a dynamic dim.""" |
| 2803 | + |
| 2804 | + def _remove_permutes(self, original: torch.fx.GraphModule) -> torch.fx.GraphModule: |
| 2805 | + result = cast(PassResult, RemovePermutesAroundElementwiseOps()(original)) |
| 2806 | + self.assertTrue(result.modified) |
| 2807 | + gm = result.graph_module |
| 2808 | + # A constant whose non-unit dims change order keeps a permute of its |
| 2809 | + # own; only the activation path has to be permute-free. |
| 2810 | + permute = exir_ops.edge.aten.permute_copy.default |
| 2811 | + (x,) = (n for n in gm.graph.nodes if n.op == "placeholder" and n.name == "x") |
| 2812 | + self.assertNotIn(permute, [user.target for user in x.users]) |
| 2813 | + (output,) = gm.graph.find_nodes(op="output") |
| 2814 | + self.assertNotEqual(output.args[0][0].target, permute) |
| 2815 | + return gm |
| 2816 | + |
| 2817 | + def _expand_size(self, gm: torch.fx.GraphModule) -> list[int]: |
| 2818 | + (expand,) = gm.graph.find_nodes( |
| 2819 | + op="call_function", target=exir_ops.edge.aten.expand_copy.default |
| 2820 | + ) |
| 2821 | + return expand.args[1] |
| 2822 | + |
| 2823 | + @parameterized.expand(_region_expand_cases()) |
| 2824 | + def test_expand_of_region_tensor( |
| 2825 | + self, name: str, broadcast: tuple[int, ...], kept_states: dict[int, str] |
| 2826 | + ) -> None: |
| 2827 | + dynamic_axes = [a for a, state in kept_states.items() if state == "dynamic"] |
| 2828 | + size = [ |
| 2829 | + -1 if kept_states.get(a) in ("minus_one", "dynamic") else _NHWC_SHAPE[a] |
| 2830 | + for a in range(4) |
| 2831 | + ] |
| 2832 | + |
| 2833 | + class PermuteExpandPermute(torch.nn.Module): |
| 2834 | + def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 2835 | + return x.permute(0, 2, 3, 1).expand(*size).permute(0, 3, 1, 2) |
| 2836 | + |
| 2837 | + input_nhwc = [1 if a in broadcast else _NHWC_SHAPE[a] for a in range(4)] |
| 2838 | + x = torch.randn(*_nchw_order(input_nhwc)) |
| 2839 | + dims = { |
| 2840 | + _NHWC_FROM_NCHW[a]: torch.export.Dim(f"d{a}", min=2, max=16) |
| 2841 | + for a in dynamic_axes |
| 2842 | + } |
| 2843 | + original = _export_to_edge(PermuteExpandPermute(), (x,), {"x": dims or None}) |
| 2844 | + gm_before = copy.deepcopy(original) |
| 2845 | + |
| 2846 | + gm = self._remove_permutes(original) |
| 2847 | + |
| 2848 | + expected_nhwc = [-1 if a in dynamic_axes else _NHWC_SHAPE[a] for a in range(4)] |
| 2849 | + self.assertEqual(self._expand_size(gm), _nchw_order(expected_nhwc)) |
| 2850 | + resized_nhwc = [ |
| 2851 | + d + 3 if a in dynamic_axes else d for a, d in enumerate(input_nhwc) |
| 2852 | + ] |
| 2853 | + for inputs in ([x], [torch.randn(*_nchw_order(resized_nhwc))]): |
| 2854 | + validate_numerics(gm_before, gm, inputs, name) |
| 2855 | + |
| 2856 | + @parameterized.expand(_constant_expand_cases()) |
| 2857 | + def test_rank_growing_expand_of_constant( |
| 2858 | + self, |
| 2859 | + name: str, |
| 2860 | + rank: int, |
| 2861 | + broadcast: tuple[int, ...], |
| 2862 | + kept_states: dict[int, str], |
| 2863 | + ) -> None: |
| 2864 | + bias_shape = [ |
| 2865 | + 1 if a in broadcast else _NHWC_SHAPE[a] for a in range(4 - rank, 4) |
| 2866 | + ] |
| 2867 | + size = [ |
| 2868 | + -1 if kept_states.get(a) == "minus_one" else _NHWC_SHAPE[a] |
| 2869 | + for a in range(4) |
| 2870 | + ] |
| 2871 | + builder = GraphBuilder() |
| 2872 | + x_data = torch.randn(*_nchw_order(_NHWC_SHAPE)) |
| 2873 | + bias_data = torch.randn(*bias_shape) |
| 2874 | + x = builder.placeholder("x", x_data) |
| 2875 | + bias = builder.placeholder("p_bias", bias_data) |
| 2876 | + nhwc = builder.call_operator( |
| 2877 | + op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 2, 3, 1]) |
| 2878 | + ) |
| 2879 | + expanded = builder.call_operator( |
| 2880 | + op=exir_ops.edge.aten.expand_copy.default, args=(bias, size) |
| 2881 | + ) |
| 2882 | + scaled = builder.call_operator( |
| 2883 | + op=exir_ops.edge.aten.mul.Tensor, args=(nhwc, expanded) |
| 2884 | + ) |
| 2885 | + nchw = builder.call_operator( |
| 2886 | + op=exir_ops.edge.aten.permute_copy.default, args=(scaled, [0, 3, 1, 2]) |
| 2887 | + ) |
| 2888 | + builder.output([nchw]) |
| 2889 | + original = builder.get_graph_module() |
| 2890 | + gm_before = copy.deepcopy(original) |
| 2891 | + |
| 2892 | + gm = self._remove_permutes(original) |
| 2893 | + |
| 2894 | + self.assertEqual(self._expand_size(gm), _nchw_order(_NHWC_SHAPE)) |
| 2895 | + validate_numerics(gm_before, gm, [x_data, bias_data], name) |
| 2896 | + |
| 2897 | + def test_keep_entry_moved_onto_an_added_dim(self) -> None: |
| 2898 | + # Undoing the [1, 0, 2, 3] swap moves the -1 onto the dim the expand |
| 2899 | + # adds, so it has to be spelled out as 1 first. |
| 2900 | + builder = GraphBuilder() |
| 2901 | + x_data = torch.randn(1, 4, 5, 3) |
| 2902 | + bias_data = torch.randn(1, 5, 3) |
| 2903 | + x = builder.placeholder("x", x_data) |
| 2904 | + bias = builder.placeholder("p_bias", bias_data) |
| 2905 | + swapped = builder.call_operator( |
| 2906 | + op=exir_ops.edge.aten.permute_copy.default, args=(x, [1, 0, 2, 3]) |
| 2907 | + ) |
| 2908 | + expanded = builder.call_operator( |
| 2909 | + op=exir_ops.edge.aten.expand_copy.default, args=(bias, [4, -1, 5, 3]) |
| 2910 | + ) |
| 2911 | + scaled = builder.call_operator( |
| 2912 | + op=exir_ops.edge.aten.mul.Tensor, args=(swapped, expanded) |
| 2913 | + ) |
| 2914 | + back = builder.call_operator( |
| 2915 | + op=exir_ops.edge.aten.permute_copy.default, args=(scaled, [1, 0, 2, 3]) |
| 2916 | + ) |
| 2917 | + builder.output([back]) |
| 2918 | + original = builder.get_graph_module() |
| 2919 | + gm_before = copy.deepcopy(original) |
| 2920 | + |
| 2921 | + gm = self._remove_permutes(original) |
| 2922 | + |
| 2923 | + self.assertEqual(self._expand_size(gm), [1, 4, 5, 3]) |
| 2924 | + validate_numerics(gm_before, gm, [x_data, bias_data], "added_dim") |
| 2925 | + |
| 2926 | + @parameterized.expand([(f"symbolic_target_{_AXIS_NAMES[a]}", a) for a in range(4)]) |
| 2927 | + def test_symbolic_broadcast_target_leaves_region( |
| 2928 | + self, name: str, axis: int |
| 2929 | + ) -> None: |
| 2930 | + class ExpandToSymbolicSize(torch.nn.Module): |
| 2931 | + def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor: |
| 2932 | + size = [-1] * 4 |
| 2933 | + size[axis] = t.shape[0] |
| 2934 | + return x.permute(0, 2, 3, 1).expand(*size).permute(0, 3, 1, 2) |
| 2935 | + |
| 2936 | + input_nhwc = [1 if a == axis else _NHWC_SHAPE[a] for a in range(4)] |
| 2937 | + x = torch.randn(*_nchw_order(input_nhwc)) |
| 2938 | + original = _export_to_edge( |
| 2939 | + ExpandToSymbolicSize(), |
| 2940 | + (x, torch.randn(6)), |
| 2941 | + {"x": None, "t": {0: torch.export.Dim("t", min=2, max=16)}}, |
| 2942 | + ) |
| 2943 | + gm_before = copy.deepcopy(original) |
| 2944 | + |
| 2945 | + result = cast(PassResult, RemovePermutesAroundElementwiseOps()(original)) |
| 2946 | + |
| 2947 | + self.assertFalse(result.modified) |
| 2948 | + self.assertEqual( |
| 2949 | + count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 2 |
| 2950 | + ) |
| 2951 | + for t_size in (6, 9): |
| 2952 | + validate_numerics( |
| 2953 | + gm_before, result.graph_module, [x, torch.randn(t_size)], name |
| 2954 | + ) |
| 2955 | + |
| 2956 | + |
| 2957 | +class SymbolicScalarInputTest(unittest.TestCase): |
| 2958 | + """A region whose ops take a symbolic size is left alone rather than |
| 2959 | + crashing the pass.""" |
| 2960 | + |
| 2961 | + def test_symbolic_arguments_leave_the_region_unchanged(self) -> None: |
| 2962 | + def nhwc(x: torch.Tensor) -> torch.Tensor: |
| 2963 | + return x.permute(0, 2, 3, 1) |
| 2964 | + |
| 2965 | + def nchw(x: torch.Tensor) -> torch.Tensor: |
| 2966 | + return x.permute(0, 3, 1, 2) |
| 2967 | + |
| 2968 | + class SymbolicArgument(torch.nn.Module): |
| 2969 | + def __init__(self, op: str) -> None: |
| 2970 | + super().__init__() |
| 2971 | + self.op = op |
| 2972 | + |
| 2973 | + def forward(self, x: torch.Tensor, n: torch.Tensor) -> torch.Tensor: |
| 2974 | + y = nhwc(x) |
| 2975 | + if self.op == "expand": |
| 2976 | + y = y.expand(2, n.shape[0], 8, 3) |
| 2977 | + elif self.op == "slice": |
| 2978 | + y = y[:, :, : n.shape[0], :] |
| 2979 | + else: |
| 2980 | + y = torch.nn.functional.pad(y, (0, 0, 0, n.shape[0])) |
| 2981 | + return nchw(y) |
| 2982 | + |
| 2983 | + for op in ("expand", "slice", "pad"): |
| 2984 | + with self.subTest(op=op): |
| 2985 | + x = torch.randn(2, 3, 1 if op == "expand" else 4, 8) |
| 2986 | + original = _export_to_edge( |
| 2987 | + SymbolicArgument(op), |
| 2988 | + (x, torch.randn(4)), |
| 2989 | + {"x": None, "n": {0: torch.export.Dim("n", min=2, max=7)}}, |
| 2990 | + ) |
| 2991 | + gm_before = copy.deepcopy(original) |
| 2992 | + |
| 2993 | + result = cast( |
| 2994 | + PassResult, RemovePermutesAroundElementwiseOps()(original) |
| 2995 | + ) |
| 2996 | + |
| 2997 | + self.assertFalse(result.modified) |
| 2998 | + self.assertEqual( |
| 2999 | + count_node( |
| 3000 | + result.graph_module, exir_ops.edge.aten.permute_copy.default |
| 3001 | + ), |
| 3002 | + 2, |
| 3003 | + ) |
| 3004 | + validate_numerics( |
| 3005 | + gm_before, result.graph_module, [x, torch.randn(6)], op |
| 3006 | + ) |
| 3007 | + |
| 3008 | + |
2737 | 3009 | class LayoutPermuteVisibilityTest(unittest.TestCase): |
2738 | 3010 | """The data-movement passes must see both permute dialects. |
2739 | 3011 |
|
|
0 commit comments