Merge branch 'main' into add-q5k-mlx

Author: JaynouOliver

backends/arm/runtime/VGFSetup.cpp

@@ -66,6 +66,7 @@ enum class FormatScalarKind {
   Uint,
   Sint,
   Float,
+  BFloat,
 };
 
 struct FormatInfo {
@@ -157,6 +158,7 @@ static uint32_t get_format_component_count(VkFormat format) {
     case VK_FORMAT_R16_UINT:
     case VK_FORMAT_R16_SINT:
     case VK_FORMAT_R16_SFLOAT:
+    case VK_FORMAT_R16_SFLOAT_FPENCODING_BFLOAT16_ARM:
     case VK_FORMAT_R32_UINT:
     case VK_FORMAT_R32_SINT:
     case VK_FORMAT_R32_SFLOAT:
@@ -209,6 +211,9 @@ static bool get_format_info(VkFormat format, FormatInfo* info) {
     case VK_FORMAT_R16_SFLOAT:
       *info = FormatInfo{1, 2, FormatScalarKind::Float};
       return true;
+    case VK_FORMAT_R16_SFLOAT_FPENCODING_BFLOAT16_ARM:
+      *info = FormatInfo{1, 2, FormatScalarKind::BFloat};
+      return true;
     case VK_FORMAT_R32_UINT:
       *info = FormatInfo{1, 4, FormatScalarKind::Uint};
       return true;
@@ -3696,6 +3701,7 @@ static uint32_t get_format_size(VkFormat format) {
     case VK_FORMAT_R16_UINT:
     case VK_FORMAT_R16_SINT:
     case VK_FORMAT_R16_SFLOAT:
+    case VK_FORMAT_R16_SFLOAT_FPENCODING_BFLOAT16_ARM:
     case VK_FORMAT_R8G8_UINT:
     case VK_FORMAT_R8G8_SINT:
       return 2;

backends/arm/scripts/vulkan_utils.sh

@@ -26,21 +26,21 @@ vulkan_sdk_arch="${ARCH}"
 # macOS and Linux x86_64 use the official LunarG SDK tarballs. Linux ARM64
 # uses a separately repackaged mirror of the same SDK version.
 if [[ "${os_name}" == "Darwin" ]]; then
-    vulkan_sdk_version="1.4.341.1"
+    vulkan_sdk_version="1.4.350.0"
     vulkan_sdk_arch="macOS"
     vulkan_sdk_url="https://sdk.lunarg.com/sdk/download/${vulkan_sdk_version}/mac/vulkansdk-macos-${vulkan_sdk_version}.zip"
-    vulkan_sdk_sha256="632cbe96c8ed6ed00c6ce25e3a7738c466134f76586e1c51f1419410d7f9042e"
+    vulkan_sdk_sha256="7acc181b8fd9b4781bf51ed086222ec95d22004b85b3d0a6683a7e48ca5a1679"
 elif [[ "${os_name}" == "Linux" ]] && [[ "${ARCH}" == "x86_64" ]]; then
-    vulkan_sdk_version="1.4.341.1"
+    vulkan_sdk_version="1.4.350.0"
     vulkan_sdk_url="https://sdk.lunarg.com/sdk/download/${vulkan_sdk_version}/linux/vulkansdk-linux-x86_64-${vulkan_sdk_version}.tar.xz"
-    vulkan_sdk_sha256="3bf0f762afb6c79bc6a9d9fb5998745ccff928800a29619b501ed9de7fd9789b"
+    vulkan_sdk_sha256="b65f068ab36263559da49d7cacd7e7b9df23824ca8b68ccc522a2b06f5725df2"
 elif [[ "${os_name}" == "Linux" ]] && ([[ "${ARCH}" == "aarch64" ]] || [[ "${ARCH}" == "arm64" ]]); then
-    vulkan_sdk_version="1.4.341.1"
+    vulkan_sdk_version="1.4.350.0"
     if [[ "${vulkan_sdk_arch}" == "arm64" ]]; then
         vulkan_sdk_arch="aarch64"
     fi
     vulkan_sdk_url="https://github.com/jakoch/vulkan-sdk-arm/releases/download/${vulkan_sdk_version}/vulkansdk-ubuntu-22.04-arm-${vulkan_sdk_version}.tar.xz"
-    vulkan_sdk_sha256="345312aee2c835e128b30653278593f899a659a7ba287c571cafb22acb708b8f"
+    vulkan_sdk_sha256="9e403d444219bb7c17e9231b580d704453e2afa30a1c2fdd568d1776dc68790b"
 else
     log_step "vulkan" "Error: only macOS and Linux are supported (detected ${os_name}); architecture must be x86-64 or aarch64/arm64"
     exit 1

backends/arm/test/ops/test_matmul.py

@@ -455,7 +455,7 @@ def test_matmul_u85_INT(test_case: test_case_t):
     pipeline.run()
 
 
-@common.parametrize("test_case", test_suite | test_suite_fp16)
+@common.parametrize("test_case", test_suite | test_suite_fp16 | test_suite_bf16)
 @common.SkipIfNoModelConverter
 def test_matmul_vgf_no_quant(test_case: test_case_t):
     test_data = test_case()

backends/arm/test/runner_utils.py

@@ -243,8 +243,7 @@ def is_concrete_shape(shape_like) -> bool:
         return all(isinstance(dim, numbers.Integral) for dim in shape_like)
 
     def to_torch_tensor() -> torch.Tensor:
-        if array.dtype.type is np.void:
-            # If dtype is void, "cheat" and use the output_tensor dtype.
+        if output_tensor.dtype == torch.bfloat16 or array.dtype.type is np.void:
             return torch.frombuffer(array, dtype=output_tensor.dtype)
         return torch.from_numpy(array)
 

backends/arm/test/tester/test_pipeline.py

@@ -1258,13 +1258,15 @@ def __init__(
     ):
         if tosa_spec is None:
             if tosa_version is None:
-                tosa_spec = VgfCompileSpec().tosa_spec
-            else:
-                if tosa_extensions is None:
+                tosa_version = str(VgfCompileSpec().tosa_spec)
+            if tosa_extensions is None:
+                if "FP" in tosa_version:
+                    tosa_extensions = ["bf16"]
+                else:
                     tosa_extensions = []
-                tosa_spec = TosaSpecification.create_from_string(
-                    tosa_version + "".join([f"+{ext}" for ext in tosa_extensions])
-                )
+            tosa_spec = TosaSpecification.create_from_string(
+                tosa_version + "".join([f"+{ext}" for ext in tosa_extensions])
+            )
         elif isinstance(tosa_spec, str):
             tosa_spec = TosaSpecification.create_from_string(tosa_spec)
         compile_spec = common.get_vgf_compile_spec(

backends/nxp/backend/edge_helper.py

@@ -109,6 +109,43 @@ def node_is_effectively_static_tensor(
     )
 
 
+def weights_are_effectively_static(
+    node: Node, parameters_mapping: dict[str, Parameter], weight_index: int = 1
+) -> bool:
+    """Neutron IR sometimes requires some weights to be static. This method checks if this is the case for the
+        provided `node`.
+
+    Sometimes a `permute_copy` is inserted to transpose the weights during edge lowering. The `permute_copy` is
+    then removed during conversion to Neutron IR if it transposes static data. In those cases, the weights will be
+    static. Therefore, it is ok if the weights are produced by a `permute_copy` with a static input.
+
+    :param node: Tensor node to check for data.
+    :param parameters_mapping: Dict mapping tensor names to their static data. Should be inferred from the
+                                `state_dict` attribute of an edge program.
+    :param weight_index: Index to the `node.args` where the weight is located. Defaults to 1.
+    :return: True if the weight at the given index is effectively static.
+    """
+
+    def _is_permute_copy(node_: Node) -> bool:
+        return hasattr(node_, "target") and node_.target == PermuteCopy
+
+    if (
+        _is_dequantize(dq_node := node.args[weight_index])
+        and _is_quantize(q_node := dq_node.args[0])
+        and _is_permute_copy(permute_copy_node := q_node.args[0])
+    ):
+        # The weights are produced by a `permute_copy`. Its input (the weights) must be static.
+        return node_is_effectively_static_tensor(
+            permute_copy_node.args[0], parameters_mapping
+        )
+
+    else:
+        # There is no `permute_copy`. The weights must be static directly.
+        return node_is_effectively_static_tensor(
+            node.args[weight_index], parameters_mapping
+        )
+
+
 def try_get_tensor_constant_from_node(
     graph_module: GraphModule, node: Node
 ) -> Parameter | None:

backends/nxp/backend/edge_program_converter.py

@@ -113,6 +113,9 @@ def convert_program(
         self._convert_qdq_cluster_q_dq_nodes(edge_program.graph.nodes, cc)
         self._process_nodes(edge_program.graph.nodes, cc)
 
+        # Make sure the operators are correctly ordered.
+        cc.tflite_builder.sort_operators_topologically()
+
         # Assign the model its inputs and outputs.
         cc.tflite_builder.assign_model_io_to_subgraph(edge_program.graph_signature)
 

backends/nxp/backend/ir/converter/builder/model_builder.py

@@ -6,6 +6,8 @@
 # See the LICENSE_MIT for more details.
 #
 
+import logging
+from collections import defaultdict, deque
 from copy import deepcopy
 from typing import List, Optional, Union
 
@@ -111,6 +113,67 @@ def __init__(
         self._zeros_tensor_map = {}
         self.edge_to_tflite_map = {}
 
+    def sort_operators_topologically(self):
+        """Sort the operators in the model graph in topological order using Kahn's algorithm.
+
+        Ensures that each operator appears after all operators that produce its input tensors.
+        This is required for correct behavior of optimization passes and by the Neutron IR format.
+
+        The algorithm works in three steps:
+            1. Build a mapping from each tensor to the operator that produces it.
+            2. Compute the number of unresolved input dependencies for each operator.
+            3. Iteratively add operators with no remaining dependencies to the sorted list,
+               decrementing the number of dependencies of their consumers until all operators are processed.
+
+        If the sort succeeds, the operator list in the model's subgraph is replaced with the topologically sorted list.
+        If it fails (e.g. due to an invalid graph), a warning is logged and the original operator order is preserved.
+        """
+
+        operators = self.get_operators()
+
+        # Map tensors to operators that produce them.
+        tensor_to_producer_op = {}
+        for producer_op in operators:
+            for output in producer_op.tmp_outputs:
+                tensor_to_producer_op[output] = producer_op
+
+        # Map operators to operators that consume their outputs.
+        op_to_child_op = defaultdict(list)
+        num_unresolved_input_dependencies = {op: 0 for op in operators}
+        for child_op in operators:
+            for input_ in child_op.tmp_inputs:
+                parent_op = tensor_to_producer_op.get(input_)
+                if parent_op is not None:
+                    op_to_child_op[parent_op].append(child_op)
+                    num_unresolved_input_dependencies[child_op] += 1
+
+        # Gradually build the graph.
+        # Use `deque` for fast appends and pops.
+        queue = deque(
+            op
+            for op, degree in num_unresolved_input_dependencies.items()
+            if degree == 0
+        )
+        sorted_ops = []
+        while len(queue) != 0:
+            op = queue.popleft()
+            sorted_ops.append(op)
+            for child_op in op_to_child_op[op]:
+                num_unresolved_input_dependencies[child_op] -= 1
+                if num_unresolved_input_dependencies[child_op] == 0:
+                    # All input of the `child_op` are produced by operators that are already in `sorted_ops`. So we can
+                    #  insert it into the graph.
+                    queue.append(child_op)
+
+        if len(sorted_ops) != operators.len():
+            logging.warning(
+                "NXP backend: ModelBuilder.sort_operators_topologically() failed. Please report this."
+            )
+
+        else:
+            # The topological sort was successful.
+            operators.vector = sorted_ops
+
     def create_zeros_tensor(
         self, dims: List[int], name: str, dtype: np.dtype, can_reuse: bool = False
     ) -> tflite_model.Tensor:

backends/nxp/backend/ir/converter/node_converter.py

@@ -253,7 +253,9 @@ def _has_shared_q_params_if_quantized(node: Node) -> bool:
 
     @staticmethod
     def is_node_alone_in_partition(
-        node: Node, partition_list: list[Partition], filter_fn: Callable[[Node], bool]
+        node: Node,
+        partition_list: list[Partition],
+        filter_fn: Callable[[Node], bool] = is_not_qdq_node,
     ) -> bool:
         """Return True if `node` is the only node in its partition for which `filter_fn`
         returns True.
@@ -264,9 +266,8 @@ def is_node_alone_in_partition(
 
         :param node: The torch.fx.Node to check.
         :param partition_list: List of proposed partitions.
-        :param filter_fn: Predicate applied to nodes in the partition.
-                        `node` is considered alone if it is the only node
-                        for which this predicate returns True.
+        :param filter_fn: Predicate applied to nodes in the partition. `node` is considered alone if it is the only node
+                           for which this predicate returns True. By default, Q/Dq nodes are ignored.
         """
         partitions = [p for p in partition_list if node in p.nodes]
         if len(partitions) != 1:

backends/nxp/backend/ir/converter/node_converters/ops_converters/addmm_converter.py

@@ -1,9 +1,15 @@
-# Copyright 2024-2025 NXP
+# Copyright 2024-2026 NXP
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from executorch.backends.nxp.backend.edge_helper import input_rank
+import torch
+
+from executorch.backends.nxp.backend.edge_helper import (
+    input_rank,
+    node_is_effectively_static_tensor,
+    weights_are_effectively_static,
+)
 from executorch.backends.nxp.backend.ir.converter.conversion.common import OpsList
 from executorch.backends.nxp.backend.ir.converter.node_converter import (
     CustomDelegationOptions,
@@ -12,10 +18,18 @@
 from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options import (
     fully_connected_options,
 )
+
+from executorch.backends.nxp.backend.neutron_target_spec import NeutronTargetSpec
 from torch.fx import Node
 from torch.nn import Parameter
 
 
+# The edge operator signature is: aten.addmm(bias, input, weight, *, beta=1, alpha=1)
+MAIN_INPUT_IDX = 1
+WEIGHT_IDX = 2
+BIAS_IDX = 0
+
+
 class AddMMConverter(NodeConverter):
     """Convert the `aten.addmm` operator to TFLite `FullyConnected` with a bias input."""
 
@@ -29,12 +43,76 @@ def _is_supported_in_IR(
             return False
 
         # The weights must be 2D.
-        if input_rank(node, 2) != 2:
+        if input_rank(node, WEIGHT_IDX) != 2:
+            return False
+
+        alpha, beta = node.kwargs.get("alpha", 1), node.kwargs.get("beta", 1)
+        if alpha != 1 or beta != 1:
+            # As these cases seem rare, conversion is not implemented for the time being.
+            return False
+
+        # The `aten.addmm` operator allows any bias shape, as long as it is broadcastable with the result of the matrix
+        #  multiplication. That means it supports 4 different shapes: [N, P], [1, P], [P], [1] (provided the MM result
+        #  has shape [N, P]). Out of these 4, Neutron IR allows only [1, P] and [P], both of which are supported on
+        #  Neutron.
+        bias_shape = list(node.args[BIAS_IDX].meta["val"].shape)
+        _, p = node.meta["val"].shape
+        if bias_shape not in [[1, p], [p]]:
+            return False
+
+        return True
+
+    @staticmethod
+    def _is_supported_on_target(
+        node: Node,
+        neutron_target_spec: NeutronTargetSpec,
+        parameters_mapping: dict[str, Parameter],
+        custom_delegation_options: CustomDelegationOptions,
+    ) -> bool:
+        # Main input and output must be `int8` or `uint8`.
+        if not NodeConverter.uses_quantization_type_for_io(
+            node, [torch.int8, torch.uint8], [MAIN_INPUT_IDX], [0]
+        ):
+            return False
+
+        # Weights must be `int8`.
+        if not NodeConverter.uses_quantization_type_for_io(
+            node, [torch.int8], [WEIGHT_IDX], []
+        ):
+            return False
+
+        # Bias must be `int32`.
+        if not NodeConverter.uses_quantization_type_for_io(
+            node, [torch.int32], [BIAS_IDX], []
+        ):
+            return False
+
+        # Weights must be constant.
+        if not weights_are_effectively_static(
+            node, parameters_mapping, weight_index=WEIGHT_IDX
+        ):
+            return False
+
+        # The bias must be constant.
+        if not node_is_effectively_static_tensor(
+            node.args[BIAS_IDX], parameters_mapping
+        ):
             return False
 
         return True
 
     def convert(self, node: Node):
+        """Convert the `aten.addmm` operator to NeutronIR `FullyConnected`.
+        The schema is:
+            addmm(
+                Tensor self,
+                Tensor mat1,
+                Tensor mat2,
+                *,
+                Scalar beta=1,
+                Scalar alpha=1
+            ) -> Tensor
+        """
         self.assert_convertible(node)
 
         t_op = self._create_tflite_op_with_io_tensors(node)
@@ -47,14 +125,14 @@ def convert(self, node: Node):
         w = t_op.tmp_inputs[2]
         y = t_op.tmp_outputs[0]
 
-        # Assign the operator its TFLite inputs and outputs
+        # Assign the operator its Neutron IR inputs and outputs
         t_op.tmp_inputs = [x, w, bias]
         t_op.tmp_outputs = [y]
 
         ops = OpsList(middle_op=t_op)
 
         # The `aten.addmm` uses main input with shape [M, N] and the weights have the shape [N, O].
-        # TFLite `FullyConnected` requires the weights to have shape [O, N] (if the main input has shape [M, N]).
+        # Neutron IR `FullyConnected` requires the weights to have shape [O, N] (if the main input has shape [M, N]).
         # Insert a `Transpose` operator to permute the weights to achieve correct conversion. (The `Transpose` will not
         #  be present in the output model if the weights are static.)
         ops.add_pre(self.builder.create_transpose_operator_before(t_op, 1, [1, 0]))