Arm backend: Add BF16 support to operators pt. 4-5 (pytorch#17003)

Erik-Lundell · web-flow · commit ec4c4629b92b · 2026-01-30T11:05:23.000+01:00
ge, gt, alias_copy, index_tensor mapping to TOSA operators GATHER, GREATER, GREATER_EQUAL, IDENTITY log, maxpool_2d, log, matmul, maximum, minimum mapping to TOSA operators LOG, MATMUL, MAXIMUM, MINIMUM, MAX_POOL2D cc @freddan80 @per @zingo @oscarandersson8218 @digantdesai --------- Signed-off-by: Erik Lundell <erik.lundell@arm.com>
diff --git a/backends/arm/_passes/rewrite_matmul.py b/backends/arm/_passes/rewrite_matmul.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.
@@ -91,6 +91,22 @@ def call(self, graph_module):
                     tosa_matmul_node.meta[TosaSpecialDtype.meta_key()] = (
                         TosaSpecialDtype.INT48
                     )
+            elif (
+                x1_fake_tensor.dtype == torch.bfloat16
+                and x2_fake_tensor.dtype == torch.bfloat16
+                and output_fake_tensor.dtype != torch.bfloat16
+            ):
+                # A TOSA BF16 MATMUL outputs FP32 wheras pytorch outputs BF16.
+                # Cast back to BF16 to get matching semantics.
+                with graph_module.graph.inserting_after(tosa_matmul_node):
+                    cast_node = create_node(
+                        graph_module.graph,
+                        op_target=exir_ops.edge.dim_order_ops._to_dim_order_copy.default,
+                        kwargs={"dtype": torch.bfloat16},
+                        from_node=tosa_matmul_node,
+                    )
+                    tosa_matmul_node.replace_all_uses_with(cast_node)
+                    cast_node.args = (tosa_matmul_node,)
 
         if modified:
             graph_module.recompile()
diff --git a/backends/arm/operators/op_ge.py b/backends/arm/operators/op_ge.py
@@ -41,7 +41,7 @@ def define_node(
         validate_valid_dtype(
             self.target,
             inputs,
-            [ts.DType.INT32, ts.DType.FP32],
+            [ts.DType.INT32, ts.DType.FP32, ts.DType.BF16],
             self.tosa_spec,
         )
         validate_valid_dtype(self.target, output, ts.DType.BOOL, self.tosa_spec)
diff --git a/backends/arm/operators/op_gt.py b/backends/arm/operators/op_gt.py
@@ -41,7 +41,7 @@ def define_node(
         validate_valid_dtype(
             self.target,
             inputs,
-            [ts.DType.INT32, ts.DType.FP32],
+            [ts.DType.INT32, ts.DType.FP32, ts.DType.BF16],
             self.tosa_spec,
         )
         validate_valid_dtype(self.target, output, ts.DType.BOOL, self.tosa_spec)
diff --git a/backends/arm/operators/op_log.py b/backends/arm/operators/op_log.py
@@ -41,7 +41,10 @@ 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, self.tosa_spec
+            self.target,
+            [*inputs, output],
+            [ts.DType.FP32, ts.DType.BF16],
+            self.tosa_spec,
         )
         attr = ts.TosaSerializerAttribute()
         attr.LogAttribute()
diff --git a/backends/arm/operators/op_max_pool2d.py b/backends/arm/operators/op_max_pool2d.py
@@ -38,7 +38,7 @@ def define_node(
     ) -> None:
         validate_num_inputs(self.target, inputs, [3, 4, 5, 6])
         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(
diff --git a/backends/arm/operators/op_mul.py b/backends/arm/operators/op_mul.py
@@ -38,7 +38,13 @@ def define_node(
         validate_valid_dtype(
             self.target,
             [*inputs, output],
-            [ts.DType.INT8, ts.DType.INT16, ts.DType.INT32, ts.DType.FP32],
+            [
+                ts.DType.INT8,
+                ts.DType.INT16,
+                ts.DType.INT32,
+                ts.DType.FP32,
+                ts.DType.BF16,
+            ],
             self.tosa_spec,
         )
 
diff --git a/backends/arm/operators/op_tosa_matmul.py b/backends/arm/operators/op_tosa_matmul.py
@@ -46,6 +46,8 @@ def define_node(
         validate_num_inputs(self.target, inputs, 2)
         validate_same_dtype(self.target, [*inputs], ts)
         supported_input_dtypes = [ts.DType.INT8, ts.DType.INT32, ts.DType.FP32]
+        if self.tosa_spec.support_extension("bf16"):
+            supported_input_dtypes.append(ts.DType.BF16)
         if self.tosa_spec.support_extension("int16"):
             supported_input_dtypes.append(ts.DType.INT16)
         validate_valid_dtype(
diff --git a/backends/arm/operators/op_view.py b/backends/arm/operators/op_view.py
@@ -46,6 +46,7 @@ def define_node(
                 ts.DType.INT16,
                 ts.DType.INT32,
                 ts.DType.FP32,
+                ts.DType.BF16,
                 ts.DType.BOOL,
             ],
             self.tosa_spec,
diff --git a/backends/arm/operators/ops_identity.py b/backends/arm/operators/ops_identity.py
@@ -49,6 +49,8 @@ def define_node(
             ]
             if self.tosa_spec.support_float():
                 supported_dtypes += [ts.DType.FP32]
+            if self.tosa_spec.support_extension("bf16"):
+                supported_dtypes += [ts.DType.BF16]
             if self.tosa_spec.support_extension("int16"):
                 supported_dtypes += [ts.DType.INT48]
             if self.tosa_spec.support_extension("int4"):
diff --git a/backends/arm/test/ops/test_alias_copy.py b/backends/arm/test/ops/test_alias_copy.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.
@@ -36,6 +36,9 @@ class AliasCopy(torch.nn.Module):
         "3d_rand": lambda: (torch.rand(3, 5, 5),),
         "4d_zeros": lambda: (torch.zeros(1, 10, 10, 10),),
     }
+    test_data_bf16 = {
+        "3d_rand_bf16": lambda: (torch.rand(3, 5, 2, dtype=torch.bfloat16),)
+    }
 
     def __init__(self):
         super().__init__()
@@ -46,13 +49,14 @@ def forward(self, x: torch.Tensor):
         )  # Multiply by one to make sure it is partitioned.
 
 
-@common.parametrize("test_data", AliasCopy.test_data)
+@common.parametrize("test_data", AliasCopy.test_data | AliasCopy.test_data_bf16)
 def test_alias_tosa_FP(test_data: input_t1):
     TosaPipelineFP[input_t1](
         AliasCopy(),
         test_data(),
         AliasCopy.aten_op,
         AliasCopy.exir_op,
+        tosa_extensions=["bf16"],
     ).run()
 
 
diff --git a/backends/arm/test/ops/test_ge.py b/backends/arm/test/ops/test_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.
@@ -69,32 +69,54 @@ def get_inputs(self):
     "ge_tensor_rank4_randn": lambda: op_ge_tensor_rank4_randn,
 }
 
+test_data_tensor_bf16 = {
+    "ge_tensor_rank2_rand_bf16": lambda: GreaterEqual(
+        torch.rand(4, 5, dtype=torch.bfloat16),
+        torch.rand(1, 5, dtype=torch.bfloat16),
+    ),
+    "ge_tensor_rank4_randn_bf16": lambda: GreaterEqual(
+        torch.randn(2, 3, 4, 2, dtype=torch.bfloat16),
+        torch.randn(2, 3, 4, 1, dtype=torch.bfloat16),
+    ),
+}
+
 test_data_scalar = {
     "ge_scalar_rank1_ones": lambda: op_ge_scalar_rank1_ones,
     "ge_scalar_rank2_rand": lambda: op_ge_scalar_rank2_rand,
     "ge_scalar_rank3_randn": lambda: op_ge_scalar_rank3_randn,
     "ge_scalar_rank4_randn": lambda: op_ge_scalar_rank4_randn,
 }
 
+test_data_scalar_bf16 = {
+    "ge_scalar_rank2_rand_bf16": lambda: GreaterEqual(
+        torch.rand(4, 5, dtype=torch.bfloat16), 0.2
+    ),
+    "ge_scalar_rank3_randn_bf16": lambda: GreaterEqual(
+        torch.randn(2, 3, 4, dtype=torch.bfloat16), -0.1
+    ),
+}
+
 
-@common.parametrize("test_module", test_data_tensor)
+@common.parametrize("test_module", test_data_tensor | test_data_tensor_bf16)
 def test_ge_tensor_tosa_FP(test_module):
     pipeline = TosaPipelineFP[input_t](
         test_module(),
         test_module().get_inputs(),
         GreaterEqual.aten_op_tensor,
         GreaterEqual.exir_op,
+        tosa_extensions=["bf16"],
     )
     pipeline.run()
 
 
-@common.parametrize("test_module", test_data_scalar)
+@common.parametrize("test_module", test_data_scalar | test_data_scalar_bf16)
 def test_ge_scalar_tosa_FP(test_module):
     pipeline = TosaPipelineFP[input_t](
         test_module(),
         test_module().get_inputs(),
         GreaterEqual.aten_op_scalar,
         GreaterEqual.exir_op,
+        tosa_extensions=["bf16"],
     )
     pipeline.run()
 
diff --git a/backends/arm/test/ops/test_gt.py b/backends/arm/test/ops/test_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.
@@ -70,32 +70,54 @@ def get_inputs(self):
     "gt_tensor_rank4_randn": lambda: op_gt_tensor_rank4_randn,
 }
 
+test_data_tensor_bf16 = {
+    "gt_tensor_rank2_rand_bf16": lambda: Greater(
+        torch.rand(4, 5, dtype=torch.bfloat16),
+        torch.rand(1, 5, dtype=torch.bfloat16),
+    ),
+    "gt_tensor_rank3_randn_bf16": lambda: Greater(
+        torch.randn(2, 3, 4, dtype=torch.bfloat16),
+        torch.randn(2, 3, 4, dtype=torch.bfloat16),
+    ),
+}
+
 test_data_scalar = {
     "gt_scalar_rank1_ones": lambda: op_gt_scalar_rank1_ones,
     "gt_scalar_rank2_rand": lambda: op_gt_scalar_rank2_rand,
     "gt_scalar_rank3_randn": lambda: op_gt_scalar_rank3_randn,
     "gt_scalar_rank4_randn": lambda: op_gt_scalar_rank4_randn,
 }
 
+test_data_scalar_bf16 = {
+    "gt_scalar_rank2_rand_bf16": lambda: Greater(
+        torch.rand(4, 5, dtype=torch.bfloat16), 0.2
+    ),
+    "gt_scalar_rank3_randn_bf16": lambda: Greater(
+        torch.randn(2, 3, 4, dtype=torch.bfloat16), -0.1
+    ),
+}
+
 
-@common.parametrize("test_module", test_data_tensor)
+@common.parametrize("test_module", test_data_tensor | test_data_tensor_bf16)
 def test_gt_tensor_tosa_FP(test_module):
     pipeline = TosaPipelineFP[input_t](
         test_module(),
         test_module().get_inputs(),
         Greater.aten_op_tensor,
         Greater.exir_op,
+        tosa_extensions=["bf16"],
     )
     pipeline.run()
 
 
-@common.parametrize("test_module", test_data_scalar)
+@common.parametrize("test_module", test_data_scalar | test_data_scalar_bf16)
 def test_gt_scalar_tosa_FP(test_module):
     pipeline = TosaPipelineFP[input_t](
         test_module(),
         test_module().get_inputs(),
         Greater.aten_op_scalar,
         Greater.exir_op,
+        tosa_extensions=["bf16"],
     )
     pipeline.run()
 
diff --git a/backends/arm/test/ops/test_index_tensor.py b/backends/arm/test/ops/test_index_tensor.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.
@@ -321,6 +321,12 @@ class IndexTensor(torch.nn.Module):
             ),
         ),
     }
