Arm backend: Initial bfloat16 support.

backends/arm/_passes/decompose_avg_pool2d_pass.py

@@ -51,7 +51,7 @@ def call_operator(self, op, args, kwargs, meta):
         full_op, cat_op, avgpool_op, mul_op = get_decomposition(op)
 
         x = args[0]
-        full_kwargs = {"device": x.data.device}
+        full_kwargs = {"device": x.data.device, "dtype": x.data.dtype}
         kernel_h, kernel_w = args[1]
         kernel_size = kernel_h * kernel_w
         if len(args) > 2 and args[2] is not None:

backends/arm/_passes/fuse_equal_placeholders_pass.py

@@ -54,15 +54,15 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
             # ensure we don't merge any special case int48_t tensors with int32_t tensors
             # since int48_t tensors needs to be instantiated separately.
             is_special_dtype = node.meta.get(TosaSpecialDtype.meta_key(), None)
-            t_cpu = tensor.detach().cpu().contiguous()
+            t_cpu = tensor.cpu().contiguous().flatten().view(dtype=torch.uint8)
             data_bytes = t_cpu.numpy().tobytes()
             key = (
                 is_special_dtype,
-                str(t_cpu.dtype),
-                tuple(t_cpu.shape),
+                str(tensor.dtype),
+                tuple(tensor.shape),
                 hashlib.sha1(data_bytes, usedforsecurity=False).hexdigest(),
             )
-            hash_buckets[key].append((node, t_cpu))
+            hash_buckets[key].append((node, tensor))
 
         # For each bucket with more than one entry, fuse:
         for nodes_tensors in hash_buckets.values():

backends/arm/_passes/match_arg_dtype_pass.py

@@ -1,4 +1,4 @@
-# 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.
@@ -19,8 +19,9 @@
     torch.int32: 4,
     torch.int64: 5,
     torch.float16: 6,
-    torch.float32: 7,
-    torch.float64: 8,
+    torch.bfloat16: 7,
+    torch.float32: 8,
+    torch.float64: 9,
 }
 
 

backends/arm/operators/op_abs.py

