[Qualcomm] Scope ExpandBroadcastTensorShape rewrite to the broadcast node (#21583)

# [Qualcomm] Scope `ExpandBroadcastTensorShape` rewrite to the broadcast
node (fix rank-0 mutable-buffer `to_executorch()` failure)

> Draft for discussion with the Qualcomm backend team before opening.
Standalone: one file, +6/-3.

## Summary
`ExpandBroadcastTensorShape` reshapes a lower-rank broadcast input up to
the broadcast op's output
rank so `LayoutTransform` sees equal ranks. It redirected **every** user
of that input to the
rank-promoted view, instead of only the broadcast node it is processing.
This over-broad rewrite
breaks any *other* consumer of the same tensor — most visibly an
in-place mutation of a rank-0 user
input.

## Symptom
A model that mutates a rank-0 tensor **input** in place (a scalar
counter, `counter += N`) fails at
`to_executorch()` on the QNN path:

```
RuntimeError: expand: the requested shape has too few dimensions ...
  at exir/passes/spec_prop_pass.py -> meta_copy_ -> expand_copy
```

XNNPACK / CoreML / generic `to_edge` are unaffected —
`ExpandBroadcastTensorShape` is QNN-only, so
only the QNN lowering produces the rank-mismatched graph.

## Root cause
The scalar counter feeds two things: (1) a broadcast `add` (`arange(N) +
counter`) and (2) the
in-place `counter.add_`. The pass promotes `counter` to `(1,)` for the
broadcast, then redirects
**all** users to the `(1,)` view — including the mutation. ExecuTorch
records the mutation as a
`USER_INPUT_MUTATION`, so at the end it writes the value back into the
input buffer via
`copy_(counter, new_value)`. That write-back is now `copy_(() <- (1,))`;
`SpecPropPass` re-traces it
and `expand_copy` cannot shrink `(1,)` into `()` → the error above.

Confirmed by graph dump: input placeholder `counter` = `()` vs the
`USER_INPUT_MUTATION` output =
`(1,)`; toggling this pass off makes both rank-0 and rank-1 models lower
cleanly.

## Fix
Redirect only the current broadcast node's input, not every user of
`arg`:

```diff
-                            users = list(arg.users.keys())
                             reshape_node = graph_module.graph.create_node( ... )
                             ...
-                            for user in users:
-                                user.replace_input_with(arg, reshape_node)
+                            # Redirect ONLY the current broadcast node's input to the reshaped
+                            # view, not every user of `arg`.
+                            node.replace_input_with(arg, reshape_node)
```

The mutation and the input placeholder both stay rank-0 → write-back
stays `copy_(() <- ())`. This
also fixes a latent correctness bug: blanket replacement is wrong when
one tensor feeds multiple
broadcasts that require different output ranks. HTP's no-rank-0
requirement is still satisfied by the
runtime `dims=[1]` promotion in `builders/node_visitor.py` (blob level),
independent of this
edge-graph rewrite.

## Reproduction
```python
import torch
from executorch.backends.qualcomm.serialization.qc_schema import QcomChipset
from executorch.backends.qualcomm.utils.utils import (
    generate_qnn_executorch_compiler_spec, generate_htp_compiler_spec,
    to_edge_transform_and_lower_to_qnn,
)

class CounterModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.register_buffer("counter", torch.zeros((), dtype=torch.long))  # rank-0
    def forward(self, x):
        self.counter.add_(x.shape[-1])          # in-place mutation of rank-0 input
        return x + self.counter                 # broadcast add -> triggers the pass

m = CounterModel().eval()
ex = (torch.randn(4),)
ep = torch.export.export(m, ex)
backend = generate_htp_compiler_spec(use_fp16=True)
specs = generate_qnn_executorch_compiler_spec(soc_model=QcomChipset.SM8650, backend_options=backend)
edge = to_edge_transform_and_lower_to_qnn(ep, ex, specs)
edge.to_executorch()   # BEFORE: RuntimeError (expand: too few dimensions);  AFTER: OK
```

## Testing
- Repro above: rank-0 `to_executorch()` **FAILED → OK**; rank-1 counter
still OK.
- Stress cases (all pass): one operand feeding two broadcasts of
different rank (each gets its own
correctly-shaped `view_copy`); non-mutated rank-0 broadcast;
tensor-valued increment. Negative control
(revert to the all-users redirect) reproduces the crash → change is
necessary and sufficient.
- **TODO before merge:** graph-level unit test, a real-model QNN
regression export, `lintrunner`.

## Questions for the Qualcomm team
1. Was the all-users rewrite intentional for any case (e.g. a shared
reshape feeding several
broadcasts), or is scoping to the single node always correct? I believe
single-node is correct and
   strictly safer, but want your read.
2. The original all-users redirect also silently rebased
**non-broadcast** consumers (e.g. a `relu`) of
the same operand onto the rank-promoted view — a latent aliasing bug
this scoping also fixes. Confirm
that was never relied upon. (A rank>0 value written back into a rank-0
input is a separate,
   model-level mismatch no operand fix addresses — out of scope here.)
3. Relation to the transformers ExecuTorch-exporter QNN work:
`StaticLayer.cumulative_length` is
exactly this rank-0-counter pattern, so this is a real blocker for that
flow.


cc @cbilgin

Author: psiddh

backends/qualcomm/_passes/expand_broadcast_tensor_shape.py

@@ -27,7 +27,9 @@ def __init__(self):
             exir_ops.edge.aten.expand_copy.default,
         ]
 
