Merge branch 'main' into uint8-io

Author: per

backends/arm/ethosu/partitioner.py

@@ -3,9 +3,9 @@
 # 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 final, Optional, Sequence
 
+import torch
 from executorch.backends.arm.ethosu import EthosUBackend, EthosUCompileSpec
 from executorch.backends.arm.tosa.partitioner import TOSAPartitioner
 from executorch.exir.backend.partitioner import DelegationSpec
@@ -33,3 +33,4 @@ def __init__(
         )
         self.additional_checks = additional_checks
         self.tosa_spec = compile_spec.tosa_spec
+        self._custom_partition_ops: set[torch._ops.OpOverload] = set()

backends/arm/operators/__init__.py

@@ -51,6 +51,7 @@
     op_to_dim_order_copy,
     op_tosa_conv2d,
     op_tosa_conv3d,
+    op_tosa_custom,
     op_tosa_depthwise_conv2d,
     op_tosa_gather,
     op_tosa_matmul,

backends/arm/operators/op_tosa_custom.py

@@ -0,0 +1,85 @@
+# Copyright 2026 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 Any, List
+
+import torch
+import tosa_serializer as ts
+
+from executorch.backends.arm.operators.node_visitor import (
+    NodeVisitor,
+    register_node_visitor,
+)
+from executorch.backends.arm.tosa.mapping import TosaArg
+
+
+@register_node_visitor
+class CustomVisitor(NodeVisitor):
+    """Lower the TOSA CUSTOM op from the TOSA backend dialect."""
+
+    target = "tosa.CUSTOM.default"
+
+    def define_node(
+        self,
+        node: torch.fx.Node,
+        tosa_graph: Any,
+        inputs: List[TosaArg],
+        output: TosaArg,
+    ) -> None:
+        allowed_kwargs = {"operator_name", "domain_name", "implementation_attrs"}
+        unexpected = set(node.kwargs.keys()) - allowed_kwargs
+        if unexpected:
+            raise ValueError(
+                f"tosa.CUSTOM received unexpected kwargs: {sorted(unexpected)}"
+            )
+
+        operator_name = node.kwargs.get("operator_name")
+        domain_name = node.kwargs.get("domain_name")
+        implementation_attrs = node.kwargs.get("implementation_attrs")
+
+        if operator_name is None or domain_name is None:
+            raise ValueError(
+                "tosa.CUSTOM requires operator_name and domain_name in kwargs"
+            )
+
+        if implementation_attrs is None:
+            impl_list = []
+        elif isinstance(implementation_attrs, list):
+            # NOTE: PyTorch schemas do not support a bytes type; we pass
+            # implementation_attrs as int[] representing raw bytes.
+            impl_list = [int(x) for x in implementation_attrs]
+        else:
+            raise TypeError(
+                "implementation_attrs must be None or list[int]; "
+                f"got {type(implementation_attrs)}"
+            )
+
+        attr = ts.TosaSerializerAttribute()
+        attr.CustomAttribute(
+            operator_name=operator_name,
+            domain_name=domain_name,
+            implementation_attrs=impl_list,
+        )
+
+        expanded = [TosaArg(item, self.tosa_spec) for item in inputs[0].special]
+        input_names = [arg.name for arg in expanded]
+        output_names = (
+            output.multiple_output_names
+            if getattr(output, "multiple_output_names", None)
+            else [output.name]
+        )
+        if len(output_names) != 1:
+            # TODO: Support multi-output CUSTOM ops with per-output meta/shape.
+            raise ValueError(
+                f"tosa.CUSTOM currently requires a single output, got {len(output_names)}"
+            )
+        self._serialize_operator(
+            node,
+            tosa_graph,
+            ts.Op.CUSTOM,
+            input_names,
+            output_names,
+            attr,
+        )

backends/arm/public_api_manifests/api_manifest_running.toml

@@ -56,6 +56,10 @@ signature = "EthosUPartitioner.ops_to_not_decompose(self, ep: torch.export.expor
 kind = "function"
 signature = "EthosUPartitioner.partition(self, exported_program: torch.export.exported_program.ExportedProgram) -> executorch.exir.backend.partitioner.PartitionResult"
 
+[python.EthosUPartitioner.register_custom_partition_op]
+kind = "function"
+signature = "EthosUPartitioner.register_custom_partition_op(self, op: torch._ops.OpOverload) -> None"
+
 [python.EthosUQuantizer]
 kind = "class"
 signature = "EthosUQuantizer(compile_spec: 'EthosUCompileSpec', use_composable_quantizer: 'bool' = False) -> 'None'"
@@ -136,6 +140,10 @@ signature = "VgfPartitioner.ops_to_not_decompose(self, ep: torch.export.exported
 kind = "function"
 signature = "VgfPartitioner.partition(self, exported_program: torch.export.exported_program.ExportedProgram) -> executorch.exir.backend.partitioner.PartitionResult"
 
+[python.VgfPartitioner.register_custom_partition_op]
+kind = "function"
+signature = "VgfPartitioner.register_custom_partition_op(self, op: torch._ops.OpOverload) -> None"
+
 [python.VgfQuantizer]
 kind = "class"
 signature = "VgfQuantizer(compile_spec: 'VgfCompileSpec', use_composable_quantizer: 'bool' = False) -> 'None'"

