Add export_add.py

Author: larryliu0820

backends/aoti/aoti_partitioner.py

@@ -28,7 +28,7 @@ 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,
+            exir_ops.edge.dim_order_ops._to_dim_order_copy.default,
         ]
 
         return supported

exir/backend/backend_api.py

@@ -711,6 +711,9 @@ def to_backend(
             fake_edge_program = copy.deepcopy(edge_program)
         partitioner_result = partitioner_instance(fake_edge_program)
         tagged_exported_program = partitioner_result.tagged_exported_program
+        # Make sure tagged_exported_program has the same example_inputs as edge_program
+        tagged_exported_program.example_inputs = edge_program.example_inputs
+
         method_to_tagged_exported_program[method_name] = tagged_exported_program
 
         # Check that the partitioner did not modify the original graph

exir/lowered_backend_module.py

@@ -734,6 +734,7 @@ def create_exported_program_from_submodule(
                     ),
                 )
             ],
+            example_inputs=owning_program.example_inputs,
             constants=subgraph_constants,
             verifiers=[owning_program.verifier],
         ),

export_add.py

@@ -0,0 +1,31 @@
+import torch
+from executorch.backends.aoti.aoti_partitioner import AotiPartitioner
+from executorch.exir import to_edge
+from torch.export import export
+
+
+# Start with a PyTorch model that adds two input tensors (matrices)
+class Add(torch.nn.Module):
+    def __init__(self):
+        super(Add, self).__init__()
+
+    def forward(self, x: torch.Tensor, y: torch.Tensor):
+        # return triton_transpose_acc(x, y)
+        return (x.cuda() + y.cuda()).cpu()
+
+
+# 1. torch.export: Defines the program with the ATen operator set.
+aten_dialect = export(
+    Add(), (torch.ones(10, device="cpu"), torch.ones(10, device="cpu"))
+)
+# 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("add.pte", "wb") as file:
+    file.write(executorch_program.buffer)

install_requirements.py

@@ -59,7 +59,7 @@ def python_is_compatible():
 
 
 # The pip repository that hosts nightly torch packages.
-TORCH_NIGHTLY_URL = "https://download.pytorch.org/whl/nightly/cu121"
+TORCH_NIGHTLY_URL = "https://download.pytorch.org/whl/nightly/cu126"
 
 
 # Since ExecuTorch often uses main-branch features of pytorch, only the nightly