Add skeleton code

Author: larryliu0820

CMakeLists.txt

@@ -603,6 +603,12 @@ endif()
 
 if(EXECUTORCH_BUILD_CORTEX_M)
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backends/cortex_m)
+  list(APPEND _executorch_backends coretex_m_backend)
+endif()
+
+if(EXECUTORCH_BUILD_AOTI)
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backends/aoti)
+  list(APPEND _executorch_backends aoti_backend)
 endif()
 
 if(EXECUTORCH_BUILD_DEVTOOLS)

backends/aoti/CMakeLists.txt

@@ -0,0 +1,47 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+#
+# Build AOTI backend for runtime.
+#
+# ### Editing this file ###
+#
+# This file should be formatted with
+# ~~~
+# cmake-format -i CMakeLists.txt
+# ~~~
+# It should also be cmake-lint clean.
+#
+set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
+
+# Source root directory for executorch.
+if(NOT EXECUTORCH_ROOT)
+  set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../..)
+endif()
+
+include(${EXECUTORCH_ROOT}/build/Utils.cmake)
+
+find_package(CUDA)
+
+set(_common_include_directories ${EXECUTORCH_ROOT}/..)
+
+set(_aoti_sources runtime/AotiBackend.cpp)
+
+add_library(aoti_backend STATIC ${_aoti_sources})
+target_include_directories(
+  aoti_backend PUBLIC ${_common_include_directories} ${CUDA_INCLUDE_DIRS}
+)
+
+target_compile_options(aoti_backend PUBLIC -fexceptions -frtti -fPIC)
+target_link_libraries(aoti_backend PUBLIC extension_tensor ${CUDA_LIBRARIES})
+executorch_target_link_options_shared_lib(aoti_backend)
+
+install(
+  TARGETS aoti_backend
+  EXPORT ExecuTorchTargets
+  DESTINATION lib
+  INCLUDES
+  DESTINATION ${_common_include_directories}
+)

backends/aoti/README.md

@@ -0,0 +1,2 @@
+## Experimental AOTI backend
+Proceed with caution. This is an experimental backend that is not yet ready for production use.

backends/aoti/aoti_backend.py

@@ -0,0 +1,43 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import copy
+
+from subprocess import check_call
+from typing import final, List
+
+import torch
+from executorch.exir.backend.backend_details import (
+    BackendDetails,
+    ExportedProgram,
+    PreprocessResult,
+)
+from executorch.exir.backend.compile_spec_schema import CompileSpec
+
+
+@final
+class AotiBackend(BackendDetails):
+    @staticmethod
+    def preprocess(
+        edge_program: ExportedProgram,
+        compile_specs: List[CompileSpec],
+    ) -> PreprocessResult:
+        print("entering  the lowerable parts in AotiBackend.preprocess....")
+
+        print("here", edge_program.example_inputs)
+        copy_edge_program = copy.deepcopy(edge_program)
+        graph_module = copy_edge_program.graph_module
+        args, kwargs = copy_edge_program.example_inputs
+        so_path = torch._inductor.aot_compile(graph_module, args, kwargs, options={})  # type: ignore[arg-type]
+        print(so_path)
+        check_call(
+            f"patchelf --remove-needed libtorch.so --remove-needed libtorch_cuda.so --remove-needed libc10_cuda.so --remove-needed libtorch_cpu.so --add-needed libcudart.so {so_path}",
+            shell=True,
+        )
+
+        with open(so_path, "rb") as f:
+            data = f.read()
+        return PreprocessResult(data)

backends/aoti/aoti_partitioner.py

@@ -0,0 +1,74 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-unsafe
+
+from typing import cast, final, List
+
+import torch
+from executorch.backends.aoti.aoti_backend import AotiBackend  # usort: skip
+from executorch.exir.backend.compile_spec_schema import CompileSpec
+from executorch.exir.backend.partitioner import (
+    DelegationSpec,
+    Partitioner,
+    PartitionResult,
+)
+from executorch.exir.backend.utils import tag_constant_data
+from executorch.exir.dialects._ops import ops as exir_ops
+from torch.export.exported_program import ExportedProgram
+from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner
+
+from torch.fx.passes.operator_support import OperatorSupportBase
+
+
+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.aten._to_copy.default,
+        ]
+
+        return supported
+
+    def is_node_supported_custom(self, node: torch.fx.Node) -> bool:
+        if node.target == exir_ops.edge.aten.mean.dim:
+            keep_dim = node.args[2] if len(node.args) > 2 else False
+            return cast(bool, keep_dim)
+        if node.target == exir_ops.edge.aten.var.correction:
+            keep_dim = node.kwargs.get("keepdim", False)
+            return cast(bool, keep_dim)
+        return True
+
+
+@final
+class AotiPartitioner(Partitioner):
+    def __init__(self, compile_spec: List[CompileSpec]) -> None:
+        self.delegation_spec = DelegationSpec(AotiBackend.__name__, compile_spec)
+        print(self.delegation_spec)
+
+    def partition(self, exported_program: ExportedProgram) -> PartitionResult:
+        # Run the CapabilityBasedPartitioner to return the largest possible
+        # subgraphs containing the nodes with the tags
+        # logger.info("AotiPartitioner::partition")
+        partition_tags = {}
+
+        capability_partitioner = CapabilityBasedPartitioner(
+            exported_program.graph_module,
+            AOTISupportedOperators(),
+            allows_single_node_partition=True,
+        )
+        partition_list = capability_partitioner.propose_partitions()
+        for partition in partition_list:
+            for node in partition.nodes:
+                tag = f"tag{partition.id}"
+                node.meta["delegation_tag"] = tag
+                partition_tags[tag] = self.delegation_spec
+
+        tag_constant_data(exported_program)
+
+        return PartitionResult(
+            tagged_exported_program=exported_program, partition_tags=partition_tags
+        )