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