From 2fbdf6b07b531c903a4888ff4d0b9d338df18798 Mon Sep 17 00:00:00 2001 From: Stephen Jia Date: Mon, 9 Nov 2026 14:47:50 -0800 Subject: [PATCH] [ET-VK][conv2d] Fuse residual adds and activations into conv2d during lowering Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- backends/vulkan/_passes/fuse_patterns.py | 7 +- backends/vulkan/custom_ops_lib.py | 58 +++++ backends/vulkan/op_registry.py | 27 +++ .../vulkan/partitioner/vulkan_partitioner.py | 15 +- backends/vulkan/patterns/BUCK | 1 + backends/vulkan/patterns/__init__.py | 43 +++- .../vulkan/patterns/conv2d_with_epilogue.py | 225 ++++++++++++++++++ backends/vulkan/patterns/pattern_registry.py | 3 +- .../test/test_vulkan_compile_options.py | 6 + backends/vulkan/test/test_vulkan_delegate.py | 105 ++++++++ backends/vulkan/test/test_vulkan_passes.py | 195 +++++++++++++++ backends/vulkan/vulkan_preprocess.py | 11 +- backends/webgpu/README.md | 8 +- backends/webgpu/partitioner/partitioner.py | 9 + backends/webgpu/test/BUCK | 1 + .../webgpu/test/op_tests/generate_op_tests.py | 4 +- .../webgpu/test/ops/test_conv_with_clamp.py | 8 +- .../webgpu/test/test_webgpu_partitioner.py | 45 +++- .../backends/webgpu/webgpu-op-support.md | 5 +- .../source/backends/webgpu/webgpu-overview.md | 14 +- 20 files changed, 753 insertions(+), 37 deletions(-) create mode 100644 backends/vulkan/patterns/conv2d_with_epilogue.py diff --git a/backends/vulkan/_passes/fuse_patterns.py b/backends/vulkan/_passes/fuse_patterns.py index 1575dd6a4f6..403d9fc65c0 100644 --- a/backends/vulkan/_passes/fuse_patterns.py +++ b/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: diff --git a/backends/vulkan/custom_ops_lib.py b/backends/vulkan/custom_ops_lib.py index c3823532fef..023cb663500 100644 --- a/backends/vulkan/custom_ops_lib.py +++ b/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 ## ################# diff --git a/backends/vulkan/op_registry.py b/backends/vulkan/op_registry.py index 3a965c642cb..bcfa29fe89c 100644 --- a/backends/vulkan/op_registry.py +++ b/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 # ============================================================================= diff --git a/backends/vulkan/partitioner/vulkan_partitioner.py b/backends/vulkan/partitioner/vulkan_partitioner.py index e5df76aef54..36492f48cbb 100644 --- a/backends/vulkan/partitioner/vulkan_partitioner.py +++ b/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 diff --git a/backends/vulkan/patterns/BUCK b/backends/vulkan/patterns/BUCK index b83b2401428..c8d61b9c88e 100644 --- a/backends/vulkan/patterns/BUCK +++ b/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", diff --git a/backends/vulkan/patterns/__init__.py b/backends/vulkan/patterns/__init__.py index 4b3e8d87f62..3ca00b4af8e 100644 --- a/backends/vulkan/patterns/__init__.py +++ b/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,7 +154,8 @@ 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 @@ -146,8 +163,14 @@ def replace_all_fusable_subgraphs( 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 diff --git a/backends/vulkan/patterns/conv2d_with_epilogue.py b/backends/vulkan/patterns/conv2d_with_epilogue.py new file mode 100644 index 00000000000..aa927244a55 --- /dev/null +++ b/backends/vulkan/patterns/conv2d_with_epilogue.py @@ -0,0 +1,225 @@ +# 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. + +from typing import Optional + +import executorch.backends.vulkan.utils as utils + +import torch + +from executorch.backends.vulkan.patterns.pattern_registry import ( + PatternMatch, + register_pattern_detector, + register_pattern_replacement, +) + +from executorch.exir import ExportedProgram +from executorch.exir.dialects._ops import ops as exir_ops + +residual_ops = { + exir_ops.edge.aten.add.Tensor: "add", + exir_ops.edge.aten.sub.Tensor: "sub", + exir_ops.edge.aten.mul.Tensor: "mul", +} + + +def tensor_val(node: torch.fx.node.Argument) -> Optional[torch.Tensor]: + if not isinstance(node, torch.fx.Node): + return None + val = node.meta.get("val") + return val if isinstance(val, torch.Tensor) else None + + +def fits_channels_packed_texture(val: torch.Tensor) -> bool: + return bool( + utils.filter_invalid_reprs( + val, utils.CHANNELS_PACKED_TEXTURE, utils.DEFAULT_TEXTURE_LIMITS + ).valid_texture_layouts + ) + + +def is_fusable_conv(node: torch.fx.node.Argument) -> bool: + """ + A 2D, batch 1 floating point convolution with a single user, whose input and + output fit channels-packed textures. The partitioner delegates the nodes of a + match without checking them, so this repeats the convolution's own checks. + """ + out = tensor_val(node) + if ( + out is None + or out.dtype not in (torch.float16, torch.float32) + or node.target != exir_ops.edge.aten.convolution.default + ): + return False + assert isinstance(node, torch.fx.Node) + x = tensor_val(node.args[0]) + return ( + len(node.args) == 9 + and len(node.users) == 1 + and x is not None + and x.dim() == 4 + and x.shape[0] == 1 + and fits_channels_packed_texture(x) + and fits_channels_packed_texture(out) + ) + + +def activation_of(node: torch.fx.Node) -> Optional[tuple[str, float, float]]: + """Returns the activation name and parameters of a fusable activation.""" + # The fused op takes the parameters as constants. + if any(isinstance(arg, torch.fx.Node) for arg in node.args[1:]): + return None + if node.target == exir_ops.edge.aten.relu.default: + return ("relu", 0.0, 0.0) + if node.target == exir_ops.edge.aten.hardtanh.default: + lo = node.args[1] if len(node.args) > 1 else -1.0 + hi = node.args[2] if len(node.args) > 2 else 1.0 + return ("clamp", float(lo), float(hi)) + if node.target == exir_ops.edge.aten.clamp.default: + lo, hi = [*node.args[1:], None, None][:2] + lo = float("-inf") if lo is None else float(lo) + hi = float("inf") if hi is None else float(hi) + return ("clamp", lo, hi) + if node.target == exir_ops.edge.aten.leaky_relu.default: + slope = node.args[1] if len(node.args) > 1 else 0.01 + return ("leaky_relu", float(slope), 0.0) + return None + + +class Conv2dEpilogueMatch(PatternMatch): + """ + Matches activation(residual_op(conv(...), residual [* scale])) and + activation(conv(...)), where the activation and the per-channel scale of the + residual are optional. + """ + + def __init__(self, node: torch.fx.Node) -> None: # noqa: C901 + self.match_found = False + self.residual_node: Optional[torch.fx.Node] = None + self.scaled_residual_node: Optional[torch.fx.Node] = None + self.scale_node: Optional[torch.fx.Node] = None + self.residual_op = residual_ops.get(node.target) + self.alpha = node.kwargs.get("alpha", 1) + if self.residual_op is None: + # activation(conv(...)), where node is the conv's only user. + if activation_of(node) is None or not is_fusable_conv(node.args[0]): + return + self.residual_op = "add" + self.conv_node = node.args[0] + self.all_nodes = [self.conv_node] + self.match_activation(self.conv_node) + return + + # residual_op(conv, residual) must equal the matched node. + operands = [(node.args[0], node.args[1])] + if self.residual_op == "mul" or (self.residual_op == "add" and self.alpha == 1): + operands.append((node.args[1], node.args[0])) + for conv, residual in operands: + out, res = tensor_val(conv), tensor_val(residual) + if ( + is_fusable_conv(conv) + and residual is not conv + and res is not None + and out is not None + and res.shape == out.shape + and res.dtype == out.dtype + ): + break + else: + return + assert isinstance(conv, torch.fx.Node) and isinstance(residual, torch.fx.Node) + self.conv_node = conv + self.residual_node = residual + self.scaled_residual_node = residual + self.all_nodes = [conv, node] + + channels = out.shape[1] + if ( + residual.target == exir_ops.edge.aten.mul.Tensor + and len(residual.users) == 1 + ): + for x, scale in (residual.args, residual.args[::-1]): + x_val, scale_val = tensor_val(x), tensor_val(scale) + if ( + x_val is not None + and scale_val is not None + and x_val.shape == out.shape + and x_val.dtype == out.dtype + and scale_val.dtype == out.dtype + and tuple(scale_val.shape) + in ((1, channels, 1, 1), (channels, 1, 1)) + and isinstance(scale, torch.fx.Node) + and scale.op == "placeholder" + ): + assert isinstance(x, torch.fx.Node) + self.residual_node = x + self.scale_node = scale + self.all_nodes.append(residual) + break + + self.match_activation(node) + + def match_activation(self, node: torch.fx.Node) -> None: + """Completes the match, fusing the only user of `node` if an activation.""" + self.anchor_node = node + self.activation = ("none", 0.0, 0.0) + if len(node.users) == 1: + user = next(iter(node.users)) + activation = activation_of(user) + if activation is not None: + self.anchor_node = user + self.activation = activation + self.all_nodes.append(user) + + self.input_nodes = [ + n + for n in (self.conv_node.args[0], self.residual_node) + if isinstance(n, torch.fx.Node) + ] + self.output_nodes = [self.anchor_node] + self.match_found = True + + +@register_pattern_detector("conv2d_with_epilogue") +def find_conv2d_with_epilogue_pattern( + node: torch.fx.Node, +) -> Optional[Conv2dEpilogueMatch]: + if node.target not in residual_ops and activation_of(node) is None: + return None + match = Conv2dEpilogueMatch(node) + return match if match.match_found else None + + +@register_pattern_replacement("conv2d_with_epilogue") +def replace_conv2d_with_epilogue_pattern( + ep: ExportedProgram, + graph_module: torch.fx.GraphModule, + match: Conv2dEpilogueMatch, +) -> Optional[bool]: + residual, scale = match.residual_node, match.scale_node + # The runtime binds the residual as a tensor and prepacks the scale, so a + # constant x or a computed scale is fused as part of the residual x * scale. + if scale is not None and ( + not utils.is_param_node(ep, scale) or utils.is_param_node(ep, residual) + ): + residual, scale = match.scaled_residual_node, None + if residual is not None and utils.is_param_node(ep, residual): + return False + + with graph_module.graph.inserting_before(match.anchor_node): + fused = graph_module.graph.call_function( + exir_ops.edge.et_vk.conv2d.with_epilogue, + ( + *match.conv_node.args, + residual, + scale, + match.residual_op, + match.alpha, + *match.activation, + ), + ) + fused.meta = match.anchor_node.meta.copy() + match.anchor_node.replace_all_uses_with(fused) diff --git a/backends/vulkan/patterns/pattern_registry.py b/backends/vulkan/patterns/pattern_registry.py index de5c6e77a1d..dcb2384e329 100644 --- a/backends/vulkan/patterns/pattern_registry.py +++ b/backends/vulkan/patterns/pattern_registry.py @@ -68,8 +68,9 @@ def create_pattern_match_from_internal_match( ) +# Returns False if the replacement left the graph unchanged. CreateReplacementFn = Callable[ - [ExportedProgram, torch.fx.GraphModule, PatternMatch], None + [ExportedProgram, torch.fx.GraphModule, PatternMatch], Optional[bool] ] diff --git a/backends/vulkan/test/test_vulkan_compile_options.py b/backends/vulkan/test/test_vulkan_compile_options.py index ff507bb3197..fab630450c9 100644 --- a/backends/vulkan/test/test_vulkan_compile_options.py +++ b/backends/vulkan/test/test_vulkan_compile_options.py @@ -49,6 +49,12 @@ def test_force_fp16_round_trips(self) -> None: round_tripped = self._round_trip({"force_fp16": True}) self.assertTrue(round_tripped.get("force_fp16")) + def test_skip_fusion_patterns_round_trips(self) -> None: + for patterns in ({"rms_norm", "conv2d_with_epilogue"}, set()): + with self.subTest(patterns=patterns): + round_tripped = self._round_trip({"skip_fusion_patterns": patterns}) + self.assertEqual(round_tripped.get("skip_fusion_patterns"), patterns) + def test_external_constants_max_data_bytes_round_trips_uint64_bounds( self, ) -> None: diff --git a/backends/vulkan/test/test_vulkan_delegate.py b/backends/vulkan/test/test_vulkan_delegate.py index de940dace41..e2517cf8fc0 100644 --- a/backends/vulkan/test/test_vulkan_delegate.py +++ b/backends/vulkan/test/test_vulkan_delegate.py @@ -2480,6 +2480,111 @@ def forward(self, x): (torch.randn(size=(1, 6, 40, 50), dtype=torch.float32),), ) + def test_vulkan_backend_conv2d_with_epilogue(self): + class ResidualConvModule(torch.nn.Module): + def __init__(self, conv, activation): + super().__init__() + self.conv = conv + self.scale = torch.nn.Parameter(torch.randn(1, conv.out_channels, 1, 1)) + self.activation = activation + + def forward(self, x): + return self.activation(self.conv(x) + x * self.scale) + + # Direct, pointwise, depthwise, im2col and transposed convolutions. + for conv, activation, size in ( + (torch.nn.Conv2d(16, 16, 3, padding=1), torch.nn.Identity(), 40), + (torch.nn.Conv2d(16, 16, 1), torch.nn.ReLU(), 40), + ( + torch.nn.Conv2d(16, 16, 3, padding=1, groups=16), + torch.nn.LeakyReLU(0.2), + 40, + ), + (torch.nn.Conv2d(128, 128, 3, padding=1), torch.nn.ReLU(), 12), + ( + torch.nn.ConvTranspose2d(16, 16, 3, padding=1), + torch.nn.Hardtanh(-1.0, 2.0), + 20, + ), + ): + with self.subTest(conv=conv): + self.lower_module_and_test_output( + ResidualConvModule(conv, activation), + (torch.randn(size=(1, conv.in_channels, size, size)),), + ) + + def test_vulkan_backend_conv2d_with_epilogue_residual_ops(self): + class ResidualOpModule(torch.nn.Module): + def __init__(self, residual_op): + super().__init__() + self.conv = torch.nn.Conv2d(16, 16, 3, padding=1) + self.scale = torch.nn.Parameter(torch.randn(16, 1, 1)) + self.residual_op = residual_op + + def forward(self, x): + return torch.relu(self.residual_op(self.conv(x), x * self.scale)) + + for name, residual_op in ( + ("sub", lambda y, r: torch.sub(y, r, alpha=0.5)), + ("add", lambda y, r: torch.add(y, r, alpha=2.0)), + ("mul", lambda y, r: r * y), + ): + with self.subTest(residual_op=name): + self.lower_module_and_test_output( + ResidualOpModule(residual_op), (torch.randn(1, 16, 20, 20),) + ) + + def test_vulkan_backend_conv2d_with_epilogue_activation_only(self): + class ClampConvModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d(16, 16, 3, padding=1) + + def forward(self, x): + return torch.clamp(self.conv(x), min=0.0) + + for module in ( + torch.nn.Sequential(torch.nn.Conv2d(16, 16, 1), torch.nn.ReLU()), + torch.nn.Sequential( + torch.nn.Conv2d(16, 16, 3, padding=1), torch.nn.LeakyReLU(0.2) + ), + ClampConvModule(), + ): + with self.subTest(module=module): + self.lower_module_and_test_output(module, (torch.randn(1, 16, 20, 20),)) + + def test_vulkan_backend_conv2d_with_epilogue_constant_residual(self): + class ConstantResidualModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d(16, 16, 3, padding=1) + self.residual = torch.nn.Parameter(torch.randn(1, 16, 20, 20)) + self.scale = torch.nn.Parameter(torch.randn(1, 16, 1, 1)) + + def forward(self, x): + return torch.relu(self.conv(x) + self.residual * self.scale) + + self.lower_module_and_test_output( + ConstantResidualModule(), (torch.randn(1, 16, 20, 20),) + ) + + def test_vulkan_backend_conv2d_with_epilogue_dynamic_shapes(self): + class ResidualConvModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d(16, 16, 3, padding=1) + + def forward(self, x): + return torch.relu(self.conv(x) + x) + + height = Dim("height", min=4, max=32) + self.lower_module_and_test_output( + ResidualConvModule(), + (torch.randn(1, 16, 24, 20),), + dynamic_shapes={"x": {2: height}}, + test_inputs=[(torch.randn(1, 16, 9, 20),), (torch.randn(1, 16, 32, 20),)], + ) + def test_vulkan_backend_div_with_padding_nan_propagation(self): """ Test division operations with non-multiple-of-4 channels followed by convolution. diff --git a/backends/vulkan/test/test_vulkan_passes.py b/backends/vulkan/test/test_vulkan_passes.py index e1c7e762aa1..49dfb13af4c 100644 --- a/backends/vulkan/test/test_vulkan_passes.py +++ b/backends/vulkan/test/test_vulkan_passes.py @@ -2,6 +2,7 @@ from itertools import product from typing import List, Optional, Tuple +import executorch.backends.vulkan.patterns as vk_patterns import torch from executorch.backends.vulkan._passes.conv1d_as_conv2d import Conv1dAsConv2dPass @@ -18,6 +19,7 @@ unwrap_identity_repeat, ) from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.backends.vulkan.patterns.conv2d_with_epilogue import activation_of from executorch.backends.vulkan.patterns.swiglu import find_swiglu_pattern from executorch.backends.vulkan.serialization.vulkan_graph_schema import VkDataType from executorch.backends.vulkan.serialization.vulkan_graph_serialize import ( @@ -530,6 +532,199 @@ def forward(self, x): ) +class ResidualConvModule(torch.nn.Module): + """conv(x) combined with a residual, as selected by `variant`.""" + + def __init__(self, variant: str, activation=None): + super().__init__() + self.variant = variant + self.activation = activation + conv = torch.nn.ConvTranspose2d if variant == "transposed" else torch.nn.Conv2d + self.conv = conv(16, 16, kernel_size=3, padding=1) + self.scale = torch.nn.Parameter(torch.randn(1, 16, 1, 1)) + self.residual = torch.nn.Parameter(torch.randn(1, 16, 12, 12)) + + def forward(self, x, other): + y = self.conv(x) + if self.variant == "scaled": + out = y + x * self.scale + elif self.variant == "input_scaled": + out = y + x * other + elif self.variant in ("swapped", "transposed"): + out = x + y + elif self.variant == "sub_swapped": + out = x - y + elif self.variant == "mul": + out = x * y + elif self.variant == "self": + out = y + y + elif self.variant == "param_residual": + out = y + self.residual + elif self.variant == "param_scaled": + out = y + self.residual * self.scale + elif self.variant == "broadcast": + out = y + other + elif self.variant == "shared": + return y + x, y + else: + out = torch.sub(y, x, alpha=0.5) + return out if self.activation is None else self.activation(out) + + +class TestConv2dEpilogueFusion(unittest.TestCase): + def _lowered_ops(self, model, inputs, **partitioner_kwargs): + graphs = vulkan_delegate_graphs( + lower_to_vulkan(model.eval(), inputs, **partitioner_kwargs) + ) + return [op.name for graph in graphs for op in graph.chain] + + def _op_names(self, model, other_sizes=(1, 16, 12, 12)): + inputs = (torch.randn(1, 16, 12, 12), torch.randn(other_sizes)) + return self._lowered_ops(model, inputs) + + def _conv_activation_ops(self, activation, **partitioner_kwargs): + model = torch.nn.Sequential( + torch.nn.Conv2d(16, 16, kernel_size=3, padding=1), activation + ) + inputs = (torch.randn(1, 16, 12, 12),) + return self._lowered_ops(model, inputs, **partitioner_kwargs) + + def test_fused(self): + for variant, activation, other_sizes in ( + ("scaled", None, (1, 16, 12, 12)), + ("swapped", torch.nn.ReLU(), (1, 16, 12, 12)), + ("sub", torch.nn.Hardtanh(-1.0, 2.0), (1, 16, 12, 12)), + ("scaled", torch.nn.LeakyReLU(0.2), (1, 16, 12, 12)), + ("transposed", torch.nn.ReLU(), (1, 16, 12, 12)), + ("mul", torch.nn.ReLU(), (1, 16, 12, 12)), + ): + with self.subTest(variant=variant, activation=activation): + ops = self._op_names( + ResidualConvModule(variant, activation), other_sizes + ) + self.assertEqual(ops.count("et_vk.conv2d.with_epilogue"), 1) + self.assertEqual(ops.count("aten.convolution.default"), 0) + self.assertFalse( + any( + op.startswith(("aten.add", "aten.sub", "aten.mul")) + for op in ops + ) + ) + self.assertFalse(any("relu" in op or "hardtanh" in op for op in ops)) + + def test_activation_fused_without_residual(self): + for activation in ( + torch.nn.ReLU(), + torch.nn.ReLU6(), + torch.nn.Hardtanh(-1.0, 2.0), + torch.nn.LeakyReLU(0.2), + ): + with self.subTest(activation=activation): + ops = self._conv_activation_ops(activation) + self.assertEqual(ops.count("et_vk.conv2d.with_epilogue"), 1) + self.assertEqual(ops.count("aten.convolution.default"), 0) + self.assertEqual(ops.count("et_vk.conv_with_clamp.default"), 0) + self.assertFalse(any("relu" in op or "hardtanh" in op for op in ops)) + + def test_conv1d_activation_uses_conv_with_clamp(self): + model = torch.nn.Sequential( + torch.nn.Conv1d(16, 16, kernel_size=3, padding=1, groups=16), + torch.nn.ReLU(), + ) + ops = self._lowered_ops(model, (torch.randn(1, 16, 12),)) + self.assertEqual(ops.count("et_vk.conv_with_clamp.default"), 1) + self.assertEqual(ops.count("et_vk.conv2d.with_epilogue"), 0) + + def test_skip_fusion_patterns(self): + options = {"skip_fusion_patterns": {"conv2d_with_epilogue"}} + ops = self._conv_activation_ops(torch.nn.ReLU(), compile_options=options) + self.assertEqual(ops.count("et_vk.conv_with_clamp.default"), 1) + self.assertEqual(ops.count("et_vk.conv2d.with_epilogue"), 0) + + def test_skip_fusion_patterns_rejects_invalid_values(self): + for value in ("conv2d_with_epilogue", {"no_such_pattern"}): + with self.subTest(value=value), self.assertRaises(ValueError): + VulkanPartitioner(compile_options={"skip_fusion_patterns": value}) + + def test_conv_exceeding_texture_limits_stays_unfused(self): + model = torch.nn.Sequential( + torch.nn.Conv2d(16, 16, kernel_size=1), torch.nn.ReLU() + ) + ops = self._lowered_ops(model, (torch.randn(1, 16, 1, 16385),)) + self.assertEqual(ops.count("et_vk.conv2d.with_epilogue"), 0) + self.assertEqual(ops.count("aten.convolution.default"), 0) + + def test_activation_parameters(self): + graph = torch.fx.Graph() + x = graph.placeholder("x") + bound = graph.placeholder("bound") + for target, args, expected in ( + (exir_ops.edge.aten.hardtanh.default, (x, 0.0), ("clamp", 0.0, 1.0)), + ( + exir_ops.edge.aten.clamp.default, + (x, None, 6.0), + ("clamp", float("-inf"), 6.0), + ), + (exir_ops.edge.aten.clamp.default, (x, bound), None), + ): + with self.subTest(target=target, args=args): + node = graph.call_function(target, args) + self.assertEqual(activation_of(node), expected) + + def test_quantized_conv_activation_uses_q8ta(self): + from executorch.backends.xnnpack.quantizer.xnnpack_quantizer import ( + get_symmetric_quantization_config, + XNNPACKQuantizer, + ) + + model = torch.nn.Sequential( + torch.nn.Conv2d(16, 16, kernel_size=3, padding=1), torch.nn.ReLU() + ) + quantizer = XNNPACKQuantizer() + quantizer.set_global(get_symmetric_quantization_config(is_per_channel=True)) + edge_program = quantize_and_lower_module( + model.eval(), (torch.randn(1, 16, 12, 12),), quantizer + ) + ep = edge_program.exported_program() + # The conv and relu the quantized pattern replaces must not match again. + self.assertEqual( + vk_patterns.replace_all_fusable_subgraphs(ep, ep.graph_module), 1 + ) + self.assertEqual(op_node_count(ep.graph_module, "q8ta_conv2d.default"), 1) + + def test_constant_residual_fuses_scaled_residual(self): + ops = self._op_names(ResidualConvModule("param_scaled")) + self.assertEqual(ops.count("et_vk.conv2d.with_epilogue"), 1) + self.assertEqual(ops.count("aten.mul.Tensor"), 1) + self.assertEqual(ops.count("aten.add.Tensor"), 0) + + def test_skipped_replacement_is_not_counted(self): + inputs = (torch.randn(1, 16, 12, 12), torch.randn(1, 16, 12, 12)) + ep = export_to_edge(ResidualConvModule("param_residual"), inputs) + self.assertEqual( + vk_patterns.replace_all_fusable_subgraphs(ep, ep.graph_module), 0 + ) + self.assertEqual(op_node_count(ep.graph_module, "convolution.default"), 1) + + def test_input_scale_stays_unfused(self): + ops = self._op_names(ResidualConvModule("input_scaled"), (1, 16, 1, 1)) + self.assertEqual(ops.count("et_vk.conv2d.with_epilogue"), 1) + self.assertEqual(ops.count("aten.mul.Tensor"), 1) + + def test_not_fused(self): + for variant, other_sizes in ( + ("sub_swapped", (1, 16, 12, 12)), + ("broadcast", (1, 16, 1, 1)), + ("shared", (1, 16, 12, 12)), + ("self", (1, 16, 12, 12)), + ("param_residual", (1, 16, 12, 12)), + ): + with self.subTest(variant=variant): + ops = self._op_names(ResidualConvModule(variant), other_sizes) + self.assertEqual(ops.count("et_vk.conv2d.with_epilogue"), 0) + self.assertEqual(ops.count("aten.convolution.default"), 1) + + class TestVulkanPasses(unittest.TestCase): def test_embedding_q4gsw_reference_supports_per_row_scales(self): import executorch.backends.vulkan.custom_ops_lib # noqa: F401 diff --git a/backends/vulkan/vulkan_preprocess.py b/backends/vulkan/vulkan_preprocess.py index ed5729e5a29..e5d78188c11 100644 --- a/backends/vulkan/vulkan_preprocess.py +++ b/backends/vulkan/vulkan_preprocess.py @@ -106,7 +106,9 @@ def _parse_external_constants_max_data_bytes(value_bytes: bytes) -> int: return value -def parse_compile_spec(compile_specs: List[CompileSpec]) -> Dict[str, Any]: +def parse_compile_spec( # noqa: C901 + compile_specs: List[CompileSpec], +) -> Dict[str, Any]: options = {} for spec in compile_specs: if spec.key == "storage_type_override": @@ -141,6 +143,9 @@ def parse_compile_spec(compile_specs: List[CompileSpec]) -> Dict[str, Any]: if spec.key == "external_constants_max_data_bytes": options[spec.key] = _parse_external_constants_max_data_bytes(spec.value) + if spec.key == "skip_fusion_patterns": + options[spec.key] = set(filter(None, spec.value.decode().split(","))) + # Unhandled options are ignored return options @@ -194,8 +199,10 @@ def preprocess( # noqa: C901 FuseBatchNormPass(program), AddmmToLinearTransform(), InsertDtypePromotionPass(), - FusePatternsPass(), + FusePatternsPass(compile_options.get("skip_fusion_patterns", ())), Conv1dAsConv2dPass(), + # Fuses activations into the convolutions FusePatternsPass left + # unfused, such as conv1d. FuseClampPass(), RemoveRedundantOpsTransform(), ReplaceInstanceNormPass(), diff --git a/backends/webgpu/README.md b/backends/webgpu/README.md index 5b9877b9601..fc289cb3f10 100644 --- a/backends/webgpu/README.md +++ b/backends/webgpu/README.md @@ -110,10 +110,12 @@ Key design choices: storage annotations and allocates tensors as buffers. - **Built-in WGSL shaders.** Shader source is embedded as C++ string constants, with build-time code generation and drift validation. -- **Thin Python AOT frontend.** `WebGPUPartitioner` delegates unchanged to +- **Thin Python AOT frontend.** `WebGPUPartitioner` delegates to `VulkanPartitioner`, preserving Vulkan serialization and the `VulkanBackend` - delegate id. It does not independently validate the narrower WebGPU runtime - capability set; WebGPU integrations apply that policy separately. + delegate id. It skips the Vulkan fusion patterns whose fused ops the WebGPU + runtime does not implement. It does not independently validate the narrower + WebGPU runtime capability set; WebGPU integrations apply that policy + separately. - **Dynamic shapes.** Tensors allocate at their maximum shape, with SymInt arithmetic and per-operator resize hooks for runtime dimensions. - **Shape-routed dispatch.** Runtime tensor dimensions select specialized diff --git a/backends/webgpu/partitioner/partitioner.py b/backends/webgpu/partitioner/partitioner.py index 59ab2ffb6b7..dd88741c50a 100644 --- a/backends/webgpu/partitioner/partitioner.py +++ b/backends/webgpu/partitioner/partitioner.py @@ -36,6 +36,10 @@ def _arange_bounds_are_static(node: torch.fx.Node) -> bool: } +# Vulkan fusion patterns whose fused ops the WebGPU runtime does not implement. +_WEBGPU_SKIP_FUSION_PATTERNS = {"conv2d_with_epilogue"} + + @final class WebGPUPartitioner(Partitioner): """WebGPU frontend for the Vulkan partitioning and serialization path.""" @@ -48,6 +52,11 @@ def __init__( nn_module_blocklist: Optional[List[str]] = None, nn_module_allowlist: Optional[List[str]] = None, ) -> None: + compile_options = dict(compile_options or {}) + compile_options["skip_fusion_patterns"] = { + *compile_options.get("skip_fusion_patterns", ()), + *_WEBGPU_SKIP_FUSION_PATTERNS, + } self._vulkan_partitioner = VulkanPartitioner( compile_options=compile_options, operator_blocklist=operator_blocklist, diff --git a/backends/webgpu/test/BUCK b/backends/webgpu/test/BUCK index bc6ea6863e2..51649178c1d 100644 --- a/backends/webgpu/test/BUCK +++ b/backends/webgpu/test/BUCK @@ -42,6 +42,7 @@ fbcode_target( deps = [ "//caffe2:torch", "//executorch/backends/vulkan/partitioner:vulkan_partitioner", + "//executorch/backends/vulkan/serialization:lib", "//executorch/backends/vulkan:vulkan_preprocess", "//executorch/backends/webgpu/partitioner:webgpu_partitioner", "//executorch/backends/webgpu/test:tester", diff --git a/backends/webgpu/test/op_tests/generate_op_tests.py b/backends/webgpu/test/op_tests/generate_op_tests.py index 33b7940a9be..0d225cfe70d 100644 --- a/backends/webgpu/test/op_tests/generate_op_tests.py +++ b/backends/webgpu/test/op_tests/generate_op_tests.py @@ -19,7 +19,7 @@ import os import torch -from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.backends.webgpu.partitioner import WebGPUPartitioner # Importing cases populates op_test_registry via the @register_op_test decorators. from executorch.backends.webgpu.test.op_tests import cases # noqa: F401 @@ -67,7 +67,7 @@ def export_case(suite: WebGPUTestSuite, case) -> tuple[torch.nn.Module, tuple, o export_inputs = tuple(_materialize(s) for s in case.export_inputs) ep = torch.export.export(module, export_inputs, dynamic_shapes=case.dynamic_shapes) prog = to_edge_transform_and_lower( - ep, partitioner=[VulkanPartitioner()] + ep, partitioner=[WebGPUPartitioner()] ).to_executorch() return module, inputs, prog diff --git a/backends/webgpu/test/ops/test_conv_with_clamp.py b/backends/webgpu/test/ops/test_conv_with_clamp.py index 511c5857cdc..c7573f3a3cf 100644 --- a/backends/webgpu/test/ops/test_conv_with_clamp.py +++ b/backends/webgpu/test/ops/test_conv_with_clamp.py @@ -7,8 +7,8 @@ """`et_vk.conv_with_clamp` module + configs for the WebGPU op-test framework. `ConvWithClampModule` is a plain `nn.Conv2d` followed by `F.relu6`, which the -Vulkan fusion rewrites to a delegated `et_vk.conv_with_clamp` (fp32 general conv -+ output clamp[0,6]). The conv weight/bias are baked params; only the float +WebGPU lowering rewrites to a delegated `et_vk.conv_with_clamp` (fp32 general +conv + output clamp[0,6]). The conv weight/bias are baked params; only the float tensor `x` is a runtime input. Goldened vs the module's fp32 eager. `ConvWithClampTest` is the export-delegation smoke test. """ @@ -20,7 +20,7 @@ import torch import torch.nn as nn -from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.backends.webgpu.partitioner import WebGPUPartitioner from executorch.exir import to_edge_transform_and_lower # name -> (input_shape, ic, oc, k, stride, padding, dilation, bias) @@ -63,7 +63,7 @@ def _det(shape): def _lower(shape, ic, oc, k, stride, padding, dilation, bias): mod = ConvWithClampModule(ic, oc, k, stride, padding, dilation, bias).eval() ep = torch.export.export(mod, (_det(shape),)) - return to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()]) + return to_edge_transform_and_lower(ep, partitioner=[WebGPUPartitioner()]) def _delegated(et) -> bool: diff --git a/backends/webgpu/test/test_webgpu_partitioner.py b/backends/webgpu/test/test_webgpu_partitioner.py index e94ee51aa35..7249df02e66 100644 --- a/backends/webgpu/test/test_webgpu_partitioner.py +++ b/backends/webgpu/test/test_webgpu_partitioner.py @@ -11,7 +11,14 @@ import torch from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.backends.vulkan.serialization.vulkan_graph_serialize import ( + extract_vk_flatbuffer, + flatbuffer_to_vk_graph, +) from executorch.backends.webgpu.partitioner import WebGPUPartitioner +from executorch.backends.webgpu.partitioner.partitioner import ( + _WEBGPU_SKIP_FUSION_PATTERNS, +) from executorch.backends.webgpu.test.tester import Partition, ToEdgeTransformAndLower from executorch.exir import ( ExecutorchProgramManager, @@ -28,6 +35,15 @@ def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: return a + b +class ConvReluModule(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(4, 4, 3, padding=1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.relu(self.conv(x)) + + class ArangeModule(torch.nn.Module): def __init__(self, dtype: torch.dtype) -> None: super().__init__() @@ -76,18 +92,43 @@ def _assert_same_partitioning( def test_lowers_identically_with_vulkan_backend_id(self) -> None: webgpu_program = self._lower(WebGPUPartitioner()) - vulkan_program = self._lower(VulkanPartitioner()) + vulkan_program = self._lower( + VulkanPartitioner( + compile_options={"skip_fusion_patterns": _WEBGPU_SKIP_FUSION_PATTERNS} + ) + ) self.assertEqual(webgpu_program.buffer, vulkan_program.buffer) self.assertEqual(self._delegate_ids(webgpu_program), ["VulkanBackend"]) + def test_fuses_conv_activation_into_conv_with_clamp(self) -> None: + # The WebGPU runtime implements conv_with_clamp, not conv2d.with_epilogue. + edge = to_edge_transform_and_lower( + torch.export.export(ConvReluModule(), (torch.randn(1, 4, 8, 8),)), + partitioner=[WebGPUPartitioner()], + ) + ops = [ + op.name + for _, module in edge.exported_program().graph_module.named_children() + for op in flatbuffer_to_vk_graph( + extract_vk_flatbuffer(module.processed_bytes) + ).chain + ] + self.assertIn("et_vk.conv_with_clamp.default", ops) + self.assertNotIn("et_vk.conv2d.with_epilogue", ops) + def test_forwards_compile_options(self) -> None: compile_options = {"skip_bool_tensors": True} expected_options = compile_options.copy() result = self._assert_same_partitioning( WebGPUPartitioner(compile_options=compile_options), - VulkanPartitioner(compile_options=compile_options), + VulkanPartitioner( + compile_options={ + **compile_options, + "skip_fusion_patterns": _WEBGPU_SKIP_FUSION_PATTERNS, + } + ), ) self.assertGreater(len(result.partition_tags), 0) diff --git a/docs/source/backends/webgpu/webgpu-op-support.md b/docs/source/backends/webgpu/webgpu-op-support.md index b5ee43252d8..f7486248670 100644 --- a/docs/source/backends/webgpu/webgpu-op-support.md +++ b/docs/source/backends/webgpu/webgpu-op-support.md @@ -9,8 +9,9 @@ and the generated native test catalog is defined in ## Compatibility Considerations -- `WebGPUPartitioner` is a thin wrapper over `VulkanPartitioner`; it preserves - Vulkan partitioning behavior and does not add WebGPU capability policy. +- `WebGPUPartitioner` is a thin wrapper over `VulkanPartitioner`; apart from + skipping fusions the WebGPU runtime does not implement, it preserves Vulkan + partitioning behavior and does not add WebGPU capability policy. - Support can depend on tensor rank, shape, layout, and data type in addition to the operator name. - FP16 shader paths require a WebGPU adapter exposing `shader-f16`. diff --git a/docs/source/backends/webgpu/webgpu-overview.md b/docs/source/backends/webgpu/webgpu-overview.md index ebc46031a53..2c5fd7fd2f6 100644 --- a/docs/source/backends/webgpu/webgpu-overview.md +++ b/docs/source/backends/webgpu/webgpu-overview.md @@ -42,12 +42,12 @@ Vulkan FlatBuffer delegate payload (VH00/VK00) inside a .pte file └── Browser: Emscripten + emdawnwebgpu → browser WebGPU implementation ``` -`WebGPUPartitioner` is a thin wrapper over `VulkanPartitioner`. It forwards -partitioning unchanged and preserves Vulkan serialization and the -`VulkanBackend` delegate identifier. It does not independently validate the -narrower WebGPU runtime capability set. The runtime ignores Vulkan -texture-storage annotations and executes the graph with WebGPU buffers and -WGSL compute pipelines. +`WebGPUPartitioner` is a thin wrapper over `VulkanPartitioner`. It preserves +Vulkan serialization and the `VulkanBackend` delegate identifier, and skips the +Vulkan fusion patterns whose fused operators the WebGPU runtime does not +implement. It does not independently validate the narrower WebGPU runtime +capability set. The runtime ignores Vulkan texture-storage annotations and +executes the graph with WebGPU buffers and WGSL compute pipelines. ::::{important} The WebGPU runtime currently registers with ExecuTorch as `VulkanBackend` so it @@ -112,7 +112,7 @@ with open("add_webgpu.pte", "wb") as file: file.write(et_program.buffer) ``` -Because `WebGPUPartitioner` delegates unchanged to `VulkanPartitioner`, it can +Because `WebGPUPartitioner` otherwise partitions like `VulkanPartitioner`, it can select operators that do not yet have a WebGPU implementation. Validate the delegated graph against the WebGPU registry and operator tests. A Vulkan-supported but WebGPU-unsupported operator can fail when the WebGPU