-    def traverse_broadcast_node(self, graph_module: torch.fx.GraphModule):
+    def traverse_broadcast_node(
+        self, graph_module: torch.fx.GraphModule, reshape_cache
+    ):
         for node in graph_module.graph.nodes:
             if node.target in self.broadcast_op_targets:
                 for arg in node.args:
@@ -36,31 +38,43 @@ def traverse_broadcast_node(self, graph_module: torch.fx.GraphModule):
                     input_rank = len(arg.meta["val"].shape)
                     output_rank = len(node.meta["val"].shape)
                     if input_rank != output_rank:
-                        with graph_module.graph.inserting_after(arg):
-                            new_rank = [1] * (output_rank - input_rank) + list(
-                                arg.meta["val"].shape
-                            )
-                            users = list(arg.users.keys())
-                            reshape_node = graph_module.graph.create_node(
-                                "call_function",
-                                exir_ops.edge.aten.view_copy.default,
-                                (arg, tuple(new_rank)),
-                            )
-                            # try skip dq_ops to get correct param node if applicable
-                            arg_meta = (
-                                arg.args[0].meta if arg.target in dq_ops else arg.meta
-                            )
-                            # meta needs to be copied elementwisely for fake-tensor
-                            # to be updated correctly and not affect meta of arg
-                            for k, v in arg_meta.items():
-                                reshape_node.meta[k] = v
-                            reshape_node.meta["val"] = reshape_node.meta["val"].reshape(
-                                new_rank
-                            )
-                            for user in users:
-                                user.replace_input_with(arg, reshape_node)
+                        new_rank = [1] * (output_rank - input_rank) + list(
+                            arg.meta["val"].shape
+                        )
+                        # Redirect ONLY the current broadcast node's input to the reshaped
+                        # view, not every user of `arg`. Rewriting all users leaks the rank
+                        # promotion into unrelated consumers (e.g. an in-place mutation of a
+                        # rank-0 user input), producing a rank-mismatched USER_INPUT_MUTATION
+                        # write-back that fails in to_executorch().
+                        # Dedupe reshapes by (arg, new_rank) so that multiple broadcast ops
+                        # sharing the same operand and target rank reuse a single view_copy
+                        # instead of each creating their own.
+                        cache_key = (arg, tuple(new_rank))
+                        reshape_node = reshape_cache.get(cache_key)
+                        if reshape_node is None:
+                            with graph_module.graph.inserting_after(arg):
+                                reshape_node = graph_module.graph.create_node(
+                                    "call_function",
+                                    exir_ops.edge.aten.view_copy.default,
+                                    (arg, tuple(new_rank)),
+                                )
+                                # try skip dq_ops to get correct param node if applicable
+                                arg_meta = (
+                                    arg.args[0].meta
+                                    if arg.target in dq_ops
+                                    else arg.meta
+                                )
+                                # meta needs to be copied elementwisely for fake-tensor
+                                # to be updated correctly and not affect meta of arg
+                                for k, v in arg_meta.items():
+                                    reshape_node.meta[k] = v
+                                reshape_node.meta["val"] = reshape_node.meta[
+                                    "val"
+                                ].reshape(new_rank)
+                            reshape_cache[cache_key] = reshape_node
+                        node.replace_input_with(arg, reshape_node)
 
     def call(self, graph_module: torch.fx.GraphModule):
