Merge branch 'main' into modelLoadAndExecutionLogging

Author: leafs1

.buckconfig

@@ -39,6 +39,7 @@
 
 [buck2]
 restarter=true
+file_watcher=notify
 
 [oss]
 folly_cxx_tests = False

backends/arm/_passes/__init__.py

@@ -24,6 +24,7 @@
 from .decompose_gelu_pass import DecomposeGeluPass  # noqa
 from .decompose_layernorm_pass import DecomposeLayerNormPass  # noqa
 from .decompose_leaky_relu_pass import DecomposeLeakyReLUPass  # noqa
+from .decompose_linalg_vector_norm_pass import DecomposeLinearVectorNormPass  # noqa
 from .decompose_linear_pass import DecomposeLinearPass  # noqa
 from .decompose_meandim_pass import DecomposeMeanDimPass  # noqa
 from .decompose_ne_pass import DecomposeNotEqualPass  # noqa

backends/arm/_passes/arm_pass_manager.py

@@ -29,6 +29,7 @@
     DecomposeLayerNormPass,
     DecomposeLeakyReLUPass,
     DecomposeLinearPass,
+    DecomposeLinearVectorNormPass,
     DecomposeMeanDimPass,
     DecomposeNotEqualPass,
     DecomposeSelectPass,
@@ -86,6 +87,7 @@ def _tosa_080_BI_pipeline(self, exported_program: ExportedProgram) -> GraphModul
         self.add_pass(ConvertSplitToSlicePass())
         self.add_pass(ConvertMmToBmmPass())
         self.add_pass(DecomposeLinearPass())
+        self.add_pass(DecomposeLinearVectorNormPass())
         self.add_pass(DecomposeMeanDimPass())
         self.add_pass(ConvertFullLikeToFullPass())
         self.add_pass(ConvertToClampPass())
@@ -133,6 +135,7 @@ def _tosa_080_MI_pipeline(self, exported_program: ExportedProgram) -> GraphModul
         self.add_pass(FuseBatchnorm2DPass(exported_program))
         self.add_pass(ConvertMmToBmmPass())
         self.add_pass(DecomposeLinearPass())
+        self.add_pass(DecomposeLinearVectorNormPass())
         self.add_pass(DecomposeLeakyReLUPass())
         self.add_pass(DecomposeBatchNormPass())
         self.add_pass(DecomposeLayerNormPass())
@@ -207,6 +210,7 @@ def transform_for_annotation_pipeline(self, graph_module: GraphModule):
         self.add_pass(DecomposeCosineSimilarityPass())
         self.add_pass(DecomposeDivPass())
         self.add_pass(DecomposeLeakyReLUPass())
+        self.add_pass(DecomposeLinearVectorNormPass())
         self.add_pass(DecomposeSqrtPass())
         self.add_pass(DecomposeSiluPass())
 

backends/arm/_passes/decompose_linalg_vector_norm_pass.py

@@ -0,0 +1,78 @@
+# Copyright 2025 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 torch
+from executorch.exir.pass_base import ExportPass
+
+
+class DecomposeLinearVectorNormPass(ExportPass):
+    """
+    This pass decomposes aten.linalg_vector_norm.default into more primitive ops.
+    We need to add this pass before quantization for graph annotation.
+    By default, aten.linalg_vector_norm op is decomposed during legalization to Edge IR.
+
+    The decomposition is as follows:
+
+      For p == 1:
+          out = REDUCE_SUM(ABS(x), dims, keepdim)
+
+      For p == 2:
+          out = SQRT(REDUCE_SUM(MUL(x, x), dims, keepdim))
+
+      For arbitrary p:
+          We dont support arbitrary p, because our decomposition looks like
+          out = POW(REDUCE_SUM(POW(ABS(x), p), dims, keepdim), 1/p)
+          In this case we need to wrap p into Tensor and we need to know
+          dtype prior, but we dont know this from FX graph.
+    """
+
+    torch_linalg_vector_norm = (torch.ops.aten.linalg_vector_norm.default,)
+
+    def call_operator(self, op, args, kwargs, meta):
+        if op not in self.torch_linalg_vector_norm:
+            return super().call_operator(op, args, kwargs, meta)
+
+        # Extract inputs and optional arguments.
+        # Expected args:
+        #   args[0]: input tensor
+        #   args[1]: norm order 'p' (optional, default: 2.0)
+        #   args[2]: dimensions to reduce (should be provided)
+        #   args[3]: keepdim flag (optional, default: False)
+        input_tensor = args[0]
+        norm_order = args[1] if len(args) > 1 else 2.0
+        norm_dim = args[2] if len(args) > 2 else None
+        keepdim = args[3] if len(args) > 3 else False
+
+        if norm_order not in (1, 2):
+            raise ValueError(
+                f"The order of {norm_order}\n"
+                f"is not supported for linalg_vector_norm operator"
+            )
+
+        if norm_dim is None:
+            raise ValueError("The norm_dim for linalg_vector_norm is None.")
+
+        dims = [norm_dim] if isinstance(norm_dim, int) else list(norm_dim)
+
+        # Decomposition based on norm order.
+        if norm_order == 1:
+            op1 = super().call_operator(
+                torch.ops.aten.abs.default, (input_tensor,), {}, meta
+            )
+            op2 = super().call_operator(
+                torch.ops.aten.sum.dim_IntList, (op1, dims, keepdim), {}, meta
+            )
+            return op2
+
+        elif norm_order == 2:
+            # For p == 2, decomposition is sqrt(sum(x * x, dims, keepdim))
+            op1 = super().call_operator(
+                torch.ops.aten.mul.Tensor, (input_tensor, input_tensor), {}, meta
+            )
+            op2 = super().call_operator(
+                torch.ops.aten.sum.dim_IntList, (op1, dims, keepdim), {}, meta
+            )
+            op3 = super().call_operator(torch.ops.aten.sqrt.default, (op2,), {}, meta)
+            return op3

