Vulkan: implement instance norm by rewriting it to group norm

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",
 ]

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

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)

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
 # =============================================================================

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

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(),