Add get_delegated_payload utility function (pytorch#11651) (pytorch#16852)

## Summary
Add `get_delegated_payload()` utility function to extract delegate
payloads from a graph module. The function returns a dictionary mapping
delegate names to tuples of (backend_id, compile_specs,
processed_bytes).

Fixes pytorch#11651

## Test plan
Added 3 unit tests:

```bash
python -m unittest exir.backend.test.test_utils.TestUtils.test_get_delegated_payload_with_delegates -v
python -m unittest exir.backend.test.test_utils.TestUtils.test_get_delegated_payload_without_delegates -v
python -m unittest exir.backend.test.test_utils.TestUtils.test_get_delegated_payload_keys_match_delegates -v

Author: ahmet-f-gumustas

exir/backend/test/test_utils.py

@@ -14,6 +14,7 @@
 from executorch.exir.backend.test.op_partitioner_demo import AddMulPartitionerDemo
 from executorch.exir.backend.utils import (
     format_delegated_graph,
+    get_delegated_payload,
     get_delegates,
     get_non_lowered_nodes,
     is_identical_graph,
@@ -262,3 +263,130 @@ def forward(self, a, x, b):
             graph_str,
             "Expect to see the aten.mm in the delegated graph",
         )
+
+    def test_get_delegated_payload_with_delegates(self):
+        """Test get_delegated_payload returns correct payload for delegated modules."""
+
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, a, x, b):
+                y = torch.mm(a, x)
+                z = y + b
+                a = z - a
+                y = torch.mm(a, x)
+                z = y + b
+                return z
+
+        m = Model()
+        inputs = (torch.randn(2, 2), torch.randn(2, 2), torch.randn(2, 2))
+
+        edge = to_edge(export(m, inputs, strict=True)).to_backend(
+            AddMulPartitionerDemo()
+        )
+
+        payloads = get_delegated_payload(edge.exported_program().graph_module)
+
+        # Should have 2 delegates: (mm + add) -> sub -> (mm + add)
+        self.assertEqual(len(payloads), 2)
+
+        # Verify payload structure for each delegate
+        for name, (backend_id, compile_specs, processed_bytes) in payloads.items():
+            # Check delegate name format
+            self.assertTrue(
+                name.startswith("lowered_module_"),
+                f"Delegate name should start with 'lowered_module_', got {name}",
+            )
+
+            # Check backend_id
+            self.assertEqual(
+                backend_id,
+                "BackendWithCompilerDemo",
+                f"Expected backend_id 'BackendWithCompilerDemo', got {backend_id}",
+            )
+
+            # Check compile_specs is a list
+            self.assertIsInstance(
+                compile_specs,
+                list,
+                f"compile_specs should be a list, got {type(compile_specs)}",
+            )
+
+            # Check processed_bytes is bytes
+            self.assertIsInstance(
+                processed_bytes,
+                bytes,
+                f"processed_bytes should be bytes, got {type(processed_bytes)}",
+            )
+
+            # Verify processed_bytes is not empty (backend should produce some output)
+            self.assertGreater(
+                len(processed_bytes),
+                0,
+                "processed_bytes should not be empty",
+            )
+
+    def test_get_delegated_payload_without_delegates(self):
+        """Test get_delegated_payload returns empty dict when no delegates present."""
+
+        class SimpleModel(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, x):
+                return x + 1
+
+        m = SimpleModel()
+        inputs = (torch.randn(2, 2),)
+
+        # Create edge program without delegation
+        edge = to_edge(export(m, inputs, strict=True))
+
+        payloads = get_delegated_payload(edge.exported_program().graph_module)
+
+        # Should have no delegates
+        self.assertEqual(
+            len(payloads),
+            0,
+            "Expected empty payload dict when no delegates present",
+        )
+
+    def test_get_delegated_payload_keys_match_delegates(self):
+        """Test that get_delegated_payload keys match get_delegates node names."""
+
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, a, x, b):
+                y = torch.mm(a, x)
+                z = y + b
+                a = z - a
+                y = torch.mm(a, x)
+                z = y + b
+                return z
+
+        m = Model()
+        inputs = (torch.randn(2, 2), torch.randn(2, 2), torch.randn(2, 2))
+
+        edge = to_edge(export(m, inputs, strict=True)).to_backend(
+            AddMulPartitionerDemo()
+        )
+
+        graph_module = edge.exported_program().graph_module
+
+        # Get delegates using existing utility
+        delegate_nodes = get_delegates(graph_module.graph)
+        delegate_names = {node.name for node in delegate_nodes}
+
+        # Get payloads using new utility
+        payloads = get_delegated_payload(graph_module)
+        payload_names = set(payloads.keys())
+
+        # Names should match
+        self.assertEqual(
+            delegate_names,
+            payload_names,
+            f"Delegate names mismatch: {delegate_names} vs {payload_names}",
+        )

exir/backend/utils.py

@@ -11,7 +11,7 @@
 import operator
 from collections import defaultdict, OrderedDict
 from functools import lru_cache
-from typing import Dict, Iterable, List, Optional, Set, Tuple, Union
+from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Union
 
 import torch
 from executorch.exir.backend.backend_details import ExportedProgram
@@ -659,3 +659,53 @@ def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
         if not is_supported:
             self.reporter.report_reject(node, self.message)
         return is_supported
+
+
+def get_delegated_payload(
+    graph_module: torch.fx.GraphModule,
+) -> Dict[str, Tuple[str, List[Any], bytes]]:
+    """
+    Extracts the payload for delegates from a graph module that has been lowered
+    to one or more backends.
+
+    This function iterates through all lowered modules (delegates) in the graph
+    and returns a dictionary mapping each delegate's name to a tuple containing:
+    - backend_id: The name/identifier of the backend
+    - compile_specs: A list of backend-specific compilation specifications
+    - processed_bytes: The delegate blob created by the backend's preprocess method
+
+    Args:
+        graph_module: A torch.fx.GraphModule that may contain lowered backend modules.
+            This is typically obtained from an EdgeProgramManager or ExecutorchProgram
+            via `.exported_program().graph_module`.
+
+    Returns:
+        Dict[str, Tuple[str, List[Any], bytes]]: A dictionary where:
+            - Keys are the delegate names (e.g., "lowered_module_0", "lowered_module_1")
+            - Values are tuples of (backend_id, compile_specs, processed_bytes)
+
+    Example:
+        >>> edge = to_edge(export(model, inputs))
+        >>> edge = edge.to_backend(MyPartitioner())
+        >>> payloads = get_delegated_payload(edge.exported_program().graph_module)
+        >>> for name, (backend_id, specs, data) in payloads.items():
+        ...     print(f"{name}: backend={backend_id}, data_size={len(data)}")
+    """
+    from executorch.exir.lowered_backend_module import LoweredBackendModule
+
+    delegate_payloads: Dict[str, Tuple[str, List[Any], bytes]] = {}
+
+    # Find all lowered modules in the graph
+    for node in graph_module.graph.nodes:
+        if node.op == "get_attr" and node.name.startswith("lowered_module_"):
+            lowered_module = getattr(graph_module, node.name, None)
+            if lowered_module is not None and isinstance(
+                lowered_module, LoweredBackendModule
+            ):
+                delegate_payloads[node.name] = (
+                    lowered_module.backend_id,
+                    lowered_module.compile_specs,
+                    lowered_module.processed_bytes,
+                )
+
+    return delegate_payloads