2024-08-27 nightly release (dc66414)

backends/arm/arm_partitioner.py

@@ -39,6 +39,7 @@ def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
             exir_ops.edge.aten.add.Tensor,
             exir_ops.edge.aten.addmm.default,
             exir_ops.edge.aten.expand_copy.default,
+            exir_ops.edge.aten.cat.default,
             exir_ops.edge.aten.permute_copy.default,
             exir_ops.edge.aten.hardtanh.default,
             exir_ops.edge.aten.convolution.default,
@@ -51,12 +52,14 @@ def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
             exir_ops.edge.aten.sigmoid.default,
             exir_ops.edge.aten.mm.default,
             exir_ops.edge.aten.repeat.default,
+            exir_ops.edge.aten.relu.default,
             exir_ops.edge.aten._softmax.default,
             exir_ops.edge.aten.slice_copy.Tensor,
             exir_ops.edge.aten.sub.Tensor,
             exir_ops.edge.aten.view_copy.default,
             exir_ops.edge.aten.clone.default,
             exir_ops.edge.aten.mean.dim,
+            exir_ops.edge.aten.unsqueeze_copy.default,
             operator.getitem,
             exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
             exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,

backends/arm/operators/__init__.py

@@ -9,6 +9,7 @@
     op_addmm,
     op_avg_pool2d,
     op_batch_norm,
+    op_cat,
     op_conv2d,
     op_dequant,
     op_div,
@@ -20,10 +21,12 @@
     op_mul,
     op_permute,
     op_quant,
+    op_relu,
     op_repeat,
     op_sigmoid,
     op_slice,
     op_softmax,
     op_sub,
+    op_unsqueeze,
     op_view,
 )

backends/arm/operators/op_cat.py

@@ -0,0 +1,45 @@
+# 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.
+
+from typing import List
+
+import serializer.tosa_serializer as ts
+from executorch.backends.arm.operators.node_visitor import (
+    NodeVisitor,
+    register_node_visitor,
+)
+from executorch.backends.arm.tosa_mapping import TosaArg
+from serializer.tosa_serializer import TosaOp
+from torch.fx import Node
+
+
+@register_node_visitor
+class CatVisitor(NodeVisitor):
+    target = "aten.cat.default"
+
+    def __init__(self, *args):
+        super().__init__(*args)
+
+    def define_node(
+        self,
+        node: Node,
+        tosa_graph: ts.TosaSerializer,
+        inputs: List[TosaArg],
+        output: TosaArg,
+        is_quant_node: bool,
+    ) -> None:
+
+        tensors = inputs[0].special
+        dim = 0 if len(inputs) < 2 else inputs[1].number
+        rank = len(output.shape)
+        dim = (dim + rank) % rank
+        dim = output.dim_order.index(dim)
+
+        attr = ts.TosaSerializerAttribute()
+        attr.AxisAttribute(dim)
+
+        tosa_graph.addOperator(
+            TosaOp.Op().CONCAT, [tensor.name for tensor in tensors], [output.name], attr
+        )

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/operators/op_relu.py

@@ -0,0 +1,55 @@
+# 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 executorch.backends.arm.tosa_quant_utils as tqutils
+import serializer.tosa_serializer as ts
+import torch.fx
+from executorch.backends.arm.operators.node_visitor import (
+    NodeVisitor,
+    register_node_visitor,
+)
+from executorch.backends.arm.tosa_mapping import TosaArg
+from serializer.tosa_serializer import TosaOp
+
+
+@register_node_visitor
+class ReluVisitor(NodeVisitor):
+    target = "aten.relu.default"
+
+    def __init__(self, *args):
+        super().__init__(*args)
+
+    def define_node(
+        self,
+        node: torch.fx.Node,
+        tosa_graph: ts.TosaSerializer,
+        inputs: list[TosaArg],
+        output: TosaArg,
+        is_quant_node: bool,
+    ) -> None:
+        attr = ts.TosaSerializerAttribute()
+
+        clamp_min_fp = 0.0
+        clamp_max_fp = 0.0
+        clamp_min_qs = 0
+        clamp_max_qs = 0
+        if is_quant_node:
+            out_qargs = tqutils.get_quant_node_args(list(node.users)[0])
+            clamp_min_qs = tqutils.quantize_value(0, out_qargs)
+            clamp_max_qs = tqutils.quantize_value(float("inf"), out_qargs)
+
+        else:
+            clamp_min_fp = 0
+            clamp_max_fp = float("inf")
+
+        attr.ClampAttribute(
+            tosa_graph.builder,
+            clamp_min_qs,
+            clamp_max_qs,
+            clamp_min_fp,
+            clamp_max_fp,
+        )
+
+        tosa_graph.addOperator(TosaOp.Op().CLAMP, [inputs[0].name], [output.name], attr)

backends/arm/operators/op_unsqueeze.py

