Update on "migrate etrecord generation after to_edge_transform_and_lower to new infra"

Differential Revision: [D79420502](https://our.internmc.facebook.com/intern/diff/D79420502/)

[ghstack-poisoned]

Author: Gasoonjia

backends/apple/coreml/runtime/delegate/multiarray.mm

@@ -123,6 +123,9 @@ bool init_bnns_descriptor(BNNSNDArrayDescriptor& bnns_descriptor, const MultiArr
 }
 
 bool copy_using_bnns(const MultiArray& src, MultiArray& dst) {
+    if (src.layout().dataType() != dst.layout().dataType()) {
+        return false;
+    }
     if (dst.layout().num_bytes() < src.layout().num_bytes()) {
         return false;
     }

backends/xnnpack/recipes/xnnpack_recipe_provider.py

@@ -27,6 +27,7 @@
 from executorch.export import (
     BackendRecipeProvider,
     ExportRecipe,
+    LoweringRecipe,
     QuantizationRecipe,
     RecipeType,
 )
@@ -88,12 +89,19 @@ def create_recipe(
             )
         return None
 
+    def _get_xnnpack_lowering_recipe(
+        self, precision_type: Optional[ConfigPrecisionType] = None
+    ) -> LoweringRecipe:
+        return LoweringRecipe(
+            partitioners=[XnnpackPartitioner(precision_type=precision_type)],
+            edge_compile_config=get_xnnpack_edge_compile_config(),
+        )
+
     def _build_fp32_recipe(self, recipe_type: RecipeType) -> ExportRecipe:
         return ExportRecipe(
             name=recipe_type.value,
-            edge_compile_config=get_xnnpack_edge_compile_config(),
+            lowering_recipe=self._get_xnnpack_lowering_recipe(),
             executorch_backend_config=get_xnnpack_executorch_backend_config(),
-            partitioners=[XnnpackPartitioner()],
         )
 
     def _build_quantized_recipe(
@@ -120,9 +128,8 @@ def _build_quantized_recipe(
         return ExportRecipe(
             name=recipe_type.value,
             quantization_recipe=quant_recipe,
-            edge_compile_config=get_xnnpack_edge_compile_config(),
+            lowering_recipe=self._get_xnnpack_lowering_recipe(precision_type),
             executorch_backend_config=get_xnnpack_executorch_backend_config(),
-            partitioners=[XnnpackPartitioner(config_precision=precision_type)],
         )
 
     def _build_int8da_intx_weight_recipe(
@@ -150,9 +157,8 @@ def _build_int8da_intx_weight_recipe(
         return ExportRecipe(
             name=recipe_type.value,
             quantization_recipe=quant_recipe,
-            edge_compile_config=get_xnnpack_edge_compile_config(),
+            lowering_recipe=self._get_xnnpack_lowering_recipe(),
             executorch_backend_config=get_xnnpack_executorch_backend_config(),
-            partitioners=[XnnpackPartitioner()],
         )
 
     def _validate_recipe_kwargs(self, recipe_type: RecipeType, **kwargs: Any) -> None:

examples/xnnpack/aot_compiler.py

@@ -101,15 +101,15 @@
             _check_ir_validity=False if args.quantize else True,
             _skip_dim_order=True,  # TODO(T182187531): enable dim order in xnnpack
         ),
-        generate_etrecord=args.etrecord is not None,
+        generate_etrecord=args.etrecord is not False,
     )
     logging.info(f"Exported and lowered graph:\n{edge.exported_program().graph}")
 
     exec_prog = edge.to_executorch(
         config=ExecutorchBackendConfig(extract_delegate_segments=False)
     )
 
-    if args.etrecord is not None:
+    if args.etrecord is not False:
         edge.get_etrecord().save(args.etrecord)
         logging.info(f"Saved ETRecord to {args.etrecord}")
 

export/TARGETS

@@ -15,7 +15,6 @@ runtime.python_library(
         "//caffe2:torch",
         "//executorch/exir/backend:backend_api",
         "//executorch/exir:pass_manager",
-        "//executorch/devtools/backend_debug:delegation_info",
         "//executorch/extension/export_util:export_util",
     ]
 )
@@ -31,11 +30,35 @@ runtime.python_library(
     ],
     deps = [
         ":recipe",
+        ":stages",
+        ":types",
         "//executorch/runtime:runtime",
         ":recipe_registry"
     ]
 )
 
+
+runtime.python_library(
+    name = "stages",
+    srcs = [
+        "stages.py",
+    ],
+    visibility = [
+        "//executorch/...",
+        "@EXECUTORCH_CLIENTS",
+    ],
+    deps = [
+        ":recipe",
+        ":types",
+        "//executorch/devtools/backend_debug:delegation_info",
+        "//executorch/exir/backend:backend_api",
+        "//executorch/exir:pass_manager",
+        "//caffe2:torch",
+        "//executorch/devtools/backend_debug:delegation_info",
+    ]
+)
+
+
 runtime.python_library(
     name = "lib",
     srcs = [
@@ -48,8 +71,10 @@ runtime.python_library(
     deps = [
         ":export",
         ":recipe",
+        ":stages",
         ":recipe_registry",
-        ":recipe_provider"
+        ":recipe_provider",
+        ":types",
     ],
 )
 
@@ -78,3 +103,10 @@ runtime.python_library(
         ":recipe",
     ]
 )
+
+runtime.python_library(
+    name = "types",
+    srcs = [
+        "types.py",
+    ],
+)

export/__init__.py

@@ -15,13 +15,15 @@
 """
 
 from .export import export, ExportSession
-from .recipe import ExportRecipe, QuantizationRecipe, RecipeType
+from .recipe import ExportRecipe, LoweringRecipe, QuantizationRecipe, RecipeType
 from .recipe_provider import BackendRecipeProvider
 from .recipe_registry import recipe_registry
-
+from .types import StageType
 
 __all__ = [
+    "StageType",
     "ExportRecipe",
+    "LoweringRecipe",
     "QuantizationRecipe",
     "ExportSession",
     "export",