[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 ATen 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 graph is not modified.

Fixes #22078

This change was authored with Claude Code.

Author: JakeStevens

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

backends/xnnpack/_passes/fold_computed_params.py

@@ -0,0 +1,82 @@
+# 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,
+    is_param_node,
+)
+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:
+            # A GEMM rejected after its check, possibly by an earlier partitioner,
+            # leaves tags behind, which delegate outputs and inputs inherit. So
+            # fold a tagged node only next to its GEMM, and only from params.
+            if (
+                node.op != "call_function"
+                or not names & node.meta.get(COMPUTED_PARAM_TAG, set())
+                or not all(
+                    n in values or is_param_node(self.exported_program, n)
+                    for n in node.all_input_nodes
+                )
+            ):
+                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)

backends/xnnpack/partition/config/gemm_configs.py

@@ -28,6 +28,7 @@
     tag_as_implicit_q_dq,
 )
 from executorch.backends.xnnpack.utils.utils import (
+    COMPUTED_PARAM_TAG,
     get_input_node,
     is_depthwise_conv,
     is_getitem,
@@ -39,6 +40,7 @@
 )
 from executorch.exir.backend.utils import WhyNoPartition
 from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.dialects.edge._ops import EdgeOpOverload
 from torch.export import ExportedProgram
 from torch.fx.experimental.symbolic_shapes import free_symbols
 from torch.fx.passes.utils.source_matcher_utils import (
@@ -50,6 +52,77 @@
 why = WhyNoPartition(logger=logger)
 
 
+def get_static_param_deps(  # noqa: C901
+    ep: 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 ATen 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(ep, node):
+        return [node]
+
+    signature = ep.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(ep, dep):
+            if signature.inputs_to_buffers.get(dep.name) in mutated_buffers:
+                why(user, f"{dep} is a mutated buffer")
+                return None
+            # Folding would move data tagged for an external file into the PTE.
+            if dep.meta.get("custom", {}).get("delegate_constant_tag") is not None:
+                why(user, f"{dep} is tagged for an external data file")
+                return None
+            continue
+        # Preprocess runs these ops on CPU, and a custom op may only have a fake kernel.
+        if (
+            dep.op != "call_function"
+            or dep.is_impure()
+            or is_quant(dep)
+            or is_dequant(dep)
+            or is_qparam(dep)
+            or not (
+                is_getitem(dep)
+                or (
+                    isinstance(dep.target, (torch._ops.OpOverload, EdgeOpOverload))
+                    and dep.target.namespace in ("aten", "dim_order_ops")
+                )
+            )
+        ):
+            why(user, f"{dep} is neither a param nor an op that can be folded into one")
+            return None
+        pending.extend(dep.all_input_nodes)
+
+    for dep in deps:
+        for u in dep.users:
+            if u not in deps and not (
+                u is user and (dep is node or dep.op == "placeholder")
+            ):
+                why(user, f"{dep} is also used by {u}")
+                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):
+        why(user, f"folding {node} would grow the program")
+        return None
+
+    for dep in deps:
+        if dep.op == "call_function":
+            dep.meta.setdefault(COMPUTED_PARAM_TAG, set()).add(user.name)
+    return list(deps)
+
+
 class GEMMConfig(XNNPartitionerConfig):
     """
     GEMM-like ops like Convolution, Addmm, Linear, mostly behave in the same way, in which we
@@ -153,10 +226,10 @@ 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):
-                why(node, "Expected weight to be a static param")
+            weight_deps = get_static_param_deps(ep, weight_node, node)
+            if weight_deps is None:
                 return (False, [])
-            gemm_deps.append(weight_node)
+            gemm_deps.extend(weight_deps)
 
             return (True, gemm_deps)
         else:
@@ -248,10 +321,10 @@ 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):
-                    why(node, "Expected bias to be a static param")
+                bias_deps = get_static_param_deps(ep, bias_node, node)
+                if bias_deps is None:
                     return (False, [])
-                gemm_deps.append(bias_node)
+                gemm_deps.extend(bias_deps)
 
         return (True, gemm_deps)