From 96e5183a55445dc14ed60cdd6927570c5741a7b0 Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Fri, 30 Oct 2026 16:36:10 -0400 Subject: [PATCH] [XNNPACK] Delegate GEMMs whose weight is computed from parameters A convolution, linear, addmm or mm whose weight or bias is computed from parameters, such as a convolution under weight_norm or spectral_norm, was declined by the XNNPACK partitioner and ran on the portable kernels together with the computation of its weight. The GEMM configs now accept such a weight or bias when it is computed with pure ops from parameters, unmutated buffers or constants, the computation feeds only that GEMM, and the folded tensor is no larger than the parameters it replaces. The computation is partitioned with the GEMM, and FoldComputedParamsPass folds it into a constant at the start of XNNPACK preprocess, inside the delegate, so the top-level program is not modified. AddmmConfig now also rejects addmm with non-default alpha or beta, which the XNNPACK lowering ignored. Fixes #22078 This change was authored with Claude Code. --- backends/xnnpack/_passes/__init__.py | 4 + .../xnnpack/_passes/fold_computed_params.py | 70 ++++++ .../xnnpack/partition/config/gemm_configs.py | 17 +- .../test/passes/test_fold_computed_params.py | 235 ++++++++++++++++++ backends/xnnpack/utils/utils.py | 66 +++++ 5 files changed, 388 insertions(+), 4 deletions(-) create mode 100644 backends/xnnpack/_passes/fold_computed_params.py create mode 100644 backends/xnnpack/test/passes/test_fold_computed_params.py diff --git a/backends/xnnpack/_passes/__init__.py b/backends/xnnpack/_passes/__init__.py index 0051f8ceb88..5f17356bfea 100644 --- a/backends/xnnpack/_passes/__init__.py +++ b/backends/xnnpack/_passes/__init__.py @@ -26,6 +26,9 @@ ) from executorch.backends.xnnpack._passes.decompose_batch_norm import DecomposeBatchNorm from executorch.backends.xnnpack._passes.decompose_cat import DecomposeConcatenate +from executorch.backends.xnnpack._passes.fold_computed_params import ( + FoldComputedParamsPass, +) from executorch.backends.xnnpack._passes.fuse_activation_pass import FuseActivationPass from executorch.backends.xnnpack._passes.fuse_batch_norm import FuseBatchNormPass from executorch.backends.xnnpack._passes.fuse_pad_into_conv_pass import ( @@ -74,6 +77,7 @@ def __init__( if not passes: # All the XNNPACK passes self.passes = [ + FoldComputedParamsPass, RewriteFp16SiluPass, XNNPACKRemoveCloneOpsTransform, # TODO - remove this pass once we have a better support for dim_order ops lowering diff --git a/backends/xnnpack/_passes/fold_computed_params.py b/backends/xnnpack/_passes/fold_computed_params.py new file mode 100644 index 00000000000..80816210d6f --- /dev/null +++ b/backends/xnnpack/_passes/fold_computed_params.py @@ -0,0 +1,70 @@ +# 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. + +import torch +from executorch.backends.transforms.utils import create_constant_placeholder +from executorch.backends.xnnpack._passes.xnnpack_pass import XNNPACKPass +from executorch.backends.xnnpack.utils.utils import COMPUTED_PARAM_TAG, get_param_tensor +from executorch.exir.pass_base import PassResult +from torch.export.graph_signature import InputKind + + +class FoldComputedParamsPass(XNNPACKPass): + """ + Folds the nodes that the partitioner tagged with COMPUTED_PARAM_TAG, which + compute a GEMM weight or bias from params, into a param holding the result. + """ + + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: + graph = graph_module.graph + names = {node.name for node in graph.nodes} + values = {} + for node in graph.nodes: + # Checking a GEMM tags the nodes computing its weight or bias before + # the GEMM is accepted, so they are folded only next to that GEMM. + if not names & node.meta.get(COMPUTED_PARAM_TAG, set()): + continue + args, kwargs = torch.fx.node.map_arg( + (node.args, node.kwargs), + lambda n: ( + values[n] + if n in values + else get_param_tensor(self.exported_program, n) + ), + ) + with torch.no_grad(): + values[node] = node.target(*args, **kwargs) + + if not values: + return PassResult(graph_module, False) + + taken = ( + names + | self.exported_program.state_dict.keys() + | self.exported_program.constants.keys() + ) + for node, value in values.items(): + if all(user in values for user in node.users): + continue + name = f"{node.name}_folded" + while name in taken: + name += "_" + taken.add(name) + with graph.inserting_before(next(iter(graph.nodes))): + param = create_constant_placeholder( + exp_program=self.exported_program, + graph=graph, + kind=InputKind.PARAMETER, + name=name, + data=value, + ) + node.replace_all_uses_with(param) + + for node in reversed(values): + graph.erase_node(node) + + graph_module.recompile() + return PassResult(graph_module, True) diff --git a/backends/xnnpack/partition/config/gemm_configs.py b/backends/xnnpack/partition/config/gemm_configs.py index 239a424965b..8d5a760f89b 100644 --- a/backends/xnnpack/partition/config/gemm_configs.py +++ b/backends/xnnpack/partition/config/gemm_configs.py @@ -29,6 +29,7 @@ ) from executorch.backends.xnnpack.utils.utils import ( get_input_node, + get_static_param_deps, is_depthwise_conv, is_getitem, is_node, @@ -153,10 +154,11 @@ def _get_weight_deps( if precision == ConfigPrecisionType.FP32: # First find the weight weight_node = get_input_node(node, self.weight_idx) - if not is_param_node(ep, weight_node): + weight_deps = get_static_param_deps(ep, weight_node, node) + if weight_deps is None: why(node, "Expected weight to be a static param") return (False, []) - gemm_deps.append(weight_node) + gemm_deps.extend(weight_deps) return (True, gemm_deps) else: @@ -248,10 +250,11 @@ def _get_bias_deps( if len(node.all_input_nodes) > 2 and self.bias_idx is not None: bias_node = get_input_node(node, self.bias_idx) if bias_node: - if not is_param_node(ep, bias_node): + bias_deps = get_static_param_deps(ep, bias_node, node) + if bias_deps is None: why(node, "Expected bias to be a static param") return (False, []) - gemm_deps.append(bias_node) + gemm_deps.extend(bias_deps) return (True, gemm_deps) @@ -479,6 +482,12 @@ def __init__(self, **kwargs): self.src_partitions = None self.linear_modules = [torch.nn.functional.linear, torch.nn.Linear] + def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool: + if node.kwargs.get("alpha", 1) != 1 or node.kwargs.get("beta", 1) != 1: + why(node, "XNNPACK does not support addmm with alpha or beta") + return False + return super().check_constraints(node, ep) + def _get_weight_deps( self, node: torch.fx.Node, ep: ExportedProgram, precision: ConfigPrecisionType ) -> Tuple[bool, List[torch.fx.Node]]: diff --git a/backends/xnnpack/test/passes/test_fold_computed_params.py b/backends/xnnpack/test/passes/test_fold_computed_params.py new file mode 100644 index 00000000000..80fdafbec91 --- /dev/null +++ b/backends/xnnpack/test/passes/test_fold_computed_params.py @@ -0,0 +1,235 @@ +# 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. + +import unittest + +import torch +import torch.nn.functional as F +from executorch.backends.xnnpack.partition.config.generic_node_configs import MulConfig +from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner +from executorch.backends.xnnpack.test.tester import Tester, ToEdgeTransformAndLower +from torch.nn.utils.parametrizations import weight_norm + + +class TestFoldComputedParams(unittest.TestCase): + def setUp(self): + torch._dynamo.reset() + + class WeightNormConv1d(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv = weight_norm( + torch.nn.Conv1d(8, 8, 5, padding=2, groups=2), dim=2 + ) + + def forward(self, x): + return self.conv(x) + + class ComputedLinear(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.randn(4, 8)) + self.bias = torch.nn.Parameter(torch.randn(4)) + + def forward(self, x): + return F.linear(x, self.weight * 2, self.bias + 1) + + class FoldedNameCollision(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.randn(4, 8)) + # The name the folded weight would get by default. + self.aten_mul_tensor_folded = torch.nn.Parameter(torch.randn(4)) + + def forward(self, x): + return F.linear(x, self.weight * 2, self.aten_mul_tensor_folded) + + class MatmulWeightConv1d(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.randn(3, 3)) + + def forward(self, x): + return F.conv1d(x, (self.weight @ self.weight).reshape(3, 3, 1)) + + class NestedGemmWeightConv1d(torch.nn.Module): + def __init__(self): + super().__init__() + self.a = torch.nn.Parameter(torch.randn(3, 3)) + self.b = torch.nn.Parameter(torch.randn(3, 3)) + + def forward(self, x): + return F.conv1d(x, (self.a @ self.b.flip([1])).reshape(3, 3, 1)) + + class ScaledConv1d(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.randn(4, 4, 3)) + + def forward(self, x): + return F.conv1d(x, self.weight * 2) + + class ScaledAddmm(torch.nn.Module): + def __init__(self, computed): + super().__init__() + self.computed = computed + self.weight = torch.nn.Parameter(torch.randn(8, 4)) + self.bias = torch.nn.Parameter(torch.randn(4)) + + def forward(self, x): + weight, bias = self.weight, self.bias + if self.computed: + weight, bias = weight * 2, bias + 1 + return torch.addmm(bias, x, weight, alpha=2, beta=3) + + class SourceUsedElsewhere(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.randn(4, 8)) + + def forward(self, x): + return F.linear(x, self.weight * 2), self.weight.sum() + + class SourceRetainedByGemm(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.randn(16)) + + def forward(self, x): + return F.linear(x, (self.weight * 2).unsqueeze(1), self.weight) + + class WeightComputationFeedsActivation(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.randn(16, 16)) + + def forward(self, x): + w = self.weight * 2 + return w @ (w * 2) + x + + class LowRankWeight(torch.nn.Module): + def __init__(self): + super().__init__() + self.a = torch.nn.Parameter(torch.randn(16, 2)) + self.b = torch.nn.Parameter(torch.randn(2, 16)) + + def forward(self, x): + return F.linear(x, self.a @ self.b) + + class MutatedBufferSource(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("weight", torch.randn(4, 4, 3)) + + def forward(self, x): + y = F.conv1d(x, self.weight * 2) + self.weight.copy_(x.mean() * torch.ones(4, 4, 3)) + return y + + def _test_fully_delegated(self, module, inputs): + ( + Tester(module.eval(), inputs) + .export() + .to_edge_transform_and_lower() + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .check_not(["executorch_exir_dialects_edge__ops_"]) + .to_executorch() + .serialize() + .run_method_and_compare_outputs() + ) + + def test_fp32_weight_norm_conv1d(self): + self._test_fully_delegated(self.WeightNormConv1d(), (torch.randn(1, 8, 16),)) + + def test_fp32_computed_linear_weight_and_bias(self): + self._test_fully_delegated(self.ComputedLinear(), (torch.randn(2, 8),)) + + def test_fp32_folded_name_does_not_replace_param(self): + self._test_fully_delegated(self.FoldedNameCollision(), (torch.randn(2, 8),)) + + def test_fp32_computed_weight_of_rejected_gemm_not_folded(self): + # Conv1d with two dynamic dims is rejected after its weight was checked, + # while the matmul computing the weight is still delegated on its own. + dims = { + 0: torch.export.Dim("batch", min=1, max=4), + 2: torch.export.Dim("length", min=4, max=16), + } + for module in (self.MatmulWeightConv1d(), self.NestedGemmWeightConv1d()): + with self.subTest(module=type(module).__name__): + ( + Tester( + module.eval(), (torch.randn(2, 3, 8),), dynamic_shapes=(dims,) + ) + .export() + .to_edge_transform_and_lower() + .check( + ["executorch_exir_dialects_edge__ops_aten_convolution_default"] + ) + .to_executorch() + .serialize() + .run_method_and_compare_outputs() + ) + + def test_fp32_scaled_addmm_not_delegated(self): + for computed in (False, True): + with self.subTest(computed=computed): + ( + Tester(self.ScaledAddmm(computed).eval(), (torch.randn(2, 8),)) + .export() + .to_edge_transform_and_lower() + .check(["executorch_exir_dialects_edge__ops_aten_addmm_default"]) + .to_executorch() + .serialize() + .run_method_and_compare_outputs() + ) + + def test_fp32_weight_computed_by_prior_delegate_not_folded(self): + ( + Tester(self.ScaledConv1d().eval(), (torch.randn(1, 4, 8),)) + .export() + .to_edge_transform_and_lower( + ToEdgeTransformAndLower( + [XnnpackPartitioner(configs=[MulConfig]), XnnpackPartitioner()] + ) + ) + .check(["executorch_exir_dialects_edge__ops_aten_convolution_default"]) + .to_executorch() + .serialize() + .run_method_and_compare_outputs() + ) + + def test_fp32_computed_weight_not_folded(self): + for module, inputs, gemm in ( + (self.SourceUsedElsewhere(), (torch.randn(2, 8),), "aten_mm_default"), + (self.LowRankWeight(), (torch.randn(2, 16),), "aten_mm_default"), + (self.SourceRetainedByGemm(), (torch.randn(2, 1),), "aten_addmm_default"), + ( + self.WeightComputationFeedsActivation(), + (torch.randn(16, 16),), + "aten_mm_default", + ), + ( + self.MutatedBufferSource(), + (torch.randn(1, 4, 8),), + "aten_convolution_default", + ), + ): + with self.subTest(module=type(module).__name__): + ( + Tester(module.eval(), inputs) + .export() + .to_edge_transform_and_lower() + .check([f"executorch_exir_dialects_edge__ops_{gemm}"]) + ) + + def test_fp32_externally_stored_weight_not_folded(self): + tester = Tester(self.WeightNormConv1d().eval(), (torch.randn(1, 8, 16),)) + for node in tester.export().get_artifact().graph.nodes: + if node.op == "placeholder" and node.name.startswith("p_"): + node.meta["custom"] = {"delegate_constant_tag": "weights.ptd"} + tester.to_edge_transform_and_lower().check( + ["executorch_exir_dialects_edge__ops_aten_convolution_default"] + ) diff --git a/backends/xnnpack/utils/utils.py b/backends/xnnpack/utils/utils.py index 2d66ef75aea..6d61f94bead 100644 --- a/backends/xnnpack/utils/utils.py +++ b/backends/xnnpack/utils/utils.py @@ -8,6 +8,11 @@ from typing import Any, cast, List, Optional, Sequence, Tuple import torch +from executorch.backends.xnnpack.utils.quant_utils import ( + is_dequant, + is_qparam, + is_quant, +) from executorch.exir import ExportedProgram from executorch.exir.dialects._ops import ops as exir_ops @@ -27,6 +32,8 @@ PERM_NCHW_TO_NHWC = [0, 2, 3, 1] PERM_NHWC_TO_NCHW = [0, 3, 1, 2] +COMPUTED_PARAM_TAG = "xnn_computed_param" + def check_or_raise(condition: bool, err: str) -> None: """ @@ -164,6 +171,65 @@ def get_param_tensor( raise RuntimeError(f"unsupported param type, {node.op}.") +def get_static_param_deps( + exp_prog: ExportedProgram, node: torch.fx.Node, user: torch.fx.Node +) -> Optional[List[torch.fx.Node]]: + """ + Returns the nodes to partition with user for node to be a static param, or + None. That is node itself if it is a param. Otherwise it is node and the + unmutated params it is computed from with pure ops, such as the weight of a + conv under weight_norm. Preprocess folds node into a param that replaces + them without growing the program, so only node and the params may have + users outside of the computation, and only user, and node may not be larger + than the params it frees. The computing nodes add user's name to their + COMPUTED_PARAM_TAG set. + """ + if is_param_node(exp_prog, node): + return [node] + + signature = exp_prog.graph_signature + mutated_buffers = set(signature.buffers_to_mutate.values()) + deps = {} + pending = [node] + while pending: + dep = pending.pop() + if dep in deps: + continue + deps[dep] = None + if dep.op == "placeholder" and is_param_node(exp_prog, dep): + if ( + signature.inputs_to_buffers.get(dep.name) in mutated_buffers + # Folding would move data tagged for an external file into the PTE. + or dep.meta.get("custom", {}).get("delegate_constant_tag") is not None + ): + return None + continue + if ( + dep.op != "call_function" + or dep.is_impure() + or is_quant(dep) + or is_dequant(dep) + or is_qparam(dep) + ): + return None + pending.extend(dep.all_input_nodes) + + if any( + u not in deps and not (u is user and (dep is node or dep.op == "placeholder")) + for dep in deps + for u in dep.users + ): + return None + removed = [dep for dep in deps if dep.op == "placeholder" and user not in dep.users] + if node.meta["val"].nbytes > sum(param.meta["val"].nbytes for param in removed): + return None + + for dep in deps: + if dep.op == "call_function": + dep.meta.setdefault(COMPUTED_PARAM_TAG, set()).add(user.name) + return list(deps) + + def get_tensor_name(exp_prog: ExportedProgram, node: torch.fx.Node) -> str: if node is None: return ""