diff --git a/backends/transforms/channels_last_ops.py b/backends/transforms/channels_last_ops.py index cbb182ccb02..54a897c7b80 100644 --- a/backends/transforms/channels_last_ops.py +++ b/backends/transforms/channels_last_ops.py @@ -55,6 +55,44 @@ def _avg_pool2d( return out.permute(0, 2, 3, 1).contiguous() +def _adaptive_avg_pool2d(input, output_size): + nchw = input.permute(0, 3, 1, 2) + out = torch.ops.aten.adaptive_avg_pool2d(nchw, output_size) + return out.permute(0, 2, 3, 1).contiguous() + + +def _upsample_bilinear2d(input, output_size, align_corners, scale_factors): + nchw = input.permute(0, 3, 1, 2) + out = torch.ops.aten.upsample_bilinear2d.vec( + nchw, output_size, align_corners, scale_factors + ) + return out.permute(0, 2, 3, 1).contiguous() + + +def _upsample_nearest2d(input, output_size, scale_factors): + nchw = input.permute(0, 3, 1, 2) + out = torch.ops.aten.upsample_nearest2d.vec(nchw, output_size, scale_factors) + return out.permute(0, 2, 3, 1).contiguous() + + +def _max_pool2d_with_indices(input, kernel_size, stride, padding, dilation, ceil_mode): + nchw = input.permute(0, 3, 1, 2) + values, indices = torch.ops.aten.max_pool2d_with_indices( + nchw, kernel_size, stride, padding, dilation, ceil_mode + ) + values = values.permute(0, 2, 3, 1).contiguous() + indices = indices.permute(0, 2, 3, 1).contiguous() + return values, indices + + +def _grid_sampler_2d(input, grid, interpolation_mode, padding_mode, align_corners): + nchw = input.permute(0, 3, 1, 2) + out = torch.ops.aten.grid_sampler_2d( + nchw, grid, interpolation_mode, padding_mode, align_corners + ) + return out.permute(0, 2, 3, 1).contiguous() + + def _permute_copy(input, dims): return torch.ops.aten.permute_copy(input, dims).contiguous() @@ -74,6 +112,42 @@ def _permute_copy(input, dims): lib.impl("avg_pool2d", _avg_pool2d, "CompositeExplicitAutograd") register_fake("channels_last::avg_pool2d", _avg_pool2d, lib=lib) +lib.define("adaptive_avg_pool2d(Tensor input, int[2] output_size) -> Tensor") +lib.impl("adaptive_avg_pool2d", _adaptive_avg_pool2d, "CompositeExplicitAutograd") +register_fake("channels_last::adaptive_avg_pool2d", _adaptive_avg_pool2d, lib=lib) + +lib.define( + "upsample_bilinear2d(Tensor input, int[]? output_size, bool align_corners, " + "float[]? scale_factors) -> Tensor" +) +lib.impl("upsample_bilinear2d", _upsample_bilinear2d, "CompositeExplicitAutograd") +register_fake("channels_last::upsample_bilinear2d", _upsample_bilinear2d, lib=lib) + +lib.define( + "upsample_nearest2d(Tensor input, int[]? output_size, float[]? scale_factors) " + "-> Tensor" +) +lib.impl("upsample_nearest2d", _upsample_nearest2d, "CompositeExplicitAutograd") +register_fake("channels_last::upsample_nearest2d", _upsample_nearest2d, lib=lib) + +lib.define( + "max_pool2d_with_indices(Tensor input, int[2] kernel_size, int[2] stride, " + "int[2] padding, int[2] dilation, bool ceil_mode) -> (Tensor, Tensor)" +) +lib.impl( + "max_pool2d_with_indices", _max_pool2d_with_indices, "CompositeExplicitAutograd" +) +register_fake( + "channels_last::max_pool2d_with_indices", _max_pool2d_with_indices, lib=lib +) + +lib.define( + "grid_sampler_2d(Tensor input, Tensor grid, int interpolation_mode, " + "int padding_mode, bool align_corners) -> Tensor" +) +lib.impl("grid_sampler_2d", _grid_sampler_2d, "CompositeExplicitAutograd") +register_fake("channels_last::grid_sampler_2d", _grid_sampler_2d, lib=lib) + lib.define("permute_copy(Tensor input, int[] dims) -> Tensor") lib.impl("permute_copy", _permute_copy, "CompositeExplicitAutograd") register_fake("channels_last::permute_copy", _permute_copy, lib=lib) diff --git a/backends/transforms/decompose_channels_last_pass.py b/backends/transforms/decompose_channels_last_pass.py index db2a173ab10..69214437f7a 100644 --- a/backends/transforms/decompose_channels_last_pass.py +++ b/backends/transforms/decompose_channels_last_pass.py @@ -5,6 +5,8 @@ # LICENSE file in the root directory of this source tree. # Importing registers the channels_last dialect (and its edge overloads). +import operator + import executorch.backends.transforms.channels_last_ops # noqa: F401 from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import ExportPass @@ -13,10 +15,16 @@ _NHWC_TO_NCHW = [0, 3, 1, 2] _NCHW_TO_NHWC = [0, 2, 3, 1] -# channels_last op -> the channels-first aten op it wraps. +# Single-output channels_last op -> the channels-first aten op it wraps. The +# activation (arg 0) is permuted to NCHW, the op runs, and the result is permuted +# back; any remaining args (e.g. grid_sampler's grid) pass through unchanged. _DECOMPOSITIONS = { exir_ops.edge.channels_last.convolution.default: exir_ops.edge.aten.convolution.default, exir_ops.edge.channels_last.avg_pool2d.default: exir_ops.edge.aten.avg_pool2d.default, + exir_ops.edge.channels_last.adaptive_avg_pool2d.default: exir_ops.edge.aten._adaptive_avg_pool2d.default, + exir_ops.edge.channels_last.upsample_bilinear2d.default: exir_ops.edge.aten.upsample_bilinear2d.vec, + exir_ops.edge.channels_last.upsample_nearest2d.default: exir_ops.edge.aten.upsample_nearest2d.vec, + exir_ops.edge.channels_last.grid_sampler_2d.default: exir_ops.edge.aten.grid_sampler_2d.default, } @@ -48,6 +56,35 @@ def call_operator(self, op, args, kwargs, meta): {}, meta, ) + if op == exir_ops.edge.channels_last.max_pool2d_with_indices.default: + nchw_in = super().call_operator( + exir_ops.edge.aten.permute_copy.default, + (args[0], _NHWC_TO_NCHW), + {}, + meta, + ) + pooled = super().call_operator( + exir_ops.edge.aten.max_pool2d_with_indices.default, + (nchw_in, *args[1:]), + kwargs, + meta, + ) + values = super().call_operator(operator.getitem, (pooled, 0), {}, meta) + indices = super().call_operator(operator.getitem, (pooled, 1), {}, meta) + return ( + super().call_operator( + exir_ops.edge.aten.permute_copy.default, + (values, _NCHW_TO_NHWC), + {}, + meta, + ), + super().call_operator( + exir_ops.edge.aten.permute_copy.default, + (indices, _NCHW_TO_NHWC), + {}, + meta, + ), + ) if op == exir_ops.edge.channels_last.permute_copy.default: return super().call_operator( exir_ops.edge.aten.permute_copy.default, args, kwargs, meta diff --git a/backends/transforms/test/test_channels_last_ops.py b/backends/transforms/test/test_channels_last_ops.py index b06a2773c59..ca745e9c9ef 100644 --- a/backends/transforms/test/test_channels_last_ops.py +++ b/backends/transforms/test/test_channels_last_ops.py @@ -125,6 +125,85 @@ def test_permute_copy_moves_data(dims): assert actual.is_contiguous() +def test_adaptive_avg_pool2d_matches_aten(): + torch.manual_seed(0) + nchw = torch.randn(2, 3, 8, 8) + nhwc = _to_nhwc(nchw) + + expected = _to_nhwc(torch.ops.aten.adaptive_avg_pool2d(nchw, [4, 4])) + actual = torch.ops.channels_last.adaptive_avg_pool2d(nhwc, [4, 4]) + + assert actual.shape == expected.shape + assert torch.allclose(actual, expected, atol=1e-5) + + +def test_upsample_bilinear2d_matches_aten(): + torch.manual_seed(0) + nchw = torch.randn(2, 3, 8, 8) + nhwc = _to_nhwc(nchw) + + expected = _to_nhwc( + torch.ops.aten.upsample_bilinear2d.vec(nchw, [16, 16], False, None) + ) + actual = torch.ops.channels_last.upsample_bilinear2d(nhwc, [16, 16], False, None) + + assert actual.shape == expected.shape + assert torch.allclose(actual, expected, atol=1e-5) + + +def test_upsample_nearest2d_matches_aten(): + torch.manual_seed(0) + nchw = torch.randn(2, 3, 8, 8) + nhwc = _to_nhwc(nchw) + + expected = _to_nhwc(torch.ops.aten.upsample_nearest2d.vec(nchw, [16, 16], None)) + actual = torch.ops.channels_last.upsample_nearest2d(nhwc, [16, 16], None) + + assert actual.shape == expected.shape + assert torch.allclose(actual, expected, atol=1e-5) + + +def test_max_pool2d_with_indices_matches_aten(): + torch.manual_seed(0) + nchw = torch.randn(2, 3, 8, 8) + nhwc = _to_nhwc(nchw) + + # Oracle: PyTorch's own channels-last max_pool (a separate code path from our + # permute wrapper), so this is an independent check rather than a round-trip. + ref_v, ref_i = torch.ops.aten.max_pool2d_with_indices( + nchw.to(memory_format=torch.channels_last), + [2, 2], + [2, 2], + [0, 0], + [1, 1], + False, + ) + expected_v = ref_v.permute(0, 2, 3, 1) + expected_i = ref_i.permute(0, 2, 3, 1) + + act_v, act_i = torch.ops.channels_last.max_pool2d_with_indices( + nhwc, [2, 2], [2, 2], [0, 0], [1, 1], False + ) + + assert act_v.shape == expected_v.shape + assert torch.allclose(act_v, expected_v, atol=1e-5) + # Indices are flat spatial positions (h*W+w), layout-agnostic in PyTorch. + assert torch.equal(act_i, expected_i) + + +def test_grid_sampler_2d_matches_aten(): + torch.manual_seed(0) + nchw = torch.randn(2, 3, 8, 8) + grid = torch.rand(2, 6, 6, 2) * 2 - 1 + nhwc = _to_nhwc(nchw) + + expected = _to_nhwc(torch.ops.aten.grid_sampler_2d(nchw, grid, 0, 0, False)) + actual = torch.ops.channels_last.grid_sampler_2d(nhwc, grid, 0, 0, False) + + assert actual.shape == expected.shape + assert torch.allclose(actual, expected, atol=1e-5) + + def test_convolution_lowers_to_edge_dialect(): class M(torch.nn.Module): def forward(self, x, w, b): diff --git a/backends/transforms/test/test_decompose_channels_last_pass.py b/backends/transforms/test/test_decompose_channels_last_pass.py index 34c6dd0e7b4..625ecc1b102 100644 --- a/backends/transforms/test/test_decompose_channels_last_pass.py +++ b/backends/transforms/test/test_decompose_channels_last_pass.py @@ -48,6 +48,39 @@ def forward(self, x): ) +class _AdaptiveAvgPoolModule(torch.nn.Module): + def forward(self, x): + return torch.ops.channels_last.adaptive_avg_pool2d(x, [4, 4]) + + +class _UpsampleBilinearModule(torch.nn.Module): + def forward(self, x): + return torch.ops.channels_last.upsample_bilinear2d(x, [16, 16], False, None) + + +class _UpsampleBilinearScaleModule(torch.nn.Module): + # output_size=None with scale_factors (the other upsample.vec branch). + def forward(self, x): + return torch.ops.channels_last.upsample_bilinear2d(x, None, False, [2.0, 2.0]) + + +class _UpsampleNearestModule(torch.nn.Module): + def forward(self, x): + return torch.ops.channels_last.upsample_nearest2d(x, [16, 16], None) + + +class _GridSamplerModule(torch.nn.Module): + def forward(self, x, grid): + return torch.ops.channels_last.grid_sampler_2d(x, grid, 0, 0, False) + + +class _MaxPoolModule(torch.nn.Module): + def forward(self, x): + return torch.ops.channels_last.max_pool2d_with_indices( + x, [2, 2], [2, 2], [0, 0], [1, 1], False + ) + + class _PermuteModule(torch.nn.Module): def forward(self, x): return torch.ops.channels_last.permute_copy(x, [0, 3, 1, 2]) @@ -75,6 +108,41 @@ def forward(self, x): exir_ops.edge.aten.avg_pool2d.default, 2, ), + ( + _AdaptiveAvgPoolModule(), + (torch.randn(2, 8, 8, 3),), + exir_ops.edge.channels_last.adaptive_avg_pool2d.default, + exir_ops.edge.aten._adaptive_avg_pool2d.default, + 2, + ), + ( + _UpsampleBilinearModule(), + (torch.randn(2, 8, 8, 3),), + exir_ops.edge.channels_last.upsample_bilinear2d.default, + exir_ops.edge.aten.upsample_bilinear2d.vec, + 2, + ), + ( + _UpsampleBilinearScaleModule(), + (torch.randn(2, 8, 8, 3),), + exir_ops.edge.channels_last.upsample_bilinear2d.default, + exir_ops.edge.aten.upsample_bilinear2d.vec, + 2, + ), + ( + _UpsampleNearestModule(), + (torch.randn(2, 8, 8, 3),), + exir_ops.edge.channels_last.upsample_nearest2d.default, + exir_ops.edge.aten.upsample_nearest2d.vec, + 2, + ), + ( + _GridSamplerModule(), + (torch.randn(2, 8, 8, 3), torch.rand(2, 6, 6, 2) * 2 - 1), + exir_ops.edge.channels_last.grid_sampler_2d.default, + exir_ops.edge.aten.grid_sampler_2d.default, + 2, + ), ( _PermuteModule(), (torch.randn(2, 8, 8, 3),), @@ -108,3 +176,35 @@ def test_decomposed_program_runs_and_matches_eager( actual = method.forward(list(args))[0] torch.testing.assert_close(actual, eager, atol=1e-4, rtol=1e-4) + + +# max_pool2d_with_indices is multi-output (values, indices), so it gets dedicated +# tests rather than the single-output _CASES harness. +def test_max_pool2d_with_indices_decomposes(): + args = (torch.randn(2, 8, 8, 3),) + ep = torch.export.export(_MaxPoolModule().eval(), args, strict=True) + gm = ( + to_edge(ep) + .transform([DecomposeChannelsLastPass()]) + .exported_program() + .graph_module + ) + + assert _count(gm, exir_ops.edge.channels_last.max_pool2d_with_indices.default) == 0 + assert _count(gm, exir_ops.edge.aten.max_pool2d_with_indices.default) == 1 + # permutes: input + values + indices. + assert _count(gm, exir_ops.edge.aten.permute_copy.default) == 3 + + +def test_max_pool2d_with_indices_decomposed_runs_and_matches_eager(): + module = _MaxPoolModule() + args = (torch.randn(2, 8, 8, 3),) + expected_values, expected_indices = module(*args) + + ep = torch.export.export(module.eval(), args, strict=True) + et = to_edge(ep).transform([DecomposeChannelsLastPass()]).to_executorch() + method = _load_for_executorch_from_buffer(et.buffer) + values, indices = method.forward(list(args)) + + torch.testing.assert_close(values, expected_values, atol=1e-4, rtol=1e-4) + torch.testing.assert_close(indices, expected_indices)