Add support for lifted tensors in ArmPartitioner

Differential Revision: D61542388

Pull Request resolved: #4788

Author: Erik-Lundell

backends/arm/operators/op_placeholder.py

@@ -5,7 +5,7 @@
 
 import numpy as np
 import serializer.tosa_serializer as ts
-import torch
+import torch.fx
 from executorch.backends.arm.tosa_mapping import TosaArg
 from executorch.backends.arm.tosa_quant_utils import (
     get_quant_arg_dtype,
@@ -130,6 +130,21 @@ def process_inputs_to_buffers(
     )
 
 
+def process_inputs_to_lifted_tensor_constants(
+    node: torch.fx.Node,
+    tosa_graph: ts.TosaSerializer,
+    edge_program: ExportedProgram,
+):
+    arg = TosaArg(node)
+    tensor_name = edge_program.graph_signature.inputs_to_lifted_tensor_constants[
+        arg.name
+    ]
+    tensor = edge_program.tensor_constants[tensor_name]
+    tensor_data = tensor.detach().numpy()
+
+    tosa_graph.addConst(tensor_data.shape, arg.dtype, tensor_data, name=arg.name)
+
+
 def process_placeholder(
     node: torch.fx.Node,
     tosa_graph: ts.TosaSerializer,
@@ -145,5 +160,11 @@ def process_placeholder(
         process_inputs_to_parameters(node, tosa_graph, edge_program)
     elif node.name in edge_program.graph_signature.inputs_to_buffers:
         process_inputs_to_buffers(node, tosa_graph, edge_program)
+    elif node.name in edge_program.graph_signature.inputs_to_lifted_tensor_constants:
+        process_inputs_to_lifted_tensor_constants(node, tosa_graph, edge_program)
+    elif node.name in edge_program.graph_signature.inputs_to_lifted_custom_objs:
+        raise NotImplementedError(
+            "Placeholder is of type 'lifted custom object' which is not supported."
+        )
     else:
-        raise RuntimeError(f"Unknown placeholder {node.name}")
+        raise RuntimeError(f"Placeholder '{node.name}' is of unknown type.")

backends/arm/test/misc/test_lifted_tensor.py

@@ -0,0 +1,42 @@
+# Copyright 2024 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import unittest
+
+import torch
+from executorch.backends.arm.test import common
+from executorch.backends.arm.test.tester.arm_tester import ArmTester
+
+
+class LiftedTensor(torch.nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.lifted_tensor = torch.Tensor([[1, 2], [3, 4]])
+
+    def forward(self, x: torch.Tensor, length) -> torch.Tensor:
+        sliced = self.lifted_tensor[:, :length]
+        return sliced + x
+
+
+class TestLiftedTensor(unittest.TestCase):
+    """Tests the ArmPartitioner with a placeholder of type lifted tensor."""
+
+    def test_partition_lifted_tensor(self):
+        tester = (
+            ArmTester(
+                LiftedTensor(),
+                example_inputs=(torch.ones(2, 2), 2),
+                compile_spec=common.get_tosa_compile_spec(),
+            )
+            .export()
+            .to_edge()
+            .dump_artifact()
+        )
+        signature = tester.get_artifact().exported_program().graph_signature
+        assert len(signature.lifted_tensor_constants) > 0
+        tester.partition()
+        tester.to_executorch()
+        tester.run_method_and_compare_outputs((torch.ones(2, 2), 2))