support latest backend interface

Author: Gasoonjia

backends/aoti/aoti_backend.py

@@ -5,6 +5,8 @@
 # LICENSE file in the root directory of this source tree.
 
 import copy
+import os
+import shutil
 
 from subprocess import check_call
 from typing import final, List
@@ -16,8 +18,7 @@
     PreprocessResult,
 )
 from executorch.exir.backend.compile_spec_schema import CompileSpec
-import os
-import shutil
+
 
 @final
 class AotiBackend(BackendDetails):
@@ -33,7 +34,7 @@ def preprocess(
         graph_module = copy_edge_program.graph_module
         args, kwargs = copy_edge_program.example_inputs
         temp_so_path = torch._inductor.aot_compile(graph_module, args, kwargs, options={})  # type: ignore[arg-type]
-        so_path = os.path.join(os.getcwd(), 'aoti.so')
+        so_path = os.path.join(os.getcwd(), "aoti.so")
         print("so_path after aot_compile: ", temp_so_path)
         print("so path we will using ", so_path)
         shutil.copyfile(temp_so_path, so_path)

backends/aoti/aoti_partitioner.py

@@ -23,13 +23,15 @@
 
 from torch.fx.passes.operator_support import OperatorSupportBase
 
+supported_fallback_operators = []
+
 
 class AOTISupportedOperators(OperatorSupportBase):
     def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
-        supported = node.op == "call_function" and node.target in [
-            exir_ops.edge.aten.add.Tensor,
-            exir_ops.edge.dim_order_ops._to_dim_order_copy.default,
-        ]
+        supported = (
+            node.op == "call_function"
+            and node.target not in supported_fallback_operators
+        )
 
         return supported
 

backends/aoti/runtime/AotiBackend.cpp

@@ -50,6 +50,7 @@ using executorch::runtime::FreeableBuffer;
 using executorch::runtime::MemoryAllocator;
 using executorch::runtime::Result;
 using executorch::runtime::etensor::Tensor;
+using executorch::runtime::Span;
 
 extern "C" {
 using AOTITensorHandle = Tensor*;
@@ -494,7 +495,7 @@ class AOTIBackend final : public ::executorch::runtime::BackendInterface {
   Error execute(
       BackendExecutionContext& context,
       DelegateHandle* handle_,
-      EValue** args) const override {
+      Span<EValue*> args) const override {
     AOTIDelegateHandle* handle = (AOTIDelegateHandle*)handle_;
 
     size_t num_inputs;

export_add.py

@@ -27,5 +27,5 @@ def forward(self, x: torch.Tensor, y: torch.Tensor):
 executorch_program = edge_program.to_executorch()
 
 # 4. Save the compiled .pte program
-with open("add.pte", "wb") as file:
+with open("aoti_model.pte", "wb") as file:
     file.write(executorch_program.buffer)

export_and_run_aoti.sh

@@ -1,5 +1,5 @@
 ./install_executorch.sh
-python export_add.py
+python $1
 ./install_executorch.sh --clean
 mkdir -p cmake-out
 cd cmake-out
@@ -8,4 +8,4 @@ cmake -DEXECUTORCH_BUILD_AOTI=ON \
       ..
 cd ..
 cmake --build cmake-out -j9
-./cmake-out/executor_runner --model_path add.pte
+./cmake-out/executor_runner --model_path aoti_model.pte

export_mv2.py

@@ -0,0 +1,28 @@
+import torch
+from executorch.backends.aoti.aoti_partitioner import AotiPartitioner
+from executorch.examples.models.mobilenet_v2 import MV2Model
+from executorch.exir import to_edge
+from torch.export import export
+from torchvision import models
+from torchvision.models.mobilenetv2 import MobileNet_V2_Weights
+
+mv2 = models.mobilenetv2.mobilenet_v2(weights=MobileNet_V2_Weights)
+mv2 = mv2.eval()
+
+model_inputs = (torch.randn(1, 3, 224, 224),)
+
+
+# 1. torch.export: Defines the program with the ATen operator set.
+aten_dialect = export(mv2, model_inputs)
+
+# 2. to_edge: Make optimizations for Edge devices
+edge_program = to_edge(aten_dialect)
+
+edge_program = edge_program.to_backend(AotiPartitioner([]))
+
+# 3. to_executorch: Convert the graph to an ExecuTorch program
+executorch_program = edge_program.to_executorch()
+
+# 4. Save the compiled .pte program
+with open("aoti_model.pte", "wb") as file:
+    file.write(executorch_program.buffer)