-        self.traverse_broadcast_node(graph_module)
+        self.traverse_broadcast_node(graph_module, {})
         dead_code_elimination_pass(graph_module)
         return PassResult(graph_module, True)

backends/qualcomm/tests/models.py

@@ -3049,3 +3049,12 @@ def forward(self, x):
         input1 = torch.zeros(1)
         selected_element = torch.select(input1, 0, 0)
         return torch.add(x, selected_element)
+
+
+# rank-0 input `counter` that is both broadcast (arange + counter) and mutated in place;
+# exercises ExpandBroadcastTensorShape's rank promotion and the USER_INPUT_MUTATION write-back.
+class BroadcastAndMutate(torch.nn.Module):
+    def forward(self, x, counter):
+        position = torch.arange(x.shape[-1]) + counter
+        counter.add_(x.shape[-1])
+        return x + position.to(x.dtype)

backends/qualcomm/tests/test_passes.py

@@ -6,6 +6,7 @@
     AnnotateQuantAttrs,
     ConvertBmmToMatmul,
     ConvertMhaToSha,
+    ExpandBroadcastTensorShape,
     FoldQDQ,
     InsertIOQDQ,
     InsertReshapeForReduceOps,
@@ -23,6 +24,7 @@
     QnnExecuTorchBackendType,
 )
 from executorch.backends.qualcomm.tests.models import (
+    BroadcastAndMutate,
     HardSigmoid,
     Reciprocal,
     TopKandIndex,
@@ -35,6 +37,7 @@
 from executorch.exir import EdgeCompileConfig, to_edge
 from executorch.exir.debug_handle_utils import DEBUG_HANDLE_KEY
 from executorch.exir.dialects._ops import ops as exir_ops
+from torch.export.exported_program import OutputKind
 from torch.library import Library
 from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e
 
@@ -470,6 +473,163 @@ def test_decompose_hardsigmoid_backend_aware(self):
                         f"hardsigmoid {'should' if should_decompose else 'should NOT'} be decomposed for {backend.name}",
                     )
 
