remove unnecessary export code

Author: Gasoonjia

backends/aoti/aoti_backend.py

@@ -73,31 +73,15 @@ def preprocess(
         compile_specs: List[CompileSpec],
     ) -> PreprocessResult:
 
-        print("entering  the lowerable parts in AotiBackend.preprocess....")
         named_data_store = NamedDataStore()
 
-        # print("here", edge_program.example_inputs)
-        copy_edge_program = copy.deepcopy(edge_program)
+        # copy_edge_program = copy.deepcopy(edge_program)
 
         # Move the edge_program from CPU to CUDA for aoti compile
-        cuda_edge_program = move_to_device_pass(copy_edge_program, "cuda")
+        cuda_edge_program = move_to_device_pass(edge_program, "cuda")
 
         edge_program_module = cuda_edge_program.module()
-        args, kwargs = copy_edge_program.example_inputs
-
-        # # Deep copy args and move tensors to CUDA for aot_compile
-        # def move_to_cuda(obj):
-        #     if isinstance(obj, torch.Tensor):
-        #         return obj.cuda()
-        #     elif isinstance(obj, (list, tuple)):
-        #         return type(obj)(move_to_cuda(item) for item in obj)
-        #     elif isinstance(obj, dict):
-        #         return {key: move_to_cuda(value) for key, value in obj.items()}
-        #     else:
-        #         return obj
-
-        # args = move_to_cuda(copy.deepcopy(args))
-        # kwargs = move_to_cuda(copy.deepcopy(kwargs))
+        args, kwargs = cuda_edge_program.example_inputs
 
         output_path = os.path.join(os.getcwd(), "aoti.so")
 
@@ -122,10 +106,6 @@ def preprocess(
                     "Please add them to the AOTI backend."
                 )
 
-        assert so_path == output_path, f"Expected {output_path} but got {so_path}"
-
-        print("so_path", so_path)
-
         with open(so_path, "rb") as f:
             so_data = f.read()
 

backends/aoti/aoti_partitioner.py

@@ -6,8 +6,7 @@
 
 # pyre-unsafe
 
-import operator
-from typing import Callable, cast, Dict, final, List, Optional, Set, Tuple
+from typing import Callable, Dict, final, List, Optional, Tuple
 
 import torch
 from executorch.backends.aoti.aoti_backend import AotiBackend  # usort: skip
@@ -18,65 +17,26 @@
     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 == operator.getitem
-        #     or str(node.target._op) not in inductor_fallback_ops
-        #     or str(node.target._op) in supported_fallback_operators
-        # )
-
-        supported = node.op == "call_function"
-
-        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")
-        print("entering partitioner...")
-
-        partition_tags = {}
-
-        capability_partitioner = CapabilityBasedPartitioner(
-            exported_program.graph_module,
-            AOTISupportedOperators(),
-            allows_single_node_partition=True,
-        )
-        partition_list = capability_partitioner.propose_partitions()
-
-        assert len(partition_list) == 1, "Graph break is not supported yet"
-
-        print(f"graph breaks into {len(partition_list)} parts")
+        """
+        Fully delegate the graph to AOTInductor by tagging all nodes as a single partition.
+        """
 
-        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
+        partition_tags: Dict[str, DelegationSpec] = {}
+        for node in exported_program.graph.nodes:
+            if node.op != "call_function":
+                continue
+            tag = f"tag0"
+            node.meta["delegation_tag"] = tag
+            partition_tags[tag] = self.delegation_spec
 
         tag_constant_data(exported_program)
 
@@ -89,15 +49,13 @@ def ops_to_not_decompose(
     ) -> Tuple[List[torch._ops.OpOverload], Optional[Callable[[torch.fx.Node], bool]]]:
         """
         Return a list of operations that should not be decomposed and let the AOT compiler handle them.
+        Currently we skip decomposing all ops and let the AOT compiler handle them.
         """
         do_not_decompose = set()
-        op_support = AOTISupportedOperators()
 
         for node in ep.graph.nodes:
-            if (
-                node.op == "call_function"
-                and isinstance(node.target, torch._ops.OpOverload)
-                and op_support.is_node_supported(None, node)
+            if node.op == "call_function" and isinstance(
+                node.target, torch._ops.OpOverload
             ):
                 do_not_decompose.add(node.target)
         return list(do_not_decompose), None