[ET-VK][conv2d] Fuse residual adds and activations into conv2d during lowering

Pull Request resolved: #24808

## Ulterior Motive

Closing the DSR FP16 gap to LiteRT. DSR has 83 skip connections of the form `conv(...) + x * scale`, each lowered today to a separate `mul` and `add` dispatch. The previous commit added a conv2d epilogue that applies them inside every conv2d shader; this commit makes lowering use it.

It also makes `et_vk.conv2d.with_epilogue` the single fused op for conv2d, replacing `et_vk.conv_with_clamp` there.

## Rationale

**What**: A `conv2d_with_epilogue` pattern for `FusePatternsPass` that rewrites either of

```
activation(residual_op(convolution(...), residual [* scale]))
activation(convolution(...))
```

into `et_vk.conv2d.with_epilogue`, plus the op's definition and `op_registry` entry. A new `skip_fusion_patterns` compile option lets `WebGPUPartitioner` skip the pattern.

**Why**: Without the pattern, no exported model reaches the epilogue. On DSR, an earlier revision of this stack shrank the Vulkan delegate from 469 to 229 operators and cut end-to-end latency 9.6% on Galaxy S23 and Pixel 8 Pro. The WebGPU runtime consumes Vulkan lowering but does not implement `et_vk.conv2d.with_epilogue`.

## Details

- **Behavior**
  - conv2d + activation
    - **Before**: `FuseClampPass` rewrote it into `et_vk.conv_with_clamp`.
    - **After**: `et_vk.conv2d.with_epilogue`. At runtime both reach `add_conv2d` with an activation epilogue.
  - conv2d + residual, with or without an activation
    - **Before**: separate `aten.convolution`, `mul`, `add` and activation ops.
    - **After**: one `et_vk.conv2d.with_epilogue`.
  - conv1d + activation: unchanged. `Conv1dAsConv2dPass` and `FuseClampPass` still run after `FusePatternsPass`, so conv1d, including conv1d rewritten as conv2d, fuses into `et_vk.conv_with_clamp`, which stays registered for conv1d, WebGPU and existing `.pte` files.
  - WebGPU lowering: unchanged. `WebGPUPartitioner` skips `conv2d_with_epilogue`, so conv2d + activation still becomes `et_vk.conv_with_clamp`.
- **Match**
  - A 2D, batch 1, FP16/FP32 `aten.convolution` (direct, pointwise, depthwise, grouped or transposed) with one user, whose input and output fit channels-packed textures.
  - Feeding `aten.add.Tensor`, `aten.sub.Tensor` or `aten.mul.Tensor` with a residual of the conv output's shape and dtype. The conv may be either operand of `mul`, and either operand of `add` when `alpha` is 1; it must be the first operand of `sub`.
  - When the residual is a `mul` of `x` and a `[1, C, 1, 1]` or `[C, 1, 1]` placeholder `scale`, and the residual op is the `mul`'s only user, `x` becomes the residual and `scale` the residual scale. The runtime binds the residual as a tensor and prepacks the scale, so when `scale` is not a constant or `x` is one, the `mul` stays separate and its output is the residual.
  - The residual cannot be the conv itself (`y + y` stays unfused), and a constant residual without a scale stays unfused.
  - `relu`, `hardtanh`, `clamp` or `leaky_relu` is fused as the activation when it is the only user of the binary op, or of the conv when there is no residual. A missing `hardtanh` bound defaults to -1 or 1. Activations with non-constant parameters, such as a `clamp` bound by a symbolic size, stay unfused.
- **Partitioning**: As for the other patterns, `VulkanPartitioner` delegates a match's nodes without checking their support. The detector therefore repeats the conv's own checks (rank, batch, dtype, texture limits), so a conv that exceeds the texture limits stays on CPU unfused.
- **Pattern replacement**
  - `replace_all_fusable_subgraphs` removes dead code after the detectors replace a match at a node, so nodes another pattern already replaced (for example the conv and relu of a quantized conv) cannot match again.
  - A replacement function may return `False` when it leaves the graph unchanged, as this pattern does for a constant residual. Such matches are not counted, so `FusePatternsPass` does not retrace the graph for them.