backends/arm/operators/op_abs.py

@@ -15,6 +15,7 @@
 )
 from executorch.backends.arm.operators.operator_validation_utils import (
     validate_num_inputs,
+    validate_same_dtype,
 )
 from executorch.backends.arm.tosa_mapping import TosaArg
 from executorch.backends.arm.tosa_specification import TosaSpecification
@@ -43,13 +44,8 @@ def define_node(
         import tosa_tools.v0_80.serializer.tosa_serializer as ts  # type: ignore
 
         validate_num_inputs(self.target, inputs, 1)
-        # Specification (0.80) states that input and output types
-        # should all be the same
-        if not (inputs[0].dtype == output.dtype):
-            raise ValueError(
-                "All inputs and outputs need same dtype."
-                f"Got {inputs[0].dtype=}, {output.dtype=}"
-            )
+        validate_same_dtype(self.target, [*inputs, output])
+
         # Handle int8 (quantized) and int32
         if not (inputs[0].dtype in [ts.DType.INT8, ts.DType.INT32]):
             raise ValueError(
@@ -110,13 +106,7 @@ def define_node(
         import tosa_tools.v0_80.serializer.tosa_serializer as ts  # type: ignore
 
         validate_num_inputs(self.target, inputs, 1)
-        # Specification (0.80) states that input and output types
-        # should all be the same
-        if not (inputs[0].dtype == output.dtype):
-            raise ValueError(
-                "All inputs and output need same dtype."
-                f"Got {inputs[0].dtype=}, {output.dtype=}"
-            )
+        validate_same_dtype(self.target, [*inputs, output])
 
         if inputs[0].dtype in [ts.DType.INT8, ts.DType.INT32]:
             # Call the inherited define_node for handling integers
@@ -163,14 +153,8 @@ def define_node(
         import serializer.tosa_serializer as ts  # type: ignore
 
         validate_num_inputs(self.target, inputs, 1)
+        validate_same_dtype(self.target, [*inputs, output])
 
-        # Specification (1.0) states that input and output types
-        # should all be the same
-        if not (inputs[0].dtype == output.dtype):
-            raise ValueError(
-                "All inputs and outputs need same dtype."
-                f"Got {inputs[0].dtype=}, {output.dtype=}"
-            )
         # Handle int8 (quantized) and int32
         if not (inputs[0].dtype in [ts.DType.INT8, ts.DType.INT32]):
             raise ValueError(
@@ -232,14 +216,7 @@ def define_node(
         import serializer.tosa_serializer as ts  # type: ignore
 
         validate_num_inputs(self.target, inputs, 1)
-
-        # Specification (1.0) states that input and output types
-        # should all be the same
-        if not (inputs[0].dtype == output.dtype):
-            raise ValueError(
-                "All inputs and output need same dtype."
-                f"Got {inputs[0].dtype=}, {output.dtype=}"
-            )
+        validate_same_dtype(self.target, [*inputs, output])
 
         if inputs[0].dtype in [ts.DType.INT8, ts.DType.INT32]:
             # Call the inherited define_node for handling integers

backends/arm/operators/op_add.py

@@ -16,6 +16,7 @@
 )
 from executorch.backends.arm.operators.operator_validation_utils import (
     validate_num_inputs,
+    validate_same_dtype,
 )
 from executorch.backends.arm.tosa_mapping import TosaArg
 from executorch.backends.arm.tosa_specification import TosaSpecification
@@ -44,14 +45,8 @@ def define_node(
         import tosa_tools.v0_80.serializer.tosa_serializer as ts  # type: ignore
 
         validate_num_inputs(self.target, inputs, 2)
-        # Specification (0.80) states that input and output types
-        # should all be the same
-        if inputs[0].dtype != inputs[1].dtype or inputs[0].dtype != output.dtype:
-            raise TypeError(
-                f"All IO needs to have the same data type, got input 1: "
-                f"{inputs[0].dtype}, input 2: {inputs[1].dtype} and output: "
-                f"{output.dtype}"
-            )
+        validate_same_dtype(self.target, [*inputs, output])
+
         # Handle int8 (quantized) and int32
         supported_dtypes = [ts.DType.INT8, ts.DType.INT32]
         if inputs[0].dtype not in supported_dtypes:
@@ -123,14 +118,7 @@ def define_node(
         import tosa_tools.v0_80.serializer.tosa_serializer as ts  # type: ignore
 
         validate_num_inputs(self.target, inputs, 2)
-        # Specification (0.80) states that input and output types
-        # should all be the same
-        if inputs[0].dtype != inputs[1].dtype or inputs[0].dtype != output.dtype:
-            raise TypeError(
-                f"All IO needs to have the same data type, got input 1: "
-                f"{inputs[0].dtype}, input 2: {inputs[1].dtype} and output: "
-                f"{output.dtype}"
-            )
+        validate_same_dtype(self.target, [*inputs, output])
 
         if inputs[0].dtype in [ts.DType.INT8, ts.DType.INT32]:
             # Call the inherited define_node for handling integers
@@ -175,15 +163,8 @@ def define_node(
         import serializer.tosa_serializer as ts  # type: ignore
 
         validate_num_inputs(self.target, inputs, 2)
+        validate_same_dtype(self.target, [*inputs, output])
 
-        # Specification (1.0) states that input and output types
-        # should all be the same
-        if inputs[0].dtype != inputs[1].dtype or inputs[0].dtype != output.dtype:
-            raise TypeError(
-                f"All IO needs to have the same data type, got input 1: "
-                f"{inputs[0].dtype}, input 2: {inputs[1].dtype} and output: "
-                f"{output.dtype}"
-            )
         # Handle int8 (quantized) and int32
         supported_dtypes = [ts.DType.INT8, ts.DType.INT32]
         if inputs[0].dtype not in supported_dtypes:
@@ -245,15 +226,7 @@ def define_node(
         import serializer.tosa_serializer as ts  # type: ignore
 
         validate_num_inputs(self.target, inputs, 2)
-
-        # Specification (1.0) states that input and output types
-        # should all be the same
-        if inputs[0].dtype != inputs[1].dtype or inputs[0].dtype != output.dtype:
-            raise TypeError(
-                f"All IO needs to have the same data type, got input 1: "
-                f"{inputs[0].dtype}, input 2: {inputs[1].dtype} and output: "
-                f"{output.dtype}"
-            )
+        validate_same_dtype(self.target, [*inputs, output])
 
         if inputs[0].dtype in [ts.DType.INT8, ts.DType.INT32]:
             # Call the inherited define_node for handling integers

backends/arm/operators/op_amax.py

@@ -11,6 +11,7 @@
 )
 from executorch.backends.arm.operators.operator_validation_utils import (
     validate_num_inputs,
+    validate_same_dtype,
 )
 from executorch.backends.arm.tosa_mapping import TosaArg
 from torch.fx import Node
@@ -35,6 +36,7 @@ def define_node(
         import tosa_tools.v0_80.serializer.tosa_serializer as ts
 
         validate_num_inputs(self.target, inputs, 3)
+        validate_same_dtype(self.target, [inputs[0], output])
 
         input = inputs[0]
         dim = inputs[1].number
@@ -77,6 +79,7 @@ def define_node(
         import serializer.tosa_serializer as ts
 
         validate_num_inputs(self.target, inputs, 3)
+        validate_same_dtype(self.target, [inputs[0], output])
 
         input = inputs[0]
         dim = inputs[1].number

backends/arm/operators/op_amin.py

@@ -11,6 +11,7 @@
 )
 from executorch.backends.arm.operators.operator_validation_utils import (
     validate_num_inputs,
+    validate_same_dtype,
 )
 from executorch.backends.arm.tosa_mapping import TosaArg
 from torch.fx import Node
@@ -35,6 +36,7 @@ def define_node(
         import tosa_tools.v0_80.serializer.tosa_serializer as ts
 
         validate_num_inputs(self.target, inputs, 3)
+        validate_same_dtype(self.target, [inputs[0], output])
 
         input = inputs[0]
         dim = inputs[1].number
@@ -77,6 +79,7 @@ def define_node(
         import serializer.tosa_serializer as ts
 
         validate_num_inputs(self.target, inputs, 3)
+        validate_same_dtype(self.target, [inputs[0], output])
 
         input = inputs[0]
         dim = inputs[1].number

backends/arm/operators/op_any.py

@@ -12,6 +12,7 @@
 )
 from executorch.backends.arm.operators.operator_validation_utils import (
     validate_num_inputs,
+    validate_same_dtype,
 )
 
 from executorch.backends.arm.tosa_mapping import TosaArg  # type: ignore
@@ -34,12 +35,8 @@ def define_node(
         import tosa_tools.v0_80.serializer.tosa_serializer as ts  # type: ignore
 
         validate_num_inputs(self.target, inputs, 3)
+        validate_same_dtype(self.target, [inputs[0], output])
 
-        if not (inputs[0].dtype == output.dtype):
-            raise ValueError(
-                "All inputs and outputs need same dtype."
-                f"Got {ts.DTypeNames[inputs[0].dtype]=}, {ts.DTypeNames[output.dtype]=}."
-            )
         if not (inputs[0].dtype == ts.DType.BOOL):
             raise ValueError("All inputs need to be BOOL." f"Got {inputs[0].dtype=}")
 
@@ -75,12 +72,8 @@ def define_node(
         import serializer.tosa_serializer as ts
 
         validate_num_inputs(self.target, inputs, 3)
+        validate_same_dtype(self.target, [inputs[0], output])
 
-        if not (inputs[0].dtype == output.dtype):
-            raise ValueError(
-                "All inputs and outputs need same dtype."
-                f"Got {ts.DTypeNames[inputs[0].dtype]=}, {ts.DTypeNames[output.dtype]=}."
-            )
         if not (inputs[0].dtype == ts.DType.BOOL):
             raise ValueError("All inputs need to be BOOL." f"Got {inputs[0].dtype=}")
 

backends/arm/operators/op_avg_pool2d.py

@@ -18,6 +18,7 @@
 )
 from executorch.backends.arm.operators.operator_validation_utils import (
     validate_num_inputs,
+    validate_same_dtype,
 )
 from executorch.backends.arm.tosa_mapping import TosaArg
 from executorch.backends.arm.tosa_specification import TosaSpecification
@@ -89,6 +90,7 @@ def define_node(
         import tosa_tools.v0_80.serializer.tosa_serializer as ts  # type: ignore
 
         validate_num_inputs(self.target, inputs, [3, 4, 6])
+        validate_same_dtype(self.target, [inputs[0], output])
 
         supported_dtypes = [ts.DType.INT8]
         if inputs[0].dtype not in supported_dtypes:
@@ -128,6 +130,7 @@ def define_node(
         import tosa_tools.v0_80.serializer.tosa_serializer as ts  # type: ignore
 
         validate_num_inputs(self.target, inputs, [3, 4, 6])
+        validate_same_dtype(self.target, [inputs[0], output])
 
         supported_dtypes = [ts.DType.INT8, ts.DType.FP32]
         if inputs[0].dtype not in supported_dtypes:
@@ -220,6 +223,7 @@ def define_node(
         import serializer.tosa_serializer as ts  # type: ignore
 
         validate_num_inputs(self.target, inputs, [3, 4, 6])
+        validate_same_dtype(self.target, [inputs[0], output])
 
         supported_dtypes = [ts.DType.INT8]
         if inputs[0].dtype not in supported_dtypes:
@@ -262,6 +266,7 @@ def define_node(
         import serializer.tosa_serializer as ts  # type: ignore
 
         validate_num_inputs(self.target, inputs, [3, 4, 6])
+        validate_same_dtype(self.target, [inputs[0], output])
 
         supported_dtypes = [ts.DType.INT8, ts.DType.FP32]
         if inputs[0].dtype not in supported_dtypes: