[Mlir-commits] [mlir] [mlir][tosa] Extend cancelling block-scaled cast optimisation to bf16 (PR #212517)

[Ian Tayler Lessa]

https://github.com/IanTaylerLessa-arm created https://github.com/llvm/llvm-project/pull/212517

The CancellingBlockScaledCastsOptimization optimises away sequences where a block-scaled tensor is cast to fp32 and then cast back to its original block-scaled type.

This patch extends it to also cover cases where the intermediate type is bf16, since the 7 bits of mantissa and 8 bits in the exponent are sufficient to encode any `element * scale` product with valid block-scaled type combinations.

>From 4d5ad6011bfe63b0d0a87300a307fa31b24dc780 Mon Sep 17 00:00:00 2001
From: Ian Tayler Lessa <ian.taylerlessa at arm.com>
Date: Tue, 28 Jul 2026 15:57:27 +0100
Subject: [PATCH] [mlir][tosa] Extend cancelling block-scaled cast optimisation
 to bf16

The CancellingBlockScaledCastsOptimization optimises away sequences
where a block-scaled tensor is cast to fp32 and then cast back to its
original block-scaled type.

This patch extends it to also cover cases where the intermediate type is
bf16, since the 7 bits of mantissa and 8 bits in the exponent are
sufficient to encode any `element * scale` product with valid
block-scaled type combinations.

Signed-off-by: Ian Tayler Lessa <ian.taylerlessa at arm.com>
Change-Id: I62b2dd7e07e99dfb9ffa84dc586563658e069c71
---
 .../Dialect/Tosa/IR/TosaCanonicalizations.cpp |  3 +-
 mlir/test/Dialect/Tosa/canonicalize.mlir      | 34 +++++++++++++++++++
 2 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
index 8937875df7fc5..bb3c0f16b97aa 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
@@ -1263,7 +1263,8 @@ struct CancellingBlockScaledCastsOptimization
           castOp, "inner input type must match outer output type");
 
     const Type innerOutputElemType = innerOutputTy.getElementType();
-    const bool isLosslessCast = isa<Float32Type>(innerOutputElemType);
+    const bool isLosslessCast =
+        isa<Float32Type, BFloat16Type>(innerOutputElemType);
     if (!isLosslessCast)
       return rewriter.notifyMatchFailure(
           castOp, "avoid cancelling casts that should be lossy");
diff --git a/mlir/test/Dialect/Tosa/canonicalize.mlir b/mlir/test/Dialect/Tosa/canonicalize.mlir
index 7c9dd260ca497..24b47b79ac760 100644
--- a/mlir/test/Dialect/Tosa/canonicalize.mlir
+++ b/mlir/test/Dialect/Tosa/canonicalize.mlir
@@ -1773,6 +1773,40 @@ func.func @test_do_not_canonicalize_cast_from_cast_to_block_scaled_type_f8E5M2_f
 
 // -----
 
+// CHECK-LABEL: @test_canonicalize_cast_from_cast_to_block_scaled_type_f8E5M2_through_bf16
+// CHECK: return %arg0 : tensor<160x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f8E5M2>>
+func.func @test_canonicalize_cast_from_cast_to_block_scaled_type_f8E5M2_through_bf16(%arg0: tensor<160x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f8E5M2>>) -> tensor<160x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f8E5M2>> {
+  %0 = tosa.cast %arg0 : (tensor<160x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f8E5M2>>) -> tensor<160xbf16>
+  %1 = tosa.cast %0 : (tensor<160xbf16>) -> tensor<160x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f8E5M2>>
+  return %1 : tensor<160x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f8E5M2>>
+}
+
+// -----
+
+// CHECK-LABEL: @test_do_not_canonicalize_cast_from_cast_to_block_scaled_type_different_types_f8E5M2_f6E2M3_through_bf16
+// CHECK: %[[values:.+]] = tosa.cast %arg0
+// CHECK: %[[block_scaled:.+]] = tosa.cast %[[values]]
+// CHECK: return %[[block_scaled]] : tensor<160x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f6E2M3FN>>
+func.func @test_do_not_canonicalize_cast_from_cast_to_block_scaled_type_different_types_f8E5M2_f6E2M3_through_bf16(%arg0: tensor<160x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f8E5M2>>) -> tensor<160x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f6E2M3FN>> {
+  %0 = tosa.cast %arg0 : (tensor<160x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f8E5M2>>) -> tensor<160xbf16>
+  %1 = tosa.cast %0 : (tensor<160xbf16>) -> tensor<160x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f6E2M3FN>>
+  return %1 : tensor<160x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f6E2M3FN>>
+}
+
+// -----
+
+// CHECK-LABEL: @test_do_not_canonicalize_cast_from_cast_to_block_scaled_type_f4E2M1FN_through_f8E5M2
+// CHECK: %[[values:.+]] = tosa.cast %arg0
+// CHECK: %[[block_scaled:.+]] = tosa.cast %[[values]]
+// CHECK: return %[[block_scaled]] : tensor<160x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>
+func.func @test_do_not_canonicalize_cast_from_cast_to_block_scaled_type_f4E2M1FN_through_f8E5M2(%arg0: tensor<160x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>) -> tensor<160x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>> {
+  %0 = tosa.cast %arg0 : (tensor<160x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>) -> tensor<160xf8E5M2>
+  %1 = tosa.cast %0 : (tensor<160xf8E5M2>) -> tensor<160x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>
+  return %1 : tensor<160x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>
+}
+
+// -----
+
 // CHECK-LABEL: @test_canonicalize_cast_from_cast_to_block_scaled_f4E2M1
 // CHECK: return %arg0, %arg1 : tensor<15x3x2x256xf4E2M1FN>, tensor<15x3x2x8xf8E8M0FNU>
 func.func @test_canonicalize_cast_from_cast_to_block_scaled_f4E2M1(%arg0: tensor<15x3x2x256xf4E2M1FN>, %arg1: tensor<15x3x2x8xf8E8M0FNU>) -> (tensor<15x3x2x256xf4E2M1FN>, tensor<15x3x2x8xf8E8M0FNU>) {