+    test_data_bf16: dict[input_params] = {
+        "test_2d_1_idx_bf16": (
+            torch.rand(size=(3, 4), dtype=torch.bfloat16),
+            (torch.arange(2, dtype=torch.int32),),
+        ),
+    }
 
     # xfail - None (unsqueeze) unsupported
     test_data_none: dict[input_params] = {
@@ -372,7 +378,7 @@ def forward(self, input_: torch.Tensor, indices: Tuple[None | torch.Tensor]):
         return input_[indices]
 
 
-@common.parametrize("test_data", IndexTensor.test_data)
+@common.parametrize("test_data", IndexTensor.test_data | IndexTensor.test_data_bf16)
 def test_index_tensor_tosa_FP(test_data: input_params):
     test_input = test_data
     with torch.no_grad():
@@ -384,6 +390,7 @@ def test_index_tensor_tosa_FP(test_data: input_params):
                 IndexTensorTestCommon.exir_op,
                 atol=IndexTensorTestCommon.atol,
                 rtol=IndexTensorTestCommon.rtol,
+                tosa_extensions=["bf16"],
             ).run()
         )
 
diff --git a/backends/arm/test/ops/test_log.py b/backends/arm/test/ops/test_log.py
@@ -1,6 +1,6 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
-# 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.
@@ -34,15 +34,22 @@
     "ramp": lambda: (torch.arange(0.01, 20, 0.2)),
 }
 