- **Registry**: The residual is channels-packed texture3d like the input and output; weight, bias and scale are prepacked.
- **`skip_fusion_patterns`**: a collection of pattern names; a string or an unknown name raises `ValueError`. It is serialized as a compile spec, and both the partitioner's pattern detection and `FusePatternsPass` honor it:

```python
VulkanPartitioner(compile_options={"skip_fusion_patterns": {"conv2d_with_epilogue"}})
```

- **Compatibility**: `.pte` files exported after this commit contain `et_vk.conv2d.with_epilogue` for every FP conv2d followed by an activation or a residual op, so they need a runtime that includes the previous commit. Existing `.pte` files keep loading, since `et_vk.conv_with_clamp` stays registered.
- **WebGPU tests**: The op test generator and `test_conv_with_clamp.py` lowered through `VulkanPartitioner` directly; they now use `WebGPUPartitioner`.

```
schema: conv2d.with_epilogue(Tensor input, Tensor weight, Tensor? bias, SymInt[] stride,
  SymInt[] padding, SymInt[] dilation, bool transposed, SymInt[] output_padding,
  SymInt groups, Tensor? residual, Tensor? residual_scale, str residual_op,
  Scalar alpha, str activation, float activation_a, float activation_b) -> Tensor
```
ghstack-source-id: 469036207

