diff --git a/backends/vulkan/_passes/__init__.py b/backends/vulkan/_passes/__init__.py index 1afaf48dde7..4cf83023704 100644 --- a/backends/vulkan/_passes/__init__.py +++ b/backends/vulkan/_passes/__init__.py @@ -22,6 +22,9 @@ from executorch.backends.vulkan._passes.remove_redundant_ops import ( RemoveRedundantOpsTransform, ) +from executorch.backends.vulkan._passes.replace_instance_norm import ( + ReplaceInstanceNormPass, +) from executorch.backends.vulkan._passes.squeeze_unsqueeze_inputs import ( SqueezeUnsqueezeInputs, ) @@ -36,6 +39,7 @@ "remove_asserts", "RemoveAssertsTransform", "RemoveRedundantOpsTransform", + "ReplaceInstanceNormPass", "SqueezeUnsqueezeInputs", "TagMemoryMetaPass", ] diff --git a/backends/vulkan/_passes/remove_redundant_ops.py b/backends/vulkan/_passes/remove_redundant_ops.py index 545a25f69e9..50b922e3ac1 100644 --- a/backends/vulkan/_passes/remove_redundant_ops.py +++ b/backends/vulkan/_passes/remove_redundant_ops.py @@ -36,6 +36,9 @@ class RemoveRedundantOpsTransform(ExportPass): exir_ops.edge.aten.expand_copy.default, # copy.default(self, src): no-op when src dtype/shape matches self. exir_ops.edge.aten.copy.default, + # repeat.default: no-op when every repeat factor is 1, which the shape + # equality check below implies. + exir_ops.edge.aten.repeat.default, } # For these ops the meaningful input is args[1] (src), not args[0] (self). diff --git a/backends/vulkan/_passes/replace_instance_norm.py b/backends/vulkan/_passes/replace_instance_norm.py new file mode 100644 index 00000000000..ca5dbc7c5a3 --- /dev/null +++ b/backends/vulkan/_passes/replace_instance_norm.py @@ -0,0 +1,68 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +import executorch.backends.vulkan.utils as utils + +import torch +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass, PassResult +from executorch.exir.passes import dead_code_elimination_pass + + +class ReplaceInstanceNormPass(ExportPass): + """ + Replace ``aten._native_batch_norm_legit.no_stats`` with + ``aten.native_group_norm`` using one group per channel. + + Without this, every ``nn.InstanceNorm2d`` is a graph break. Architectures that + normalize in each block (the fast neural style transformer nets, for one) then + copy their activations out to the CPU and back once per block, which costs far + more than the normalization itself. + """ + + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: + modified = False + + for node in list(graph_module.graph.nodes): + if not utils.node_is_instance_norm(node): + continue + + input_node = node.args[0] + assert isinstance(input_node, torch.fx.Node) + input_val = input_node.meta["val"] + batches, channels, height, width = (int(d) for d in input_val.shape) + + with graph_module.graph.inserting_before(node): + group_norm_node = graph_module.graph.create_node( + "call_function", + exir_ops.edge.aten.native_group_norm.default, + args=( + input_node, + node.args[1], # weight + node.args[2], # bias + batches, + channels, + height * width, + channels, # one group per channel + node.args[5], # eps + ), + ) + + out_val, _, _ = node.meta["val"] + stats_val = input_val.new_empty((batches, channels)) + group_norm_node.meta = dict(node.meta) + group_norm_node.meta["val"] = (out_val, stats_val, stats_val) + + node.replace_all_uses_with(group_norm_node) + modified = True + + if modified: + graph_module.recompile() + dead_code_elimination_pass(graph_module) + + return PassResult(graph_module, modified) diff --git a/backends/vulkan/op_registry.py b/backends/vulkan/op_registry.py index 78270be4373..ca7990009d3 100644 --- a/backends/vulkan/op_registry.py +++ b/backends/vulkan/op_registry.py @@ -1706,6 +1706,26 @@ def register_native_batch_norm_legit_no_training(): ) +@update_features(exir_ops.edge.aten._native_batch_norm_legit.no_stats) +def register_native_batch_norm_legit_no_stats(): + """Instance norm, which ReplaceInstanceNormPass rewrites into group norm. + + ``F.instance_norm`` lowers to this overload. The pass only handles the cases + node_is_instance_norm() accepts, so gate partitioning on the same predicate. + """ + return OpFeatures( + inputs_storage=utils.CHANNELS_PACKED_TEXTURE, + inputs_dtypes=utils.FP_T, + outputs_storage=[ + utils.CHANNELS_PACKED_TEXTURE, + utils.CONTIGUOUS_BUFFER, + utils.CONTIGUOUS_BUFFER, + ], + supports_prepacking=True, + are_node_inputs_supported_fn=utils.node_is_instance_norm, + ) + + # ============================================================================= # GroupNorm.cpp # ============================================================================= diff --git a/backends/vulkan/utils.py b/backends/vulkan/utils.py index 84b901b6b6e..8aec699708e 100644 --- a/backends/vulkan/utils.py +++ b/backends/vulkan/utils.py @@ -17,6 +17,7 @@ from executorch.exir.backend.canonical_partitioners.config_partitioner import ( format_target_name, ) +from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.dialects.edge._ops import EdgeOpOverload from executorch.exir.tensor import TensorSpec from torch._export.utils import is_buffer, is_lifted_tensor_constant, is_param @@ -2014,3 +2015,46 @@ def align_width_and_update_state_dict( ) return aligned_tensor + + +def node_is_instance_norm(node: torch.fx.Node) -> bool: + """ + Whether a node is an ``F.instance_norm`` that group norm can express. + + ``F.instance_norm`` reshapes its input to ``[1, N * C, H, W]`` and lowers to + ``_native_batch_norm_legit.no_stats``, which normalizes using statistics taken + over the batch and spatial dims. When the batch dim is 1 that is exactly group + norm with one group per channel, so the existing group norm kernels cover it. + A batch dim above 1 is a different reduction and is left alone. + """ + if node.target != exir_ops.edge.aten._native_batch_norm_legit.no_stats: + return False + + input_node = node.args[0] + if not isinstance(input_node, torch.fx.Node): + return False + + val = input_node.meta.get("val") + if val is None or val.dim() != 4 or val.shape[0] != 1: + return False + + # Group norm always applies an affine transform, so both weight and bias must + # be present. add_native_group_norm_node() prepacks them, so both must also + # trace back to a constant rather than being computed at runtime. + for affine_arg in (node.args[1], node.args[2]): + if not isinstance(affine_arg, torch.fx.Node): + return False + placeholder, _ = trace_args_until_placeholder(affine_arg) + if placeholder is None: + return False + + # Only the normalized output may be consumed. Group norm returns mean and rstd + # shaped [N, group] where batch norm saves them shaped [C], so the saved + # statistics are not drop-in replacements. + for user in node.users: + if user.op != "call_function" or user.target != operator.getitem: + return False + if user.args[1] != 0: + return False + + return True diff --git a/backends/vulkan/vulkan_preprocess.py b/backends/vulkan/vulkan_preprocess.py index 34c68fe5828..12a77451a94 100644 --- a/backends/vulkan/vulkan_preprocess.py +++ b/backends/vulkan/vulkan_preprocess.py @@ -27,6 +27,9 @@ ) from executorch.backends.vulkan._passes.fuse_patterns import FusePatternsPass from executorch.backends.vulkan._passes.remove_asserts import RemoveAssertsTransform +from executorch.backends.vulkan._passes.replace_instance_norm import ( + ReplaceInstanceNormPass, +) from executorch.backends.vulkan.serialization.vulkan_graph_builder import VkGraphBuilder from executorch.backends.vulkan.serialization.vulkan_graph_schema import ( VkMemoryLayout, @@ -174,6 +177,7 @@ def preprocess( # noqa: C901 FusePatternsPass(), FuseClampPass(), RemoveRedundantOpsTransform(), + ReplaceInstanceNormPass(), FuseQuantizedOpsTransform(), FoldQDQPass(), SqueezeUnsqueezeInputs(),