+    def test_expand_broadcast_preserves_rank0_input_mutation(self):
+        """A rank-0 user input that is BOTH broadcast (needs rank promotion) AND mutated in
+        place must keep its USER_INPUT_MUTATION write-back on the rank-0 value.
+
+        Regression for the all-users rewrite: it redirected *every* consumer of the operand to
+        the promoted (1,) view, including the in-place mutation, so the write-back became
+        copy_(dst=(), src=(1,)) and to_executorch() failed with "expand: too few dimensions".
+        Asserts the broadcast consumes the promoted view while the mutation path stays rank-0.
+        """
+        add = exir_ops.edge.aten.add.Tensor
+        view_copy = exir_ops.edge.aten.view_copy.default
+
+        exported = torch.export.export(
+            BroadcastAndMutate().eval(),
+            (torch.randn(1, 4), torch.tensor(0)),
+            strict=True,
+        )
+        ep = to_edge(exported).exported_program()
+        gm = ExpandBroadcastTensorShape()(ep.graph_module).graph_module
+
+        broadcast = [
+            n for n in gm.graph.nodes if n.target == add and n.meta["val"].dim() == 1
+        ]
+        self.assertTrue(broadcast)
+        self.assertTrue(
+            any(
+                isinstance(a, torch.fx.Node) and a.target == view_copy
+                for n in broadcast
+                for a in n.args
+            ),
+            "the broadcast operand should be rank-promoted to a view_copy",
+        )
+
+        mutated = {
+            spec.arg.name
+            for spec in ep.graph_signature.output_specs
+            if spec.kind == OutputKind.USER_INPUT_MUTATION
+        }
+        self.assertTrue(mutated)
+        checked = 0
+        for node in gm.graph.nodes:
+            if node.name in mutated:
+                checked += 1
+                self.assertFalse(
+                    any(
+                        isinstance(a, torch.fx.Node) and a.target == view_copy
+                        for a in node.args
+                    ),
+                    "rank promotion leaked into the USER_INPUT_MUTATION write-back path",
+                )
+        self.assertEqual(
+            checked,
+            len(mutated),
+            "did not find every mutation-output node in the graph",
+        )
+
+    def test_expand_broadcast_promotes_per_consumer(self):
+        """One tensor feeding two broadcasts of different output rank must be promoted
+        independently per broadcast node, not shared. The all-users rewrite reshaped the
+        operand once (to the first broadcast's rank) and redirected the second broadcast too,
+        mis-shaping it."""
+        add = exir_ops.edge.aten.add.Tensor
+        view_copy = exir_ops.edge.aten.view_copy.default
+
+        class TwoBroadcasts(torch.nn.Module):
+            def forward(self, a2d, a1d, b):
+                return (
+                    a2d + b,
+                    a1d + b,
+                )  # b:(4,) -> rank-2 needs a view; rank-1 does not
+
+        exported = torch.export.export(
+            TwoBroadcasts().eval(),
+            (torch.randn(3, 4), torch.randn(4), torch.randn(4)),
+            strict=True,
+        )
+        ep = to_edge(exported).exported_program()
+        gm = ExpandBroadcastTensorShape()(ep.graph_module).graph_module
+
+        for node in gm.graph.nodes:
+            if node.target != add:
+                continue
+            has_view = any(
+                isinstance(a, torch.fx.Node) and a.target == view_copy
+                for a in node.args
+            )
+            if node.meta["val"].dim() == 2:
+                self.assertTrue(has_view, "rank-2 broadcast operand should be promoted")
+            if node.meta["val"].dim() == 1:
+                self.assertFalse(
+                    has_view, "rank-1 broadcast operand should not be promoted"
+                )
+
+    def test_expand_broadcast_dedupes_shared_same_rank(self):
+        """One operand feeding two broadcasts of the SAME output rank must be promoted with a
+        single shared view_copy, not one per broadcast node.
+
+        Regression for the per-node redirect (PR #21583): redirecting only the current node's
+        input fixed the rank-0 mutation bug but created a separate view_copy for every broadcast
+        consumer, so `a2d + b, a2d - b` (b:(4,)) produced two identical (1,4) reshapes. Deduping
+        by (operand, new_rank) collapses them to one while keeping the rank-0 mutation path (whose
+        operand is never rank-promoted) at rank-0 -- so this coexists with
+        test_expand_broadcast_preserves_rank0_input_mutation.
+        """
+        add = exir_ops.edge.aten.add.Tensor
+        sub = exir_ops.edge.aten.sub.Tensor
+        view_copy = exir_ops.edge.aten.view_copy.default
+
+        class TwoSameRankBroadcasts(torch.nn.Module):
+            def forward(self, a2d, b):
+                return a2d + b, a2d - b  # b:(4,) promoted to (1,4) for both
+
+        exported = torch.export.export(
+            TwoSameRankBroadcasts().eval(),
+            (torch.randn(3, 4), torch.randn(4)),
+            strict=True,
+        )
+        ep = to_edge(exported).exported_program()
+        gm = ExpandBroadcastTensorShape()(ep.graph_module).graph_module
+
+        views = [n for n in gm.graph.nodes if n.target == view_copy]
+        self.assertEqual(
+            len(views), 1, "the shared operand should yield exactly one view_copy"
+        )
+        shared = views[0]
+        for target in (add, sub):
+            broadcast = [n for n in gm.graph.nodes if n.target == target]
+            self.assertTrue(broadcast)
+            self.assertTrue(
+                all(shared in n.args for n in broadcast),
+                "both same-rank broadcasts should consume the shared view_copy",
+            )
+
+        exported = torch.export.export(
+            BroadcastAndMutate().eval(),
+            (torch.randn(1, 4), torch.tensor(0)),
+            strict=True,
+        )
+        ep = to_edge(exported).exported_program()
+        gm = ExpandBroadcastTensorShape()(ep.graph_module).graph_module
+
+        mutated = {
+            spec.arg.name
+            for spec in ep.graph_signature.output_specs
+            if spec.kind == OutputKind.USER_INPUT_MUTATION
+        }
+        self.assertTrue(mutated)
+        for node in gm.graph.nodes:
+            if node.name in mutated:
+                self.assertFalse(
+                    any(
+                        isinstance(a, torch.fx.Node) and a.target == view_copy
+                        for a in node.args
+                    ),
+                    "dedupe must not rank-promote the USER_INPUT_MUTATION write-back",
+                )
+
 
 if __name__ == "__main__":
     unittest.main()