+test_data_suite_bf16 = {
+    "rand_bf16": lambda: (torch.rand(6, 6, dtype=torch.bfloat16) + 1),
+    "ramp_bf16": lambda: (torch.arange(0.5, 8, 0.5, dtype=torch.bfloat16)),
+}
+
 
 class Log(torch.nn.Module):
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         return torch.log(x)
 
 
-@common.parametrize("test_data", test_data_suite)
+@common.parametrize("test_data", test_data_suite | test_data_suite_bf16)
 def test_log_tosa_FP(test_data: input_t1):
-    pipeline = TosaPipelineFP[input_t1](Log(), (test_data(),), aten_op, exir_op)
+    pipeline = TosaPipelineFP[input_t1](
+        Log(), (test_data(),), aten_op, exir_op, tosa_extensions=["bf16"]
+    )
     pipeline.run()
 
 
diff --git a/backends/arm/test/ops/test_matmul.py b/backends/arm/test/ops/test_matmul.py
diff --git a/backends/arm/test/ops/test_max_pool.py b/backends/arm/test/ops/test_max_pool.py
diff --git a/backends/arm/test/ops/test_maximum.py b/backends/arm/test/ops/test_maximum.py
diff --git a/backends/arm/test/ops/test_minimum.py b/backends/arm/test/ops/test_minimum.py
diff --git a/backends/arm/tosa/dialect/ops/matmul.py b/backends/arm/tosa/dialect/ops/matmul.py

Original file line number	Diff line number	Diff line change
`@@ -41,7 +41,7 @@ def define_node(`
`41`	`41`	`validate_valid_dtype(`
`42`	`42`	`self.target,`
`43`	`43`	`inputs,`
`44`		`- [ts.DType.INT32, ts.DType.FP32],`
	`44`	`+ [ts.DType.INT32, ts.DType.FP32, ts.DType.BF16],`
`45`	`45`	`self.tosa_spec,`
`46`	`46`	`)`
`47`	`47`	`validate_valid_dtype(self.target, output, ts.DType.BOOL, self.tosa_spec)`
Original file line number	Diff line number	Diff line change
`@@ -49,6 +49,8 @@ def define_node(`
`49`	`49`	`]`
`50`	`50`	`if self.tosa_spec.support_float():`
`51`	`51`	`supported_dtypes += [ts.DType.FP32]`
	`52`	`+ if self.tosa_spec.support_extension("bf16"):`
	`53`	`+ supported_dtypes += [ts.DType.BF16]`
`52`	`54`	`if self.tosa_spec.support_extension("int16"):`
`53`	`55`	`supported_dtypes += [ts.DType.INT48]`
`54`	`56`	`if self.tosa_spec.support_extension("int4"):`