Differential Revision: [D130216290](https://our.internmc.facebook.com/intern/diff/D130216290/)

Author: SS-JIA

backends/vulkan/_passes/fuse_patterns.py

@@ -4,7 +4,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from typing import Optional
+from typing import Collection, Optional
 
 import executorch.backends.vulkan.patterns as vk_patterns
 
@@ -15,15 +15,16 @@
 
 
 class FusePatternsPass(ExportPass):
-    def __init__(self) -> None:
+    def __init__(self, skip_patterns: Collection[str] = ()) -> None:
         super().__init__()
         self._exported_program: Optional[ExportedProgram] = None
+        self.skip_patterns = skip_patterns
 
     def call(self, graph_module: torch.fx.GraphModule):
         assert self._exported_program is not None
 
         total_replaced = vk_patterns.replace_all_fusable_subgraphs(
-            self._exported_program, graph_module
+            self._exported_program, graph_module, self.skip_patterns
         )
 
         if total_replaced > 0:

backends/vulkan/custom_ops_lib.py

@@ -110,6 +110,64 @@ def conv_with_clamp_out_impl(
 )
 lib.impl(name, conv_with_clamp_out_impl, "CompositeExplicitAutograd")
 
+##########################
+## conv2d.with_epilogue ##
+##########################
+
+
+def conv2d_with_epilogue_impl(
+    input,
+    weight,
+    bias,
+    stride,
+    padding,
+    dilation,
+    transposed,
+    output_padding,
+    groups,
+    residual=None,
+    residual_scale=None,
+    residual_op="add",
+    alpha=1,
+    activation="none",
+    activation_a=0.0,
+    activation_b=0.0,
+):
+    out = torch.convolution(
+        input,
+        weight,
+        bias,
+        stride,
+        padding,
+        dilation,
+        transposed,
+        output_padding,
+        groups,
+    )
+    if residual is not None:
+        if residual_scale is not None:
+            residual = residual * residual_scale
+        if residual_op == "add":
+            out = torch.add(out, residual, alpha=alpha)
+        elif residual_op == "sub":
+            out = torch.sub(out, residual, alpha=alpha)
+        else:
+            out = out * residual
+    if activation == "relu":
+        out = torch.relu(out)
+    elif activation == "clamp":
+        out = torch.clamp(out, activation_a, activation_b)
+    elif activation == "leaky_relu":
+        out = torch.nn.functional.leaky_relu(out, activation_a)
+    return out
+
+
+name = "conv2d.with_epilogue"
+lib.define(
+    f"{name}(Tensor input, Tensor weight, Tensor? bias, SymInt[] stride, SymInt[] padding, SymInt[] dilation, bool transposed, SymInt[] output_padding, SymInt groups, Tensor? residual, Tensor? residual_scale, str residual_op, Scalar alpha, str activation, float activation_a, float activation_b) -> Tensor"
+)
+lib.impl(name, conv2d_with_epilogue_impl, "CompositeExplicitAutograd")
+
 #################
 ## grid_priors ##
 #################

backends/vulkan/op_registry.py

@@ -1016,6 +1016,33 @@ def pick_conv_storage(
     )
 
 
+@update_features(exir_ops.edge.et_vk.conv2d.with_epilogue)
+def register_conv2d_with_epilogue():
+    return OpFeatures(
+        inputs_storage=[
+            utils.CHANNELS_PACKED_TEXTURE,  # input
+            utils.NO_STORAGE,  # weight (prepacked)
+            utils.NO_STORAGE,  # bias (prepacked)
+            utils.NO_STORAGE,  # stride (non tensor)
+            utils.NO_STORAGE,  # padding (non tensor)
+            utils.NO_STORAGE,  # dilation (non tensor)
+            utils.NO_STORAGE,  # transposed (non tensor)
+            utils.NO_STORAGE,  # output_padding (non tensor)
+            utils.NO_STORAGE,  # groups (non tensor)
+            utils.CHANNELS_PACKED_TEXTURE,  # residual
+            utils.NO_STORAGE,  # residual_scale (prepacked)
+            utils.NO_STORAGE,  # residual_op (non tensor)
+            utils.NO_STORAGE,  # alpha (non tensor)
+            utils.NO_STORAGE,  # activation (non tensor)
+            utils.NO_STORAGE,  # activation_a (non tensor)
+            utils.NO_STORAGE,  # activation_b (non tensor)
+        ],
+        inputs_dtypes=utils.FP_T,
+        supports_resize=True,
+        supports_prepacking=True,
+    )
+
+
 # =============================================================================
 # convolution/conv2d/q8ta/**/Q8taConv2d*.cpp
 # =============================================================================

backends/vulkan/partitioner/vulkan_partitioner.py

@@ -397,6 +397,18 @@ def parse_compile_options(compile_options: Dict[str, Any]) -> List[CompileSpec]:
             compile_specs.append(CompileSpec(key, value.to_bytes(8, "little")))
             continue
 
+        if key == "skip_fusion_patterns":
+            if (
+                isinstance(value, str)
+                or not set(value) <= vk_patterns.fusable_patterns.keys()
+            ):
+                raise ValueError(
+                    "skip_fusion_patterns must be a collection of pattern names "
+                    f"from {sorted(vk_patterns.fusable_patterns)}, got {value!r}"
+                )
+            compile_specs.append(CompileSpec(key, ",".join(sorted(value)).encode()))
+            continue
+
         if isinstance(value, (VkStorageType, VkMemoryLayout)):
             value_bytes = int(value).to_bytes(4, byteorder="little")
             compile_specs.append(CompileSpec(key, value_bytes))
@@ -516,7 +528,8 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult:
 
         # Get all fusable subgraphs from fuse_patterns
         fusable_subgraphs = vk_patterns.get_all_fusable_subgraphs(
-            exported_program.graph_module
+            exported_program.graph_module,
+            self.options.get("skip_fusion_patterns", ()),
         )
 
         # small_texture_limits opts into the conservative 3D texture limit that is

backends/vulkan/patterns/BUCK

@@ -9,6 +9,7 @@ fbcode_target(_kind = runtime.python_library,
     srcs = [
         "__init__.py",
         "pattern_registry.py",
+        "conv2d_with_epilogue.py",
         "rope.py",
         "rope_hf.py",
         "quantized_embedding.py",

backends/vulkan/patterns/__init__.py

@@ -4,7 +4,9 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from typing import List
+from typing import Collection, List
+
+import executorch.backends.vulkan.patterns.conv2d_with_epilogue  # noqa
 
 import executorch.backends.vulkan.patterns.quantized_binary  # noqa
 
@@ -38,6 +40,7 @@
     DetectorFn,
     fusable_patterns,
     GetGraphFn,
+    PatternEntry,
     PatternMatch,
     register_pattern_detector,
     register_pattern_graph,
@@ -66,9 +69,17 @@
 ]
 
 
-def all_fusable_graph_patterns() -> List[torch.fx.GraphModule]:
+def enabled_patterns(skip_patterns: Collection[str] = ()) -> List[PatternEntry]:
+    return [
+        entry for name, entry in fusable_patterns.items() if name not in skip_patterns
+    ]
+
+
+def all_fusable_graph_patterns(
+    skip_patterns: Collection[str] = (),
+) -> List[torch.fx.GraphModule]:
     all_patterns = []
-    for entry in fusable_patterns.values():
+    for entry in enabled_patterns(skip_patterns):
         if entry.get_graphs_fn is not None:
             all_patterns.extend(entry.get_graphs_fn())
 
@@ -77,18 +88,20 @@ def all_fusable_graph_patterns() -> List[torch.fx.GraphModule]:
 
 def get_all_fusable_subgraphs(
     graph_module: torch.fx.GraphModule,
+    skip_patterns: Collection[str] = (),
 ) -> List[PatternMatch]:
     fusable_subgraphs = []
 
-    fuse_patterns = all_fusable_graph_patterns()
+    fuse_patterns = all_fusable_graph_patterns(skip_patterns)
     for pattern in fuse_patterns:
         sm = SubgraphMatcher(pattern.graph, ignore_literals=True)
         matches = list(sm.match(graph_module.graph))
         for match in matches:
             fusable_subgraphs.append(create_pattern_match_from_internal_match(match))
 
+    entries = enabled_patterns(skip_patterns)
     for node in graph_module.graph.nodes:
-        for entry in fusable_patterns.values():
+        for entry in entries:
             if entry.detector_fn is not None:
                 maybe_match = entry.detector_fn(node)
                 if maybe_match is not None:
@@ -111,7 +124,8 @@ def create_replacement_for_pattern(
 
         for partition_to_replace in matches:
             pattern = create_pattern_match_from_internal_match(partition_to_replace)
-            create_replacement_func(ep, graph_module, pattern)
+            if create_replacement_func(ep, graph_module, pattern) is False:
+                continue
             total_replaced += 1
             # Remove dead code so they won't be matched again
             graph_module.graph.eliminate_dead_code()
@@ -122,11 +136,13 @@ def create_replacement_for_pattern(
 def replace_all_fusable_subgraphs(
     ep: ExportedProgram,
     graph_module: torch.fx.GraphModule,
+    skip_patterns: Collection[str] = (),
 ) -> int:
     total_replaced = 0
+    entries = enabled_patterns(skip_patterns)
 
     # Handle patterns identified with SubgraphMatcher
-    for entry in fusable_patterns.values():
+    for entry in entries:
         if entry.get_graphs_fn is not None and entry.create_replacement_fn is not None:
             total_replaced += create_replacement_for_pattern(
                 ep,
@@ -138,16 +154,23 @@ def replace_all_fusable_subgraphs(
 
     # Handle patterns identified with custom detector function
     for node in graph_module.graph.nodes:
-        for entry in fusable_patterns.values():
+        replaced_before = total_replaced
+        for entry in entries:
             if (
                 entry.detector_fn is not None
                 and entry.create_replacement_fn is not None
             ):
                 maybe_match = entry.detector_fn(node)
                 if maybe_match is not None:
                     assert entry.create_replacement_fn is not None
-                    entry.create_replacement_fn(ep, graph_module, maybe_match)
-                    total_replaced += 1
+                    if (
+                        entry.create_replacement_fn(ep, graph_module, maybe_match)
+                        is not False
+                    ):
+                        total_replaced += 1
+        if total_replaced > replaced_before:
+            # Remove dead code so they won't be matched again
+            graph_module.graph.eliminate_dead_code()
 
     graph_module.graph.eliminate_dead_code()
     return total_replaced