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

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