@@ -0,0 +1,51 @@
+# 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.
+#
+#  Follows this specification: https://pytorch.org/docs/stable/generated/torch.unsqueeze.html
+
+import serializer.tosa_serializer as ts
+import torch.fx
+from executorch.backends.arm.operators.node_visitor import (
+    NodeVisitor,
+    register_node_visitor,
+)
+from executorch.backends.arm.tosa_mapping import TosaArg
+from executorch.backends.arm.tosa_utils import tosa_shape
+from serializer.tosa_serializer import TosaOp
+
+
+@register_node_visitor
+class UnsqueezeVisitor(NodeVisitor):
+    target = "aten.unsqueeze_copy.default"
+
+    def __init__(self, *args):
+        super().__init__(*args)
+
+    def define_node(
+        self,
+        node: torch.fx.Node,
+        tosa_graph: ts.TosaSerializer,
+        inputs: list[TosaArg],
+        output: TosaArg,
+        is_quant_node: bool,
+    ) -> None:
+
+        dim = inputs[1].number
+        shape = inputs[0].shape
+        rank = len(shape)
+
+        assert -rank - 1 <= dim < rank + 1
+        if dim < 0:
+            dim = dim + rank + 1
+
+        new_shape = list(shape)
+        new_shape.insert(dim, 1)
+        new_shape = tosa_shape(new_shape, output.dim_order)
+
+        attr = ts.TosaSerializerAttribute()
+        attr.ReshapeAttribute(new_shape)
+        tosa_graph.addOperator(
+            TosaOp.Op().RESHAPE, [inputs[0].name], [output.name], attr
+        )

backends/arm/quantizer/arm_quantizer.py

@@ -267,6 +267,7 @@ class ArmQuantizer(Quantizer):
         "mul",
         "sigmoid",
         "mm",
+        "cat",
     ]
 
     def __init__(self) -> None:

backends/arm/quantizer/arm_quantizer_utils.py

@@ -102,12 +102,19 @@ def is_input_ok_for_quantization(input_act: Node, gm: GraphModule):
     )
 
 
+def get_node_target(module: torch.nn.Module | GraphModule, target_str: str):
+    targets = target_str.split(".")
+    for target in targets[:-1]:
+        module = module.get_submodule(target)
+    return getattr(module, targets[-1])
+
+
 def is_input_large_scalar(node: Node, gm: GraphModule):
     """Check if input is a large scalar value. So that we can skip quantization for the node
     since histc op (in HistogramObserver) only works for values up to certain upper bound
     """
     if node.op == "get_attr" and isinstance(node.target, str):
-        tensor = getattr(gm, node.target)
+        tensor = get_node_target(gm, node.target)
         # torch.histc works until this upper bound
         HISTC_UPPER_BOUND = 3.4028235e15
         return tensor.numel() == 1 and abs(tensor.item()) > HISTC_UPPER_BOUND
@@ -131,6 +138,7 @@ def is_share_obs_or_fq_op(op: Callable) -> bool:
     return op in [
         torch.ops.aten.hardtanh.default,
         torch.ops.aten.hardtanh_.default,
+        torch.ops.aten.relu.default,
         torch.ops.aten.mean.default,
         torch.ops.aten.mean.dim,
         torch.ops.aten.permute.default,

backends/arm/quantizer/quantization_annotation/__init__.py

@@ -49,6 +49,7 @@ def decorator(annotator: AnnotatorType):
 from . import (  # noqa
     adaptive_ang_pool2d_annotator,
     add_annotator,
+    cat_annotator,
     conv_annotator,
     linear_annotator,
     max_pool2d_annotator,

backends/arm/quantizer/quantization_annotation/cat_annotator.py

@@ -0,0 +1,66 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# 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 itertools
+from typing import Callable, List, Optional
+
+import torch.fx
+from executorch.backends.arm.quantizer import arm_quantizer_utils
+from executorch.backends.arm.quantizer.quantization_annotation import register_annotator
+from executorch.backends.arm.quantizer.quantization_config import QuantizationConfig
+from torch.ao.quantization.quantizer import (
+    QuantizationAnnotation,
+    SharedQuantizationSpec,
+)
+from torch.fx import Node
+from torch.fx.passes.utils.source_matcher_utils import get_source_partitions
+
+
+@register_annotator("cat")
+def _annotate_cat(
+    gm: torch.fx.GraphModule,
+    quantization_config: QuantizationConfig,
+    filter_fn: Optional[Callable[[Node], bool]] = None,
+) -> Optional[List[List[Node]]]:
+    cat_partitions = get_source_partitions(gm.graph, [torch.cat], filter_fn)
+    cat_partitions = list(itertools.chain.from_iterable(cat_partitions.values()))
+    annotated_partitions = []
+    for cat_partition in cat_partitions:
+        annotated_partitions.append(cat_partition.nodes)
+        cat_node = cat_partition.output_nodes[0]
+        if arm_quantizer_utils.is_annotated(cat_node):
+            continue
+
+        input_acts = cat_node.args[0]
+        input_act0 = input_acts[0]
+
+        input_act_qspec = quantization_config.get_input_act_qspec()
+        shared_with_input0_qspec = SharedQuantizationSpec((input_act0, cat_node))
+
+        input_qspec_map = {}
+
+        # First input is set to input qspec from the quantization config.
+        if isinstance(input_act0, Node):
+            if not arm_quantizer_utils.is_input_ok_for_quantization(input_act0, gm):
+                continue
+            input_qspec_map[input_act0] = input_act_qspec
+
+        # For the rest of the inputs, share qspec with first.
+        # If we can't quantize any of the inputs, abort annotation.
+        for input_act in input_acts[1:]:
+            if isinstance(input_act, Node):
+                if not arm_quantizer_utils.is_input_ok_for_quantization(input_act, gm):
+                    continue
+                if input_act is not input_act0:
+                    input_qspec_map[input_act] = shared_with_input0_qspec
+
+        if input_qspec_map is not None:
+            cat_node.meta["quantization_annotation"] = QuantizationAnnotation(
+                input_qspec_map=input_qspec_map,
+                output_qspec=shared_with_input0_qspec,
+                _annotated=True,
+            )
+    return annotated_partitions

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))