let backends set info on the lowered_backend.meta from preprocess (pytorch#13856)

Summary:

We want some delegates to be able to test experimental features that interact with the core runtime. Typically when doing this there is an expectation of tight coupling in the full lowering process so we arent super concerned about UX. The backend api is pretty hard to change because the code is mostly owned by 3rd parties, so we want to be really sure when we change it that its correct. By exposing this module.meta hookup we can test things in a controled maner without changing the blessed api before we are ready.

Reviewed By: hsharma35

Differential Revision: D81466391

Author: JacobSzwejbka

exir/backend/backend_api.py

@@ -126,6 +126,8 @@ def to_backend(
             lowered_module.meta = {
                 "debug_handle_map": preprocess_result.debug_handle_map
             }
+            if preprocess_result.delegate_info is not None:
+                lowered_module.meta["delegate_info"] = preprocess_result.delegate_info
             return lowered_module
     raise NotImplementedError(f"Backend {backend_id} was not found.")
 
@@ -610,6 +612,8 @@ def lower_all_submodules_to_backend(
             lowered_module.meta = {
                 "debug_handle_map": preprocess_result.debug_handle_map,
             }
+            if preprocess_result.delegate_info is not None:
+                lowered_module.meta["delegate_info"] = preprocess_result.delegate_info
             is_submodule = call_submodule_node.meta["is_submodule"]
             toplevel_input_specs_to_delete = call_submodule_node.meta[
                 "toplevel_input_specs_to_delete"

exir/backend/backend_details.py

@@ -7,7 +7,7 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 from executorch.exir._serialize._named_data_store import NamedDataStoreOutput
 
@@ -32,6 +32,9 @@ class PreprocessResult:
     # but retrieveable by delegates via the NamedDataMap at runtime.
     data_store_output: Optional[NamedDataStoreOutput] = None
 
+    # Optional delegate-specific information that will be added to the lowered_module's metadata.
+    delegate_info: Optional[Any] = None
+
 
 """
 How to create a backend (for example, BackendWithCompilerDemo):

exir/backend/test/backend_with_compiler_demo.py

@@ -138,4 +138,5 @@ def preprocess(
                 encoding="utf8",
             ),
             debug_handle_map=debug_handle_map,
+            delegate_info="test",
         )

exir/backend/test/test_backends_lifted.py

@@ -1264,3 +1264,63 @@ def forward(self, x: List[torch.Tensor]):
 
         gm = to_edge(export(ComposedM(), inputs, strict=True))
         gm.exported_program().module()(*inputs)
+
+    def test_delegate_info_full_delegate(self):
+        """
+        Test that delegate_info from BackendWithCompilerDemo ends up in the call_delegate node metadata
+        when using full delegation (to_backend directly).
+        """
+
+        class SinModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, x):
+                return torch.sin(x)
+
+        sin_module = SinModule()
+        model_inputs = (torch.ones(1),)
+        edgeir_m = to_edge(export(sin_module, model_inputs, strict=True))
+        max_value = model_inputs[0].shape[0]
+        compile_specs = [CompileSpec("max_value", bytes([max_value]))]
+        lowered_sin_module = to_backend(
+            "BackendWithCompilerDemo", edgeir_m.exported_program(), compile_specs
+        )
+
+        # Check that the lowered module has delegate_info in its meta
+        self.assertIn("delegate_info", lowered_sin_module.meta.keys())
+        self.assertEqual(lowered_sin_module.meta["delegate_info"], "test")
+
+    def test_delegate_info_partitioner(self):
+        """
+        Test that delegate_info from BackendWithCompilerDemo ends up in the call_delegate node metadata
+        when using partitioner-based delegation.
+        """
+
+        class SinModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, x):
+                return torch.sin(x)
+
+        sin_module = SinModule()
+        model_inputs = (torch.ones(1),)
+        max_value = model_inputs[0].shape[0]
+
+        partitioner = AllNodePartitioner(
+            "BackendWithCompilerDemo", [CompileSpec("max_value", bytes([max_value]))]
+        )
+
+        edgeir_m = to_edge(export(sin_module, model_inputs, strict=True))
+        lowered_m = edgeir_m.to_backend(partitioner)
+
+        # Check that the lowered submodule has delegate_info in its meta
+        lowered_submodules = get_lowered_submodules(
+            lowered_m.exported_program().graph_module
+        )
+        self.assertEqual(len(lowered_submodules), 1)
+          
+        lowered_module = lowered_submodules[0][1]
+        self.assertIn("delegate_info", lowered_module.meta)
+        self.assertEqual(lowered_module.meta["delegate_info"], "test")