Erase the cat/slice/select nop ops after memory planning (pytorch#22651)

Summary:

`_cat_nop`, `_slice_copy_nop` and `_select_copy_nop` do nothing at runtime. A
cat is only rewritten to its nop form once memory planning can place every
input at a contiguous offset inside the output, and a slice or select likewise
once its output can be colocated inside its input, so by the time planning has
run they hold by construction. The instructions that remain are pure dispatch
overhead.

This erases them, pointing their consumers at the output buffer the producers
have already written in place. A nop that cannot be erased raises rather than
being skipped: these ops exist only to carry a placement constraint from
constraint generation to here, so one reaching the emitter is a compiler bug.

**Why it has to live inside the memory planning pass.** A standalone pass is
not possible:

- Before planning the nodes cannot go, because the node is what carries the
  placement constraint. Remove it early and the aliasing never happens.
- After planning nothing may run at all:

    # WARNING: DO NOT ADD ANY MORE PASSES AFTER MEMORY PLANNING PASS.
    # THERE ARE A LOT OF ASSUMPTIONS IN THE STACK THAT MEMORY PLANNING IS
    # THE LAST PASS BEFORE THE EMITTER.
    (exir/program/_program.py)

`CadenceMemoryPlanning.run` already mutates the graph after `mem_planning.run`
- `SimplifyIdmaOpsPass` retargets nodes and runs dead code elimination there -
so that slot is the one point where placement is decided but the program is not
yet emitted. This joins it.

**The assumptions that warning refers to are real.** Spec lifetimes are node
indices, so erasing nodes leaves them pointing past the end of the graph. That
trips `find_peak_memory_usage` and, in executorch/util, the activation memory
profiler. `update_all_tensors_lifetime` alone does not fix it, because
`update_tensor_lifetime` only ever widens a lifetime:

    end = node_idx if end is None or end < node_idx else end

so a stale larger bound survives. `_refresh_lifetimes` uses the first call to
identify exactly which specs the recompute reaches, clears those, and rebuilds
them. Clearing a wider set would leave specs the recompute never revisits stuck
at None, which reads as "no lifetime" and silently drops them from the memory
reports. With that, no change is needed outside the Cadence backend - the shared
executorch diagnostics work unmodified.

The nop targets are looked up inside `call` rather than at class scope, because
their schemas come from `ops_registrations` and `memory_planning` does not
import it - resolving at import time breaks any module that imports
`CadenceMemoryPlanning` without having registered the ops first.

**The kernels stay, deliberately.**

Reviewed By: DrJessop

Differential Revision: D118922409

Pull Request resolved: pytorch#22651

Author: mcremon-meta

backends/cadence/aot/memory_planning.py

@@ -26,7 +26,11 @@
 )
 
 from executorch.exir import ExecutorchProgramManager
-from executorch.exir.memory_planning import collect_specs_from_nodes, Verifier
+from executorch.exir.memory_planning import (
+    collect_specs_from_nodes,
+    update_all_tensors_lifetime,
+    Verifier,
+)
 from executorch.exir.pass_base import PassBase
 from executorch.exir.pass_manager import PassManager
 from executorch.exir.passes import MemoryPlanningPass
@@ -388,6 +392,89 @@ def call(self, graph_module: torch.fx.GraphModule) -> Optional[PassResult]:
         return PassResult(graph_module, modified)
 
 
+class RemoveNopOpsPass(PassBase):
+    """Erase the cat/slice/select nop nodes once memory planning is done.
+
+    These ops never touch data. A cat is only rewritten to its nop form once
+    memory planning can place every input at a contiguous offset inside the
+    output, and a slice likewise once its output can be colocated inside its
+    input, so by the time planning has run both hold by construction and the
+    instructions are pure dispatch overhead.
+
+    This has to happen here rather than as its own pass. Before planning the
+    nodes cannot go, because they are what carries the placement constraint;
+    after it nothing may run at all (see the warning in
+    exir/program/_program.py). Running inside the memory planning pass, in the
+    slot SimplifyIdmaOpsPass already occupies, is the one point where the
+    placement is decided but the program is not yet emitted.
+    """
+
+    def __init__(self, graph_signature: Optional[ExportGraphSignature] = None) -> None:
+        self.graph_signature = graph_signature
+
+    def call(self, graph_module: torch.fx.GraphModule) -> Optional[PassResult]:
+        # Every nop target memory_constraints.py can produce. Keep in sync
+        # with compute_cat_contiguity_constraints and
+        # compute_slice_and_select_loc_constraints. Looked up here rather than
+        # at class scope because the schemas come from ops_registrations, which
+        # this module does not import.
+        targets = (
+            torch.ops.aten._cat_nop.out,
+            torch.ops.aten._slice_copy_nop.Tensor_out,
+            torch.ops.aten._select_copy_nop.int_out,
+        )
+        modified = False
+        for target in targets:
+            for node in graph_module.graph.find_nodes(
+                op="call_function", target=target
+            ):
+                out = node.kwargs.get("out")
+                if out is None:
+                    # These ops exist only to carry a placement constraint
+                    # between constraint generation and here. One that cannot
+                    # be erased would reach the runtime, where there is no
+                    # kernel to service it, so fail the build instead.
+                    raise RuntimeError(
+                        f"cannot erase {node.name} ({node.target}): no out "
+                        "kwarg to redirect its consumers to. A nop op must "
+                        "never reach the emitter."
+                    )
+                # Consumers read the output buffer, which the producers have
+                # already written in place. Point them straight at it.
+                node.replace_all_uses_with(out)
+                graph_module.graph.erase_node(node)
+                modified = True
+
+        if not modified:
+            return PassResult(graph_module, False)
+
+        graph_module.recompile()
+        self._refresh_lifetimes(graph_module)
+        return PassResult(graph_module, True)
+
+    def _refresh_lifetimes(self, graph_module: torch.fx.GraphModule) -> None:
+        """Recompute spec lifetimes against the graph we actually emit.
+
+        Lifetimes are node indices, so erasing nodes leaves them pointing past
+        the end of the graph, which trips the peak-memory reporter and the
+        activation profiler.
+
+        update_tensor_lifetime only ever widens a lifetime, so a single
+        recompute leaves the stale upper bounds in place. The first call
+        identifies exactly which specs the recompute reaches (it returns
+        them); those are cleared and the second call then rebuilds them from
+        scratch. Clearing a wider set than that would leave specs the
+        recompute never revisits stuck at None, which reads as "no lifetime"
+        and silently drops them from the memory reports.
+
+        Placement is already decided at this point - only lifetimes change.
+        """
+        specs = update_all_tensors_lifetime(graph_module, self.graph_signature)
+        for spec in specs:
+            spec.lifetime = [None, None]
+        update_all_tensors_lifetime(graph_module, self.graph_signature)
+
+
 ConstraintGenPassType: TypeAlias = Callable[
     [MemConstraints],
     Callable[[torch.fx.GraphModule], Optional[PassResult]],
@@ -465,8 +552,11 @@ def run(
         )
         mem_planning.run(graph_module, graph_signature)
 
-        graph_module = PassManager(passes=[SimplifyIdmaOpsPass()])(
-            graph_module
-        ).graph_module
+        graph_module = PassManager(
+            passes=[
+                SimplifyIdmaOpsPass(),
+                RemoveNopOpsPass(graph_signature),
+            ]
+        )(graph_module).graph_module
 
         return PassResult(graph_module, True)