Channels last more ops

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)

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

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):

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)