@@ -1,4 +1,4 @@
-# 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.
@@ -43,8 +43,8 @@ def define_node(
         validate_valid_dtype(
             self.target,
             [*inputs, output],
-            [ts.DType.INT32, ts.DType.FP32],
-            output.tosa_spec,
+            [ts.DType.INT32, ts.DType.FP32, ts.DType.BF16],
+            self.tosa_spec,
         )
 
         attr = ts.TosaSerializerAttribute()

backends/arm/operators/op_add.py

@@ -1,4 +1,4 @@
-# Copyright 2023-2025 Arm Limited and/or its affiliates.
+# Copyright 2023-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.
@@ -43,8 +43,8 @@ def define_node(
         validate_valid_dtype(
             self.target,
             [*inputs, output],
-            [ts.DType.INT32, ts.DType.FP32],
-            output.tosa_spec,
+            [ts.DType.INT32, ts.DType.FP32, ts.DType.BF16],
+            self.tosa_spec,
         )
 
         attr = ts.TosaSerializerAttribute()

backends/arm/operators/op_amax.py

@@ -1,4 +1,4 @@
-# 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.
@@ -42,7 +42,7 @@ def define_node(
             self.target,
             [inputs[0], output],
             [ts.DType.INT8, ts.DType.INT16, ts.DType.INT32, ts.DType.FP32],
-            output.tosa_spec,
+            self.tosa_spec,
         )
 
         input = inputs[0]

backends/arm/operators/op_amin.py

@@ -1,4 +1,4 @@
-# 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.
@@ -42,7 +42,7 @@ def define_node(
             self.target,
             [inputs[0], output],
             [ts.DType.INT8, ts.DType.INT16, ts.DType.INT32, ts.DType.FP32],
-            output.tosa_spec,
+            self.tosa_spec,
         )
 
         input = inputs[0]

backends/arm/operators/op_any.py

@@ -1,4 +1,4 @@
-# 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.
@@ -37,7 +37,7 @@ def define_node(
         validate_num_inputs(self.target, inputs, 3)
         validate_same_dtype(self.target, [inputs[0], output], ts)
         validate_valid_dtype(
-            self.target, [inputs[0], output], ts.DType.BOOL, output.tosa_spec
+            self.target, [inputs[0], output], ts.DType.BOOL, self.tosa_spec
         )
 
         input_shape = list(inputs[0].shape)

backends/arm/operators/op_avg_pool2d.py

@@ -1,4 +1,4 @@
-# Copyright 2023-2025 Arm Limited and/or its affiliates.
+# Copyright 2023-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.
@@ -115,14 +115,14 @@ def define_node(
     ) -> None:
         validate_num_inputs(self.target, inputs, [3, 4, 5, 6, 7])
         validate_same_dtype(self.target, [inputs[0], output], ts)
-        supported_dtypes = [ts.DType.INT8, ts.DType.FP32]
+        supported_dtypes = [ts.DType.INT8, ts.DType.FP32, ts.DType.BF16]
         if self.tosa_spec.support_extension("int16"):
             supported_dtypes.append(ts.DType.INT16)
         validate_valid_dtype(
             self.target,
             [inputs[0], output],
             supported_dtypes,
-            output.tosa_spec,
+            self.tosa_spec,
         )
 
         if inputs[0].dtype == ts.DType.INT8 or inputs[0].dtype == ts.DType.INT16:

backends/arm/operators/op_bitwise_not.py

@@ -1,4 +1,4 @@
-# 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.
@@ -46,7 +46,7 @@ def define_node(
             self.target,
             [*inputs, output],
             [ts.DType.INT8, ts.DType.INT16, ts.DType.INT32],
-            output.tosa_spec,
+            self.tosa_spec,
         )
 
         attr = ts.TosaSerializerAttribute()

backends/arm/operators/op_cat.py

@@ -1,4 +1,4 @@
-# Copyright 2024-2025 Arm Limited and/or its affiliates.
+# Copyright 2024-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.
@@ -37,17 +37,23 @@ def define_node(
         inputs: List[TosaArg],
         output: TosaArg,
     ) -> None:
-        supported_dtypes = [ts.DType.BOOL, ts.DType.INT8, ts.DType.INT32, ts.DType.FP32]
+        supported_dtypes = [
+            ts.DType.BOOL,
+            ts.DType.INT8,
+            ts.DType.INT32,
+            ts.DType.FP32,
+            ts.DType.BF16,
+        ]
         if self.tosa_spec.support_extension("int16"):
             supported_dtypes.append(ts.DType.INT16)
         validate_num_inputs(self.target, inputs, [1, 2])
-        input_tosa_args = [TosaArg(arg, output.tosa_spec) for arg in inputs[0].special]
+        input_tosa_args = [TosaArg(arg, self.tosa_spec) for arg in inputs[0].special]
         validate_same_dtype(self.target, [*input_tosa_args, output], ts)
         validate_valid_dtype(
             self.target,
             [*input_tosa_args, output],
             supported_dtypes,
-            output.tosa_spec,
+            self.tosa_spec,
         )
 
         dim = 0 if len(inputs) < 2 else inputs[1].number

backends/arm/operators/op_ceil.py

@@ -1,4 +1,4 @@
-# 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.
@@ -45,8 +45,8 @@ def define_node(
         validate_valid_dtype(
             self.target,
             inputs[0],
-            ts.DType.FP32,
-            output.tosa_spec,
+            [ts.DType.FP32, ts.DType.BF16],
+            self.tosa_spec,
         )
 
         attr = ts.TosaSerializerAttribute()

backends/arm/operators/op_clamp.py

@@ -1,12 +1,11 @@
-# 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
 
 
 from typing import Any, List, Tuple
 
-import numpy as np
 import torch
 import tosa_serializer as ts
 
@@ -67,16 +66,7 @@ def cast_type(value: Any) -> int | float:
         return min_arg, max_arg
 
     def _to_bytes(self, value: int | float, dtype: torch.dtype) -> bytes:
-        if dtype == torch.float32:
-            return np.frombuffer(np.float32(value).tobytes(), dtype=np.uint8).tolist()
-        elif dtype == torch.float16:
-            return np.frombuffer(np.float16(value).tobytes(), dtype=np.uint8).tolist()
-        elif dtype == torch.int8:
-            return np.frombuffer(np.int8(value).tobytes(), dtype=np.uint8).tolist()
-        elif dtype == torch.int16:
-            return np.frombuffer(np.int16(value).tobytes(), dtype=np.uint8).tolist()
-        else:
-            raise ValueError(f"Unsupported dtype for to_bytes: {dtype}")
+        return torch.full((1,), value, dtype=dtype).view(torch.uint8).numpy().tolist()
 
     def define_node(
         self,
@@ -87,14 +77,19 @@ def define_node(
     ) -> None:
         validate_num_inputs(self.target, inputs, [2, 3])
         validate_same_dtype(self.target, [inputs[0], output], ts)
-        supported_dtypes = [ts.DType.INT8, ts.DType.FP16, ts.DType.FP32]
+        supported_dtypes = [
+            ts.DType.INT8,
+            ts.DType.FP16,
+            ts.DType.BF16,
+            ts.DType.FP32,
+        ]
         if self.tosa_spec.support_extension("int16"):
             supported_dtypes.append(ts.DType.INT16)
         validate_valid_dtype(
             self.target,
             [inputs[0], output],
             supported_dtypes,
-            output.tosa_spec,
+            self.tosa_spec,
         )
 
         node_input_dtype = node.meta["val"].dtype

backends/arm/operators/op_constant_pad_nd.py

@@ -1,4 +1,4 @@
-# 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.
@@ -55,7 +55,7 @@ def define_node(
                 ts.DType.FP32,
                 ts.DType.BOOL,
             ],
-            output.tosa_spec,
+            self.tosa_spec,
         )
 
         if inputs[0].dtype == ts.DType.INT8:

backends/arm/operators/op_cos.py

@@ -1,4 +1,4 @@
-# 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.
@@ -40,7 +40,7 @@ def define_node(
         validate_num_inputs(self.target, inputs, 1)
         validate_same_dtype(self.target, [*inputs, output], ts)
         validate_valid_dtype(
-            self.target, [*inputs, output], ts.DType.FP32, output.tosa_spec
+            self.target, [*inputs, output], ts.DType.FP32, self.tosa_spec
         )
         attr = ts.TosaSerializerAttribute()
         attr.CosAttribute()

backends/arm/operators/op_eq.py

@@ -1,4 +1,4 @@
-# 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.
@@ -48,9 +48,9 @@ def define_node(
             self.target,
             inputs,
             [ts.DType.INT32, ts.DType.FP32],
-            output.tosa_spec,
+            self.tosa_spec,
         )
-        validate_valid_dtype(self.target, output, ts.DType.BOOL, output.tosa_spec)
+        validate_valid_dtype(self.target, output, ts.DType.BOOL, self.tosa_spec)
 
         attr = ts.TosaSerializerAttribute()
         attr.EqualAttribute()

backends/arm/operators/op_erf.py

@@ -1,4 +1,4 @@
-# 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.
@@ -43,7 +43,7 @@ def define_node(
             self.target,
             [*inputs, output],
             ts.DType.FP32,
-            output.tosa_spec,
+            self.tosa_spec,
         )
 
         # MI lowering

backends/arm/operators/op_exp.py

@@ -1,4 +1,4 @@
-# Copyright 2024-2025 Arm Limited and/or its affiliates.
+# Copyright 2024-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.
@@ -44,7 +44,7 @@ def define_node(
             self.target,
             [*inputs, output],
             ts.DType.FP32,
-            output.tosa_spec,
+            self.tosa_spec,
         )
 
         attr = ts.TosaSerializerAttribute()

backends/arm/operators/op_floor.py

@@ -1,4 +1,4 @@
-# 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.
@@ -46,7 +46,7 @@ def define_node(
             self.target,
             inputs[0],
             ts.DType.FP32,
-            output.tosa_spec,
+            self.tosa_spec,
         )
 
         attr = ts.TosaSerializerAttribute()

backends/arm/operators/op_ge.py

@@ -1,4 +1,4 @@
-# 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.
@@ -48,9 +48,9 @@ def define_node(
             self.target,
             inputs,
             [ts.DType.INT32, ts.DType.FP32],
-            output.tosa_spec,
+            self.tosa_spec,
         )
-        validate_valid_dtype(self.target, output, ts.DType.BOOL, output.tosa_spec)
+        validate_valid_dtype(self.target, output, ts.DType.BOOL, self.tosa_spec)
 
         attr = ts.TosaSerializerAttribute()
         attr.GreaterEqualAttribute()

backends/arm/operators/op_gt.py

@@ -1,4 +1,4 @@
-# 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.
@@ -48,9 +48,9 @@ def define_node(
             self.target,
             inputs,
             [ts.DType.INT32, ts.DType.FP32],
-            output.tosa_spec,
+            self.tosa_spec,
         )
-        validate_valid_dtype(self.target, output, ts.DType.BOOL, output.tosa_spec)
+        validate_valid_dtype(self.target, output, ts.DType.BOOL, self.tosa_spec)
 
         attr = ts.TosaSerializerAttribute()
         attr.GreaterAttribute()