From 88ff6418cd68c5f560949aea6602bf6e4cfc3795 Mon Sep 17 00:00:00 2001 From: Martin Pavella Date: Wed, 8 Jul 2026 15:00:06 +0200 Subject: [PATCH 1/3] Add default parameter values for `channels_last.avg_pool2d`. When replacing aten.avg_pool2d` with `channels_last.avg_pool2d`, some attributes of the aten operator can be left out, and must be replaced by default values. This can either be done in the "to channels last" transformation pass, or in the definition of the operator. Adding the default attributes of the aten operator to the channels_last operator seems cleaner to me. --- backends/transforms/channels_last_ops.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/backends/transforms/channels_last_ops.py b/backends/transforms/channels_last_ops.py index 54a897c7b80..04b0f245397 100644 --- a/backends/transforms/channels_last_ops.py +++ b/backends/transforms/channels_last_ops.py @@ -40,8 +40,17 @@ def _conv( def _avg_pool2d( - input, kernel_size, stride, padding, ceil_mode, count_include_pad, divisor_override + input, + kernel_size, + stride=None, + padding=0, + ceil_mode=False, + count_include_pad=True, + divisor_override=None, ): + if stride is None: + stride = [] # Default value of avg_pool2d. + nchw = input.permute(0, 3, 1, 2) out = torch.ops.aten.avg_pool2d( nchw, @@ -106,8 +115,8 @@ def _permute_copy(input, dims): register_fake("channels_last::convolution", _conv, lib=lib) lib.define( - "avg_pool2d(Tensor input, int[2] kernel_size, int[2] stride, int[2] padding, " - "bool ceil_mode, bool count_include_pad, int? divisor_override) -> Tensor" + "avg_pool2d(Tensor input, int[2] kernel_size, int[2] stride=[], int[2] padding=0, " + "bool ceil_mode=False, bool count_include_pad=True, int? divisor_override=None) -> Tensor" ) lib.impl("avg_pool2d", _avg_pool2d, "CompositeExplicitAutograd") register_fake("channels_last::avg_pool2d", _avg_pool2d, lib=lib) From 0be52e80f4a613566a557d348a4abb219c058fca Mon Sep 17 00:00:00 2001 From: Martin Pavella Date: Wed, 8 Jul 2026 15:09:00 +0200 Subject: [PATCH 2/3] Add pass to replace ops with their channels last variant. The pass currently supports `convolution` and `avg_pool2d`, as those are the only 2 operators with a channels last variant. The pass can be easily extended to more ops in the future. --- ...replace_ops_with_channels_last_variants.py | 132 ++++++++++ ...replace_ops_with_channels_last_variants.py | 240 ++++++++++++++++++ 2 files changed, 372 insertions(+) create mode 100644 backends/transforms/replace_ops_with_channels_last_variants.py create mode 100644 backends/transforms/test/test_replace_ops_with_channels_last_variants.py diff --git a/backends/transforms/replace_ops_with_channels_last_variants.py b/backends/transforms/replace_ops_with_channels_last_variants.py new file mode 100644 index 00000000000..cba8be197f5 --- /dev/null +++ b/backends/transforms/replace_ops_with_channels_last_variants.py @@ -0,0 +1,132 @@ +# Copyright 2026 NXP +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from dataclasses import dataclass +from typing import Callable + +import executorch.backends.transforms.channels_last_ops # noqa: F401 + +import torch +from executorch.backends.xnnpack._passes.xnnpack_pass import ExportPass +from executorch.exir import ExportedProgram +from executorch.exir.dialects._ops import ops as exir_ops +from torch.fx.node import Target +from torch.fx.passes.infra.pass_manager import PassResult + +_NCHW_TO_NHWC_PERM: list[int] = [0, 2, 3, 1] +_NHWC_TO_NCHW_PERM: list[int] = [0, 3, 1, 2] + + +def _requires_rank(rank: int) -> Callable[[torch.fx.Node], bool]: + def _inner(node: torch.fx.Node) -> bool: + return len(node.meta["val"].shape) == rank + + return _inner + + +@dataclass +class ChannelsLastOpSpec: + """Specification for replacing a contiguous op with its channels-last counterpart.""" + + # The channels_last dialect target to replace the contiguous op with. + target: Target + + # Positional arg indices of tensor inputs that should be permuted NCHW→NHWC. + input_indices: list[int] + + # If provided, this function must return True for a node to be replaced. + filter_fn: Callable[[torch.fx.Node], bool] | None = None + + # TODO Add output indices to support operators such as `max_pool2d_with_indices`. + + +_DEFAULT_OP_MAP: dict[Target, ChannelsLastOpSpec] = { + exir_ops.edge.aten.convolution.default: ChannelsLastOpSpec( + target=exir_ops.edge.channels_last.convolution.default, + input_indices=[0], + filter_fn=_requires_rank(4), + ), + exir_ops.edge.aten.avg_pool2d.default: ChannelsLastOpSpec( + target=exir_ops.edge.channels_last.avg_pool2d.default, + input_indices=[0], + filter_fn=_requires_rank(4), + ), +} + + +class ReplaceOpsWithChannelsLastVariants(ExportPass): + """ + Replaces contiguous (NCHW) edge-dialect ops with their channels-last (NHWC) + equivalents from the channels_last dialect, inserting permute_copy ops on + the specified inputs and on the output so that the rest of the graph remains + in contiguous layout. + + By default, all currently implemented channels_last dialect ops are replaced. + Pass a custom op_map to restrict or extend the set of replacements. + """ + + def __init__( + self, + exported_program: ExportedProgram, + op_map: dict[Target, ChannelsLastOpSpec] | None = None, + ) -> None: + super().__init__() + self.exported_program = exported_program + self.op_map: dict[Target, ChannelsLastOpSpec] = ( + op_map if op_map is not None else dict(_DEFAULT_OP_MAP) + ) + + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: + modified = False + graph = graph_module.graph + + for node in list(graph.nodes): + if node.op != "call_function": + continue + if (spec := self.op_map.get(node.target)) is None: + continue + contiguous_dim_order = tuple(range(len(node.meta["val"].shape))) + if node.meta["val"].dim_order() != contiguous_dim_order: + continue + if spec.filter_fn is not None and not spec.filter_fn(node): + continue + + args = list(node.args) + + with graph.inserting_before(node): + for idx in spec.input_indices: + perm_in = graph.create_node( + "call_function", + target=exir_ops.edge.channels_last.permute_copy.default, + args=(args[idx], _NCHW_TO_NHWC_PERM), + ) + perm_in.meta = {} + args[idx] = perm_in + + nhwc_node = graph.create_node( + "call_function", + target=spec.target, + args=tuple(args), + kwargs=node.kwargs, + ) + nhwc_node.meta = {} + + perm_out = graph.create_node( + "call_function", + target=exir_ops.edge.channels_last.permute_copy.default, + args=(nhwc_node, _NHWC_TO_NCHW_PERM), + ) + perm_out.meta = {} + + node.replace_all_uses_with(perm_out) + graph.erase_node(node) + modified = True + + if modified: + graph.eliminate_dead_code() + graph_module.recompile() + graph_module = super().call(graph_module).graph_module + + return PassResult(graph_module, modified) diff --git a/backends/transforms/test/test_replace_ops_with_channels_last_variants.py b/backends/transforms/test/test_replace_ops_with_channels_last_variants.py new file mode 100644 index 00000000000..150a50fa1ea --- /dev/null +++ b/backends/transforms/test/test_replace_ops_with_channels_last_variants.py @@ -0,0 +1,240 @@ +# Copyright 2026 NXP +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import executorch.backends.transforms.channels_last_ops # noqa: F401 +import torch +from executorch.backends.transforms.replace_ops_with_channels_last_variants import ( + _NCHW_TO_NHWC_PERM, + _NHWC_TO_NCHW_PERM, + ChannelsLastOpSpec, + ReplaceOpsWithChannelsLastVariants, +) +from executorch.exir import to_edge +from executorch.exir.dialects._ops import ops as exir_ops +from torch.export import ExportedProgram +from torch.fx import GraphModule +from torch.fx.node import Target + + +# ── helpers ──────────────────────────────────────────────────────────────────── + + +def _export_to_edge(module: torch.nn.Module, inputs: tuple) -> ExportedProgram: + ep = torch.export.export(module.eval(), inputs) + return to_edge(ep).exported_program() + + +def _run_pass(ep: ExportedProgram) -> tuple[GraphModule, bool]: + result = ReplaceOpsWithChannelsLastVariants(ep)(ep.graph_module) + return result.graph_module, result.modified + + +def _find_nodes(gm: GraphModule, target: Target) -> list[torch.fx.Node]: + return [n for n in gm.graph.nodes if n.op == "call_function" and n.target == target] + + +def _count(gm: GraphModule, target: Target) -> int: + return len(_find_nodes(gm, target)) + + +# ── modules ──────────────────────────────────────────────────────────────────── + + +class Conv2dModule(torch.nn.Module): + def __init__(self, groups: int = 1, bias: bool = True): + super().__init__() + self.conv = torch.nn.Conv2d( + 4, 4, kernel_size=3, padding=1, groups=groups, bias=bias + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.conv(x) + + +class AvgPool2dModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2) + + +class Conv2dAvgPool2dModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d(4, 4, kernel_size=3, padding=1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.avg_pool2d(self.conv(x), kernel_size=2, stride=2) + + +class Conv1dModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv1d(4, 4, kernel_size=3, padding=1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.conv(x) + + +class LinearModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.fc = torch.nn.Linear(8, 4) + + def forward(self, x): + return self.fc(x) + + +# ── tests ────────────────────────────────────────────────────────────────────── + + +class TestReplaceOpsWithChannelsLastVariants: + + def test_conv2d(self): + ep = _export_to_edge(Conv2dModule(bias=True), (torch.randn(1, 4, 8, 8),)) + assert _count(ep.graph_module, exir_ops.edge.aten.convolution.default) == 1 + + gm, modified = _run_pass(ep) + + assert modified + assert _count(gm, exir_ops.edge.aten.convolution.default) == 0 + assert _count(gm, exir_ops.edge.channels_last.convolution.default) == 1 + assert _count(gm, exir_ops.edge.channels_last.permute_copy.default) == 2 + + dialect_node = _find_nodes(gm, exir_ops.edge.channels_last.convolution.default)[ + 0 + ] + + # Input (arg 0) must be wrapped in an NCHW→NHWC permute. + input_permute = dialect_node.args[0] + assert input_permute.target == exir_ops.edge.channels_last.permute_copy.default + assert list(input_permute.args[1]) == _NCHW_TO_NHWC_PERM + + # Weight (arg 1) and bias (arg 2) must not be permuted. + for arg in (dialect_node.args[1], dialect_node.args[2]): + if isinstance(arg, torch.fx.Node): + assert arg.target != exir_ops.edge.channels_last.permute_copy.default + + # Output must be wrapped in an NHWC→NCHW permute. + users = list(dialect_node.users) + assert len(users) == 1 + output_permute = users[0] + assert output_permute.target == exir_ops.edge.channels_last.permute_copy.default + assert list(output_permute.args[1]) == _NHWC_TO_NCHW_PERM + + def test_depthwise_conv2d(self): + ep = _export_to_edge(Conv2dModule(groups=4), (torch.randn(1, 4, 8, 8),)) + gm, modified = _run_pass(ep) + + assert modified + assert _count(gm, exir_ops.edge.channels_last.convolution.default) == 1 + assert _count(gm, exir_ops.edge.channels_last.permute_copy.default) == 2 + + def test_avg_pool2d(self): + ep = _export_to_edge(AvgPool2dModule(), (torch.randn(1, 4, 8, 8),)) + assert _count(ep.graph_module, exir_ops.edge.aten.avg_pool2d.default) == 1 + + gm, modified = _run_pass(ep) + + assert modified + assert _count(gm, exir_ops.edge.aten.avg_pool2d.default) == 0 + assert _count(gm, exir_ops.edge.channels_last.avg_pool2d.default) == 1 + assert _count(gm, exir_ops.edge.channels_last.permute_copy.default) == 2 + + dialect_node = _find_nodes(gm, exir_ops.edge.channels_last.avg_pool2d.default)[ + 0 + ] + + input_permute = dialect_node.args[0] + assert input_permute.target == exir_ops.edge.channels_last.permute_copy.default + assert list(input_permute.args[1]) == _NCHW_TO_NHWC_PERM + + users = list(dialect_node.users) + assert len(users) == 1 + output_permute = users[0] + assert output_permute.target == exir_ops.edge.channels_last.permute_copy.default + assert list(output_permute.args[1]) == _NHWC_TO_NCHW_PERM + + def test_conv2d_and_avg_pool2d(self): + # 2 permutes per op. The redundant adjacent pair (NHWC→NCHW then NCHW→NHWC) between the two ops will be + # optimized out by a different pass. + ep = _export_to_edge(Conv2dAvgPool2dModule(), (torch.randn(1, 4, 8, 8),)) + gm, modified = _run_pass(ep) + + assert modified + assert _count(gm, exir_ops.edge.aten.convolution.default) == 0 + assert _count(gm, exir_ops.edge.aten.avg_pool2d.default) == 0 + assert _count(gm, exir_ops.edge.channels_last.convolution.default) == 1 + assert _count(gm, exir_ops.edge.channels_last.avg_pool2d.default) == 1 + assert _count(gm, exir_ops.edge.channels_last.permute_copy.default) == 4 + + def test_modified_false_when_no_matching_ops(self): + ep = _export_to_edge(LinearModule(), (torch.randn(2, 8),)) + _, modified = _run_pass(ep) + assert not modified + + def test_empty_op_map_leaves_graph_unchanged(self): + ep = _export_to_edge(Conv2dModule(), (torch.randn(1, 4, 8, 8),)) + + result = ReplaceOpsWithChannelsLastVariants(ep, op_map={})(ep.graph_module) + + assert not result.modified + assert _count(result.graph_module, exir_ops.edge.aten.convolution.default) == 1 + assert ( + _count(result.graph_module, exir_ops.edge.channels_last.convolution.default) + == 0 + ) + + def test_custom_op_map_only_replaces_specified_op(self): + ep = _export_to_edge(Conv2dAvgPool2dModule(), (torch.randn(1, 4, 8, 8),)) + + custom_map = { + exir_ops.edge.aten.avg_pool2d.default: ChannelsLastOpSpec( + target=exir_ops.edge.channels_last.avg_pool2d.default, + input_indices=[0], + ) + } + result = ReplaceOpsWithChannelsLastVariants(ep, op_map=custom_map)( + ep.graph_module + ) + gm = result.graph_module + + assert result.modified + assert _count(gm, exir_ops.edge.aten.convolution.default) == 1 + assert _count(gm, exir_ops.edge.channels_last.convolution.default) == 0 + assert _count(gm, exir_ops.edge.aten.avg_pool2d.default) == 0 + assert _count(gm, exir_ops.edge.channels_last.avg_pool2d.default) == 1 + + def test_conv1d_is_not_replaced(self): + ep = _export_to_edge(Conv1dModule(), (torch.randn(1, 4, 16),)) + gm, modified = _run_pass(ep) + + assert not modified + assert _count(gm, exir_ops.edge.aten.convolution.default) == 1 + assert _count(gm, exir_ops.edge.channels_last.convolution.default) == 0 + + def test_convolution_numerical_correctness(self): + torch.manual_seed(2) + x = torch.randn(1, 4, 8, 8) + model = Conv2dModule(bias=False) + reference_output = model(x) + + ep = _export_to_edge(model, (x,)) + _run_pass(ep) + assert _count(ep.module(), exir_ops.edge.channels_last.convolution.default) == 1 + + channels_last_output = ep.module()(x)[0] + assert torch.allclose(channels_last_output, reference_output) + + def test_avg_pool2d_numerical_correctness(self): + torch.manual_seed(2) + x = torch.randn(1, 4, 8, 8) + model = AvgPool2dModule() + reference_output = model(x) + + ep = _export_to_edge(model, (x,)) + _run_pass(ep) + assert _count(ep.module(), exir_ops.edge.channels_last.avg_pool2d.default) == 1 + + channels_last_output = ep.module()(x)[0] + assert torch.allclose(channels_last_output, reference_output) From 7e23e2dd5c3010e070ae16a087188df7fd4ddc37 Mon Sep 17 00:00:00 2001 From: Martin Pavella Date: Tue, 21 Jul 2026 09:32:35 +0200 Subject: [PATCH 3/3] Update based on comments and add support for new channels last operators. --- backends/transforms/channels_last_ops.py | 11 +- ...replace_ops_with_channels_last_variants.py | 189 +++++++++-- backends/transforms/targets.bzl | 30 ++ ...replace_ops_with_channels_last_variants.py | 299 +++++++++++++++++- 4 files changed, 483 insertions(+), 46 deletions(-) diff --git a/backends/transforms/channels_last_ops.py b/backends/transforms/channels_last_ops.py index 04b0f245397..8a2d7401bb3 100644 --- a/backends/transforms/channels_last_ops.py +++ b/backends/transforms/channels_last_ops.py @@ -84,7 +84,12 @@ def _upsample_nearest2d(input, 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): +def _max_pool2d_with_indices( + input, kernel_size, stride=None, padding=0, dilation=1, ceil_mode=False +): + if stride is None: + stride = [] # Default value of max_pool2d_with_indices. + nchw = input.permute(0, 3, 1, 2) values, indices = torch.ops.aten.max_pool2d_with_indices( nchw, kernel_size, stride, padding, dilation, ceil_mode @@ -140,8 +145,8 @@ def _permute_copy(input, dims): 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)" + "max_pool2d_with_indices(Tensor input, int[2] kernel_size, int[2] stride=[], " + "int[2] padding=0, int[2] dilation=1, bool ceil_mode=False) -> (Tensor, Tensor)" ) lib.impl( "max_pool2d_with_indices", _max_pool2d_with_indices, "CompositeExplicitAutograd" diff --git a/backends/transforms/replace_ops_with_channels_last_variants.py b/backends/transforms/replace_ops_with_channels_last_variants.py index cba8be197f5..63a1ca5a014 100644 --- a/backends/transforms/replace_ops_with_channels_last_variants.py +++ b/backends/transforms/replace_ops_with_channels_last_variants.py @@ -3,15 +3,18 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import operator from dataclasses import dataclass from typing import Callable import executorch.backends.transforms.channels_last_ops # noqa: F401 import torch -from executorch.backends.xnnpack._passes.xnnpack_pass import ExportPass + from executorch.exir import ExportedProgram from executorch.exir.dialects._ops import ops as exir_ops + +from executorch.exir.pass_base import ExportPass from torch.fx.node import Target from torch.fx.passes.infra.pass_manager import PassResult @@ -19,9 +22,33 @@ _NHWC_TO_NCHW_PERM: list[int] = [0, 3, 1, 2] -def _requires_rank(rank: int) -> Callable[[torch.fx.Node], bool]: - def _inner(node: torch.fx.Node) -> bool: - return len(node.meta["val"].shape) == rank +FilterFnType = Callable[[torch.fx.Node], bool] + + +def _requires_rank(allowed_ranks: list[int]) -> FilterFnType: + def _inner(node) -> bool: + val = node.meta["val"] + val = val[0] if isinstance(val, (list, tuple)) else val + return val.dim() in allowed_ranks + + return _inner + + +def _has_2d_convolution_kernel() -> FilterFnType: + def _inner(node) -> bool: + # Extract kernel size from the weights + try: + w = node.args[1] + return w.meta["val"].dim() == 4 + except AttributeError: + return False + + return _inner + + +def _and(*filter_fns: FilterFnType) -> FilterFnType: + def _inner(*args) -> bool: + return all(fn(*args) for fn in filter_fns) return _inner @@ -36,22 +63,55 @@ class ChannelsLastOpSpec: # Positional arg indices of tensor inputs that should be permuted NCHW→NHWC. input_indices: list[int] - # If provided, this function must return True for a node to be replaced. - filter_fn: Callable[[torch.fx.Node], bool] | None = None + # Indices of the outputs that should be permuted NCHW→NHWC. + output_indices: list[int] - # TODO Add output indices to support operators such as `max_pool2d_with_indices`. + # If provided, this function must return True for a node to be replaced. + filter_fn: FilterFnType | None = None _DEFAULT_OP_MAP: dict[Target, ChannelsLastOpSpec] = { - exir_ops.edge.aten.convolution.default: ChannelsLastOpSpec( - target=exir_ops.edge.channels_last.convolution.default, + exir_ops.edge.aten._adaptive_avg_pool2d.default: ChannelsLastOpSpec( + target=exir_ops.edge.channels_last.adaptive_avg_pool2d.default, input_indices=[0], - filter_fn=_requires_rank(4), + output_indices=[0], + filter_fn=_requires_rank([3, 4]), ), exir_ops.edge.aten.avg_pool2d.default: ChannelsLastOpSpec( target=exir_ops.edge.channels_last.avg_pool2d.default, input_indices=[0], - filter_fn=_requires_rank(4), + output_indices=[0], + filter_fn=_requires_rank([3, 4]), + ), + exir_ops.edge.aten.convolution.default: ChannelsLastOpSpec( + target=exir_ops.edge.channels_last.convolution.default, + input_indices=[0], + output_indices=[0], + filter_fn=_and(_requires_rank([3, 4]), _has_2d_convolution_kernel()), + ), + exir_ops.edge.aten.grid_sampler_2d.default: ChannelsLastOpSpec( + target=exir_ops.edge.channels_last.grid_sampler_2d.default, + input_indices=[0], + output_indices=[0], + filter_fn=_requires_rank([4]), + ), + exir_ops.edge.aten.max_pool2d_with_indices.default: ChannelsLastOpSpec( + target=exir_ops.edge.channels_last.max_pool2d_with_indices.default, + input_indices=[0], + output_indices=[0, 1], + filter_fn=_requires_rank([3, 4]), + ), + exir_ops.edge.aten.upsample_bilinear2d.vec: ChannelsLastOpSpec( + target=exir_ops.edge.channels_last.upsample_bilinear2d.default, + input_indices=[0], + output_indices=[0], + filter_fn=_requires_rank([4]), + ), + exir_ops.edge.aten.upsample_nearest2d.vec: ChannelsLastOpSpec( + target=exir_ops.edge.channels_last.upsample_nearest2d.default, + input_indices=[0], + output_indices=[0], + filter_fn=_requires_rank([4]), ), } @@ -78,6 +138,53 @@ def __init__( op_map if op_map is not None else dict(_DEFAULT_OP_MAP) ) + @staticmethod + def _permute_node_input( + graph: torch.fx.Graph, node_input: torch.fx.Node, implicit_batch: bool + ) -> torch.fx.Node: + if implicit_batch: + # Explicitly add batch size of `1`. + node_input = graph.create_node( + "call_function", + target=exir_ops.edge.aten.unsqueeze_copy.default, + args=(node_input, 0), + ) + node_input.meta = {} + + res = graph.create_node( + "call_function", + target=exir_ops.edge.channels_last.permute_copy.default, + args=(node_input, _NCHW_TO_NHWC_PERM), + ) + res.meta = {} + + return res + + @staticmethod + def _permute_node_output( + graph: torch.fx.Graph, + node_output: torch.fx.Node, + original_node_output: torch.fx.Node, + implicit_batch: bool, + ): + output = graph.create_node( + "call_function", + target=exir_ops.edge.channels_last.permute_copy.default, + args=(node_output, _NHWC_TO_NCHW_PERM), + ) + output.meta = {} + + if implicit_batch: + # Remove the explicitly added batch size of `1`. + output = graph.create_node( + "call_function", + target=exir_ops.edge.aten.squeeze_copy.dims, + args=(output, [0]), + ) + output.meta = {} + + original_node_output.replace_all_uses_with(output) + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: modified = False graph = graph_module.graph @@ -87,23 +194,26 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: continue if (spec := self.op_map.get(node.target)) is None: continue - contiguous_dim_order = tuple(range(len(node.meta["val"].shape))) - if node.meta["val"].dim_order() != contiguous_dim_order: + val = node.meta["val"] + val = val[0] if isinstance(val, (list, tuple)) else val + contiguous_dim_order = tuple(range(val.dim())) + if val.dim_order() != contiguous_dim_order: continue if spec.filter_fn is not None and not spec.filter_fn(node): continue + # In case of implicit batch size, insert also `unsqueeze_copy.default` and `squeeze_copy.dims` operators. + # With `convolution`, this already happens during lowering to edge. But it doesn't happen for example with + # `avg_pool2d`. Therefore, we do it manually here to achieve consistency between the operators. + implicit_batch = val.dim() == 3 + args = list(node.args) with graph.inserting_before(node): for idx in spec.input_indices: - perm_in = graph.create_node( - "call_function", - target=exir_ops.edge.channels_last.permute_copy.default, - args=(args[idx], _NCHW_TO_NHWC_PERM), + args[idx] = self._permute_node_input( + graph, args[idx], implicit_batch ) - perm_in.meta = {} - args[idx] = perm_in nhwc_node = graph.create_node( "call_function", @@ -113,14 +223,41 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: ) nhwc_node.meta = {} - perm_out = graph.create_node( - "call_function", - target=exir_ops.edge.channels_last.permute_copy.default, - args=(nhwc_node, _NHWC_TO_NCHW_PERM), - ) - perm_out.meta = {} + users = list(node.users) + if all( + u.op == "call_function" and u.target == operator.getitem + for u in users + ): + # `node` produces multiple outputs which are extracted by following `getitem` nodes. + for idx in spec.output_indices: + old_getitem_nodes = [u for u in users if u.args[1] == idx] + if not old_getitem_nodes: + continue # The output is not used. + + # Add a new `getitem` node. + new_getitem_node = graph.create_node( + "call_function", + target=operator.getitem, + args=(nhwc_node, idx), + ) + new_getitem_node.meta = {} + + self._permute_node_output( + graph, + new_getitem_node, + old_getitem_nodes[0], + implicit_batch, + ) + + # Remove the old `getitem` node. + graph.erase_node(old_getitem_nodes[0]) + else: + # Regular node with a single output. + assert spec.output_indices == [ + 0 + ], f"Incorrect use of `output_indices` for op `{spec.target}`." + self._permute_node_output(graph, nhwc_node, node, implicit_batch) - node.replace_all_uses_with(perm_out) graph.erase_node(node) modified = True diff --git a/backends/transforms/targets.bzl b/backends/transforms/targets.bzl index 253a7b57bc9..61297e11cb4 100644 --- a/backends/transforms/targets.bzl +++ b/backends/transforms/targets.bzl @@ -463,3 +463,33 @@ def define_common_targets(): ":replace_nop_transpose_or_permute_with_view", ], ) + + runtime.python_library( + name = "replace_ops_with_channels_last_variants", + srcs = [ + "replace_ops_with_channels_last_variants.py", + ], + visibility = [ + "//executorch/backends/...", + ], + deps = [ + "//caffe2:torch", + ":channels_last_ops", + "//executorch/exir:pass_base", + "//executorch/exir:lib", + ], + ) + + runtime.python_test( + name = "test_replace_ops_with_channels_last_variants", + srcs = [ + "test/test_replace_ops_with_channels_last_variants.py", + ], + deps = [ + "//caffe2:torch", + ":channels_last_ops", + ":replace_ops_with_channels_last_variants", + "//executorch/exir:lib", + "fbsource//third-party/pypi/pytest:pytest", + ], + ) diff --git a/backends/transforms/test/test_replace_ops_with_channels_last_variants.py b/backends/transforms/test/test_replace_ops_with_channels_last_variants.py index 150a50fa1ea..6f1807205a4 100644 --- a/backends/transforms/test/test_replace_ops_with_channels_last_variants.py +++ b/backends/transforms/test/test_replace_ops_with_channels_last_variants.py @@ -3,7 +3,10 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import operator + import executorch.backends.transforms.channels_last_ops # noqa: F401 +import pytest import torch from executorch.backends.transforms.replace_ops_with_channels_last_variants import ( _NCHW_TO_NHWC_PERM, @@ -58,6 +61,11 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2) +class AdaptiveAvgPool2dModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.adaptive_avg_pool2d(x, (2, 2)) + + class Conv2dAvgPool2dModule(torch.nn.Module): def __init__(self): super().__init__() @@ -67,6 +75,11 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return torch.nn.functional.avg_pool2d(self.conv(x), kernel_size=2, stride=2) +class GridSampler2DModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.grid_sample(x, torch.ones(1, 2, 2, 2)) + + class Conv1dModule(torch.nn.Module): def __init__(self): super().__init__() @@ -85,6 +98,35 @@ def forward(self, x): return self.fc(x) +class MaxPool2DModule(torch.nn.Module): + def __init__(self, kernel_size: int | tuple[int, ...] = 3, **kwargs): + super().__init__() + self.max_pool2d = torch.nn.MaxPool2d(kernel_size, **kwargs) + + def forward(self, x): + return self.max_pool2d(x) + + +class UpsampleBilinearModule(torch.nn.Module): + + def __init__(self): + super().__init__() + self.upsample = torch.nn.Upsample(scale_factor=2, mode="bilinear") + + def forward(self, x): + return self.upsample(x) + + +class UpsampleNearestModule(torch.nn.Module): + + def __init__(self): + super().__init__() + self.upsample = torch.nn.Upsample(scale_factor=2, mode="nearest") + + def forward(self, x): + return self.upsample(x) + + # ── tests ────────────────────────────────────────────────────────────────────── @@ -122,6 +164,47 @@ def test_conv2d(self): assert output_permute.target == exir_ops.edge.channels_last.permute_copy.default assert list(output_permute.args[1]) == _NHWC_TO_NCHW_PERM + def test_conv2d__implicit_batch(self): + input_shape = (4, 8, 8) # Use implicit batch size of `1`. + ep = _export_to_edge(Conv2dModule(bias=True), (torch.randn(*input_shape),)) + assert _count(ep.graph_module, exir_ops.edge.aten.convolution.default) == 1 + + gm, modified = _run_pass(ep) + + assert modified + assert _count(gm, exir_ops.edge.aten.convolution.default) == 0 + assert _count(gm, exir_ops.edge.channels_last.convolution.default) == 1 + assert _count(gm, exir_ops.edge.channels_last.permute_copy.default) == 2 + + # These ops are added automatically during lowering to edge. + assert _count(gm, exir_ops.edge.aten.unsqueeze_copy.default) == 1 + assert _count(gm, exir_ops.edge.aten.squeeze_copy.dims) == 1 + + dialect_node = _find_nodes(gm, exir_ops.edge.channels_last.convolution.default)[ + 0 + ] + + # Input (arg 0) must be wrapped in an NCHW→NHWC permute. + input_permute = dialect_node.args[0] + assert input_permute.target == exir_ops.edge.channels_last.permute_copy.default + assert list(input_permute.args[1]) == _NCHW_TO_NHWC_PERM + assert input_permute.args[0].target == exir_ops.edge.aten.unsqueeze_copy.default + + # Weight (arg 1) and bias (arg 2) must not be permuted. + for arg in (dialect_node.args[1], dialect_node.args[2]): + if isinstance(arg, torch.fx.Node): + assert arg.target != exir_ops.edge.channels_last.permute_copy.default + + # Output must be wrapped in an NHWC→NCHW permute. + users = list(dialect_node.users) + assert len(users) == 1 + output_permute = users[0] + assert output_permute.target == exir_ops.edge.channels_last.permute_copy.default + assert list(output_permute.args[1]) == _NHWC_TO_NCHW_PERM + assert ( + list(output_permute.users)[0].target == exir_ops.edge.aten.squeeze_copy.dims + ) + def test_depthwise_conv2d(self): ep = _export_to_edge(Conv2dModule(groups=4), (torch.randn(1, 4, 8, 8),)) gm, modified = _run_pass(ep) @@ -130,20 +213,135 @@ def test_depthwise_conv2d(self): assert _count(gm, exir_ops.edge.channels_last.convolution.default) == 1 assert _count(gm, exir_ops.edge.channels_last.permute_copy.default) == 2 - def test_avg_pool2d(self): - ep = _export_to_edge(AvgPool2dModule(), (torch.randn(1, 4, 8, 8),)) - assert _count(ep.graph_module, exir_ops.edge.aten.avg_pool2d.default) == 1 + def test_max_pool2d(self): + ep = _export_to_edge(MaxPool2DModule(), (torch.randn(1, 4, 8, 8),)) + assert ( + _count(ep.graph_module, exir_ops.edge.aten.max_pool2d_with_indices.default) + == 1 + ) gm, modified = _run_pass(ep) assert modified - assert _count(gm, exir_ops.edge.aten.avg_pool2d.default) == 0 - assert _count(gm, exir_ops.edge.channels_last.avg_pool2d.default) == 1 + assert _count(gm, exir_ops.edge.aten.max_pool2d_with_indices.default) == 0 + assert ( + _count(gm, exir_ops.edge.channels_last.max_pool2d_with_indices.default) == 1 + ) assert _count(gm, exir_ops.edge.channels_last.permute_copy.default) == 2 + assert _count(gm, operator.getitem) == 1 - dialect_node = _find_nodes(gm, exir_ops.edge.channels_last.avg_pool2d.default)[ - 0 - ] + dialect_node = _find_nodes( + gm, exir_ops.edge.channels_last.max_pool2d_with_indices.default + )[0] + + # Input (arg 0) must be wrapped in an NCHW→NHWC permute. + input_permute = dialect_node.args[0] + assert input_permute.target == exir_ops.edge.channels_last.permute_copy.default + assert list(input_permute.args[1]) == _NCHW_TO_NHWC_PERM + + users = list(dialect_node.users) + assert len(users) == 1 + output_getitem = users[0] + assert output_getitem.target == operator.getitem + + # Output must be wrapped in an NHWC→NCHW permute. + getitem_users = list(output_getitem.users) + assert len(getitem_users) == 1 + output_permute = getitem_users[0] + assert output_permute.target == exir_ops.edge.channels_last.permute_copy.default + assert list(output_permute.args[1]) == _NHWC_TO_NCHW_PERM + + def test_max_pool2d__implicit_batch(self): + input_shape = (4, 8, 8) # Use implicit batch size of `1`. + ep = _export_to_edge(MaxPool2DModule(), (torch.randn(*input_shape),)) + assert ( + _count(ep.graph_module, exir_ops.edge.aten.max_pool2d_with_indices.default) + == 1 + ) + + gm, modified = _run_pass(ep) + + assert modified + assert _count(gm, exir_ops.edge.aten.max_pool2d_with_indices.default) == 0 + assert ( + _count(gm, exir_ops.edge.channels_last.max_pool2d_with_indices.default) == 1 + ) + assert _count(gm, exir_ops.edge.channels_last.permute_copy.default) == 2 + assert _count(gm, exir_ops.edge.aten.unsqueeze_copy.default) == 1 + assert _count(gm, exir_ops.edge.aten.squeeze_copy.dims) == 1 + + dialect_node = _find_nodes( + gm, exir_ops.edge.channels_last.max_pool2d_with_indices.default + )[0] + + # Input (arg 0) must be wrapped in an NCHW→NHWC permute + an unsqueeze. + input_permute = dialect_node.args[0] + assert input_permute.target == exir_ops.edge.channels_last.permute_copy.default + assert list(input_permute.args[1]) == _NCHW_TO_NHWC_PERM + assert input_permute.args[0].target == exir_ops.edge.aten.unsqueeze_copy.default + + users = list(dialect_node.users) + assert len(users) == 1 + output_getitem = users[0] + assert output_getitem.target == operator.getitem + + # Output must be wrapped in an NHWC→NCHW permute + a squeeze. + getitem_users = list(output_getitem.users) + assert len(getitem_users) == 1 + output_permute = getitem_users[0] + assert output_permute.target == exir_ops.edge.channels_last.permute_copy.default + assert list(output_permute.args[1]) == _NHWC_TO_NCHW_PERM + assert ( + list(output_permute.users)[0].target == exir_ops.edge.aten.squeeze_copy.dims + ) + + @pytest.mark.parametrize( + "module_cls, aten_op, channels_last_op", + [ + pytest.param( + AvgPool2dModule, + exir_ops.edge.aten.avg_pool2d.default, + exir_ops.edge.channels_last.avg_pool2d.default, + id="AvgPool2D", + ), + pytest.param( + AdaptiveAvgPool2dModule, + exir_ops.edge.aten._adaptive_avg_pool2d.default, + exir_ops.edge.channels_last.adaptive_avg_pool2d.default, + id="AdaptiveAvgPool2D", + ), + pytest.param( + GridSampler2DModule, + exir_ops.edge.aten.grid_sampler_2d.default, + exir_ops.edge.channels_last.grid_sampler_2d.default, + id="GridSampler2D", + ), + pytest.param( + UpsampleBilinearModule, + exir_ops.edge.aten.upsample_bilinear2d.vec, + exir_ops.edge.channels_last.upsample_bilinear2d.default, + id="UpsampleBilinear2D", + ), + pytest.param( + UpsampleNearestModule, + exir_ops.edge.aten.upsample_nearest2d.vec, + exir_ops.edge.channels_last.upsample_nearest2d.default, + id="UpsampleNearest2D", + ), + ], + ) + def test_ops_without_weights(self, module_cls, aten_op, channels_last_op): + ep = _export_to_edge(module_cls(), (torch.randn(1, 4, 8, 8),)) + assert _count(ep.graph_module, aten_op) == 1 + + gm, modified = _run_pass(ep) + + assert modified + assert _count(gm, aten_op) == 0 + assert _count(gm, channels_last_op) == 1 + assert _count(gm, exir_ops.edge.channels_last.permute_copy.default) == 2 + + dialect_node = _find_nodes(gm, channels_last_op)[0] input_permute = dialect_node.args[0] assert input_permute.target == exir_ops.edge.channels_last.permute_copy.default @@ -155,6 +353,56 @@ def test_avg_pool2d(self): assert output_permute.target == exir_ops.edge.channels_last.permute_copy.default assert list(output_permute.args[1]) == _NHWC_TO_NCHW_PERM + @pytest.mark.parametrize( + "module_cls, aten_op, channels_last_op", + [ + # GridSampler2D, UpsampleBilinear2D and UpsampleNearest2D don't support implicit batch. + pytest.param( + AvgPool2dModule, + exir_ops.edge.aten.avg_pool2d.default, + exir_ops.edge.channels_last.avg_pool2d.default, + id="AvgPool2D", + ), + pytest.param( + AdaptiveAvgPool2dModule, + exir_ops.edge.aten._adaptive_avg_pool2d.default, + exir_ops.edge.channels_last.adaptive_avg_pool2d.default, + id="AdaptiveAvgPool2D", + ), + ], + ) + def test_ops_without_weights__implicit_batch( + self, module_cls, aten_op, channels_last_op + ): + input_shape = (4, 8, 8) # Use implicit batch size of `1`. + ep = _export_to_edge(module_cls(), (torch.randn(*input_shape),)) + assert _count(ep.graph_module, aten_op) == 1 + + gm, modified = _run_pass(ep) + + assert modified + assert _count(gm, aten_op) == 0 + assert _count(gm, channels_last_op) == 1 + assert _count(gm, exir_ops.edge.channels_last.permute_copy.default) == 2 + assert _count(gm, exir_ops.edge.aten.unsqueeze_copy.default) == 1 + assert _count(gm, exir_ops.edge.aten.squeeze_copy.dims) == 1 + + dialect_node = _find_nodes(gm, channels_last_op)[0] + + input_permute = dialect_node.args[0] + assert input_permute.target == exir_ops.edge.channels_last.permute_copy.default + assert list(input_permute.args[1]) == _NCHW_TO_NHWC_PERM + assert input_permute.args[0].target == exir_ops.edge.aten.unsqueeze_copy.default + + users = list(dialect_node.users) + assert len(users) == 1 + output_permute = users[0] + assert output_permute.target == exir_ops.edge.channels_last.permute_copy.default + assert list(output_permute.args[1]) == _NHWC_TO_NCHW_PERM + assert ( + list(output_permute.users)[0].target == exir_ops.edge.aten.squeeze_copy.dims + ) + def test_conv2d_and_avg_pool2d(self): # 2 permutes per op. The redundant adjacent pair (NHWC→NCHW then NCHW→NHWC) between the two ops will be # optimized out by a different pass. @@ -192,6 +440,7 @@ def test_custom_op_map_only_replaces_specified_op(self): exir_ops.edge.aten.avg_pool2d.default: ChannelsLastOpSpec( target=exir_ops.edge.channels_last.avg_pool2d.default, input_indices=[0], + output_indices=[0], ) } result = ReplaceOpsWithChannelsLastVariants(ep, op_map=custom_map)( @@ -205,7 +454,7 @@ def test_custom_op_map_only_replaces_specified_op(self): assert _count(gm, exir_ops.edge.aten.avg_pool2d.default) == 0 assert _count(gm, exir_ops.edge.channels_last.avg_pool2d.default) == 1 - def test_conv1d_is_not_replaced(self): + def test_conv1d__not_replaced(self): ep = _export_to_edge(Conv1dModule(), (torch.randn(1, 4, 16),)) gm, modified = _run_pass(ep) @@ -213,28 +462,44 @@ def test_conv1d_is_not_replaced(self): assert _count(gm, exir_ops.edge.aten.convolution.default) == 1 assert _count(gm, exir_ops.edge.channels_last.convolution.default) == 0 - def test_convolution_numerical_correctness(self): + def test_convolution__numerical_correctness(self): torch.manual_seed(2) x = torch.randn(1, 4, 8, 8) model = Conv2dModule(bias=False) reference_output = model(x) ep = _export_to_edge(model, (x,)) - _run_pass(ep) - assert _count(ep.module(), exir_ops.edge.channels_last.convolution.default) == 1 + gm, modified = _run_pass(ep) + assert modified + assert _count(gm, exir_ops.edge.channels_last.convolution.default) == 1 - channels_last_output = ep.module()(x)[0] + channels_last_output = gm(ep.state_dict["conv.weight"], x)[0] assert torch.allclose(channels_last_output, reference_output) - def test_avg_pool2d_numerical_correctness(self): + def test_avg_pool2d__numerical_correctness(self): torch.manual_seed(2) x = torch.randn(1, 4, 8, 8) model = AvgPool2dModule() reference_output = model(x) ep = _export_to_edge(model, (x,)) - _run_pass(ep) - assert _count(ep.module(), exir_ops.edge.channels_last.avg_pool2d.default) == 1 + gm, modified = _run_pass(ep) + assert modified + assert _count(gm, exir_ops.edge.channels_last.avg_pool2d.default) == 1 + + channels_last_output = gm(x)[0] + assert torch.allclose(channels_last_output, reference_output) + + def test_avg_pool2d__implicit_batch__numerical_correctness(self): + torch.manual_seed(2) + x = torch.randn(4, 8, 8) + model = AvgPool2dModule() + reference_output = model(x) + + ep = _export_to_edge(model, (x,)) + gm, modified = _run_pass(ep) + assert modified + assert _count(gm, exir_ops.edge.channels_last.avg_pool2d.default) == 1 - channels_last_output = ep.module()(x)[0] + channels_last_output = gm(x)[0] assert torch.allclose(channels_last_output, reference_output)