From a0c502440a377a9612edb8763f5701aa20c03f9e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mateusz=20S=C5=82uszniak?= Date: Sun, 13 Sep 2026 14:49:24 +0200 Subject: [PATCH] Vulkan: implement instance norm by rewriting it to group norm `aten._native_batch_norm_legit.no_stats` has no Vulkan implementation, so every `nn.InstanceNorm2d` is a graph break. Architectures that normalize in each block pay for that repeatedly: the fast neural style TransformerNet at 640x640 splits into 16 delegate subgraphs with 126 nodes left on the CPU, and copies its activations out to the CPU and back 15 times per inference. `F.instance_norm` reshapes its input to [1, N*C, H, W] before lowering to this overload, so with a batch dim of 1 it is the same reduction as group norm with one group per channel. ReplaceInstanceNormPass rewrites it that way and the existing group norm kernels cover it, so no new shader is needed. A batch dim above 1 is a different reduction and is left alone. add_native_group_norm_node() prepacks weight and bias, so the pass only fires when both trace back to a constant. Identity `aten.repeat` (every repeat factor 1, which the existing shape equality check implies) now folds away as well: `F.instance_norm` emits `weight.repeat(b)`, which would otherwise leave the affine params as runtime tensors that group norm cannot prepack. Measured on a Mali-G76 (Galaxy S10+), style transfer at 640x640, fp16, per-exec slope interleaved over 3 rounds on a warm device: 2265.1 ms -> 1469.4 ms median, a 1.54x speedup, with non-overlapping ranges. The partition goes from 16 subgraphs to 1 and from 126 CPU nodes to 0. Output accuracy against the fp32 reference improves slightly as well, cosine 0.999993 -> 0.999994 and max abs diff 0.041 -> 0.031, because the fp16 round trips at the delegate boundaries are gone. --- backends/vulkan/_passes/__init__.py | 4 ++ .../vulkan/_passes/remove_redundant_ops.py | 3 + .../vulkan/_passes/replace_instance_norm.py | 68 +++++++++++++++++++ backends/vulkan/op_registry.py | 20 ++++++ backends/vulkan/utils.py | 44 ++++++++++++ backends/vulkan/vulkan_preprocess.py | 4 ++ 6 files changed, 143 insertions(+) create mode 100644 backends/vulkan/_passes/replace_instance_norm.py 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(),