backends/arm/requirements-arm-tosa.txt

@@ -7,6 +7,5 @@ ml_dtypes == 0.5.1
 flatbuffers == 24.3.25
 tosa-adapter-model-explorer == 0.1.0
 ai-edge-model-explorer >= 0.1.16
-# NOTE: Will be removed when tosa-tools is installed via pypi
-pybind11 == 2.10.4
-pytest-timeout == 2.4.0
+pytest-timeout == 2.4.0
+tosa-tools == 2026.2.1

backends/arm/requirements-arm-vgf.txt

@@ -3,6 +3,6 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-ai_ml_emulation_layer_for_vulkan == 0.8.0
-ai_ml_sdk_model_converter == 0.8.0
-ai_ml_sdk_vgf_library == 0.8.0
+ai_ml_emulation_layer_for_vulkan == 0.9.0
+ai_ml_sdk_model_converter == 0.9.0
+ai_ml_sdk_vgf_library == 0.9.0

backends/arm/runtime/VGFBackend.cpp

@@ -157,7 +157,8 @@ class VGFBackend final : public ::executorch::runtime::BackendInterface {
     new (repr) VgfRepr(
         vk_instance, vk_physical_device, vk_device, vk_queue, vk_command_pool);
 
-    auto valid_vgf = repr->process_vgf(vgf_data, compile_specs);
+    auto valid_vgf =
+        repr->process_vgf(vgf_data, processed->size(), compile_specs);
     if (!valid_vgf) {
       ET_LOG(Error, "Failed to process VGF blob.");
       return Error::Internal;

backends/arm/runtime/VGFSetup.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright 2025 Arm Limited and/or its affiliates.
+ * Copyright 2025-2026 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.
@@ -324,26 +324,38 @@ static void debug_print_modules(
   }
 }
 
-bool VgfRepr::process_vgf(const char* vgf_data, ArrayRef<CompileSpec> specs) {
+bool VgfRepr::process_vgf(
+    const char* vgf_data,
+    size_t vgf_size,
+    ArrayRef<CompileSpec> specs) {
   ET_LOG(Info, "Preparing VGF as Vulkan objects");
 
   VkResult result;
 
   // Prepare temporary decoders
   unique_ptr<vgflib::HeaderDecoder> header_decoder =
-      vgflib::CreateHeaderDecoder(vgf_data);
+      vgflib::CreateHeaderDecoder(vgf_data, vgflib::HeaderSize(), vgf_size);
+  if (!header_decoder) {
+    ET_LOG(Error, "Failed to create VGF header decoder");
+    return false;
+  }
+
   unique_ptr<vgflib::ModelSequenceTableDecoder> sequence_decoder =
       vgflib::CreateModelSequenceTableDecoder(
-          vgf_data + header_decoder->GetModelSequenceTableOffset());
+          vgf_data + header_decoder->GetModelSequenceTableOffset(),
+          header_decoder->GetModelSequenceTableSize());
   unique_ptr<vgflib::ModuleTableDecoder> module_decoder =
       vgflib::CreateModuleTableDecoder(
-          vgf_data + header_decoder->GetModuleTableOffset());
+          vgf_data + header_decoder->GetModuleTableOffset(),
+          header_decoder->GetModuleTableSize());
   unique_ptr<vgflib::ModelResourceTableDecoder> resource_decoder =
       vgflib::CreateModelResourceTableDecoder(
-          vgf_data + header_decoder->GetModelResourceTableOffset());
+          vgf_data + header_decoder->GetModelResourceTableOffset(),
+          header_decoder->GetModelResourceTableSize());
   unique_ptr<vgflib::ConstantDecoder> constant_decoder =
       vgflib::CreateConstantDecoder(
-          vgf_data + header_decoder->GetConstantsOffset());
+          vgf_data + header_decoder->GetConstantsOffset(),
+          header_decoder->GetConstantsSize());
   // Check the VGF decoders
   if (not(header_decoder && module_decoder && sequence_decoder &&
           resource_decoder && constant_decoder && header_decoder->IsValid() &&

backends/arm/runtime/VGFSetup.h

@@ -1,5 +1,5 @@
 /*
- * Copyright 2025 Arm Limited and/or its affiliates.
+ * Copyright 2025-2026 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.
@@ -58,7 +58,10 @@ class VgfRepr {
   /*
    * Process a VGF ready for execution, allocate necessary Vulkan objects.
    */
-  bool process_vgf(const char* vgf_data, ArrayRef<CompileSpec> specs);
+  bool process_vgf(
+      const char* vgf_data,
+      size_t vgf_size,
+      ArrayRef<CompileSpec> specs);
 
   /*
    * Execute the VGF we've previously processed.

backends/arm/scripts/aot_arm_compiler.py

@@ -691,7 +691,7 @@ def _get_args():
     if args.evaluate is not None or args.evaluate_config is not None:
         logging.error(
             "Model evaluation is no longer supported in this script."
-            " Ignore and continue."
+            " Use evaluate_model.py instead. Ignore and continue."
         )
 
     return args