[Mlir-commits] [mlir] f7e70bd - [mlir][tosa] Avoid introducing int <-> float casts (#195882)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu May 7 10:53:23 PDT 2026
Author: Ian Tayler Lessa
Date: 2026-05-07T18:53:18+01:00
New Revision: f7e70bd7a9fc9a6cd818698f4f8bd4dcbdcfa3f8
URL: https://github.com/llvm/llvm-project/commit/f7e70bd7a9fc9a6cd818698f4f8bd4dcbdcfa3f8
DIFF: https://github.com/llvm/llvm-project/commit/f7e70bd7a9fc9a6cd818698f4f8bd4dcbdcfa3f8.diff
LOG: [mlir][tosa] Avoid introducing int <-> float casts (#195882)
As part of the NonNarrowingCastsOptimization we were optimizing away
some cases where the inner input was an integer and the outer output was
a float.
Not all of the resulting dtype combinations for these cases are
supported by TOSA, so these scenarios are no longer optimized as part of
canonicalizations.
---------
Signed-off-by: Ian Tayler Lessa <ian.taylerlessa at arm.com>
Added:
Modified:
mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
mlir/test/Dialect/Tosa/canonicalize.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
index 642ee4b98e216..4af185a6e534b 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
@@ -967,40 +967,98 @@ struct NonNarrowingCastsOptimization : public OpRewritePattern<tosa::CastOp> {
const Value innerCastInput = innerCastOp.getInput();
- const auto innerInputType =
+ const ShapedType innerInputType =
llvm::cast<ShapedType>(innerCastInput.getType());
- const auto innerOutputType = llvm::cast<ShapedType>(innerCastOp.getType());
- const auto outerOutputType = llvm::cast<ShapedType>(castOp.getType());
+ const ShapedType innerOutputType =
+ llvm::cast<ShapedType>(innerCastOp.getType());
+ const ShapedType outerOutputType = llvm::cast<ShapedType>(castOp.getType());
- const SmallVector<ShapedType, 3> types = {innerInputType, innerOutputType,
- outerOutputType};
+ const Type innerInputElemType = innerInputType.getElementType();
+ const Type innerOutputElemType = innerOutputType.getElementType();
+ const Type outerOutputElemType = outerOutputType.getElementType();
- if (llvm::any_of(types, [](const ShapedType type) {
- const auto elemTy = type.getElementType();
+ const SmallVector<Type, 3> types = {innerInputElemType, innerOutputElemType,
+ outerOutputElemType};
+
+ if (llvm::any_of(types, [](const Type type) {
// Support a specific set of floating point types since we need to be
// careful in not introducing unsupported type combinations
- return !(elemTy.isInteger() ||
+ return !(type.isInteger() ||
llvm::isa<Float8E4M3FNType, Float8E5M2Type, BFloat16Type,
- Float16Type, Float32Type>(elemTy));
+ Float16Type, Float32Type>(type));
}))
return rewriter.notifyMatchFailure(
castOp, "only integer and f32, f16, bf16, f8E4M3FN, f8E5M2 types are "
"supported");
- if (llvm::isa<Float8E5M2Type>(innerInputType.getElementType()) &&
- llvm::isa<Float8E4M3FNType>(outerOutputType.getElementType())) {
+ if (llvm::isa<Float8E5M2Type>(innerInputElemType) &&
+ llvm::isa<Float8E4M3FNType>(outerOutputElemType)) {
return rewriter.notifyMatchFailure(
castOp, "avoid introducing f8E5M2 -> f8E4M3FN casts which are not "
"legal in TOSA");
}
- if (llvm::isa<Float8E4M3FNType>(innerInputType.getElementType()) &&
- llvm::isa<Float8E5M2Type>(outerOutputType.getElementType())) {
+ if (llvm::isa<Float8E4M3FNType>(innerInputElemType) &&
+ llvm::isa<Float8E5M2Type>(outerOutputElemType)) {
return rewriter.notifyMatchFailure(
castOp, "avoid introducing f8E4M3FN -> f8E5M2 casts which are not "
"legal in TOSA");
}
+ if (llvm::isa<Float8E5M2Type, Float8E4M3FNType>(innerInputElemType) &&
+ outerOutputElemType.isInteger()) {
+ return rewriter.notifyMatchFailure(
+ castOp, "avoid introducing fp8 -> integer casts which are not "
+ "legal in TOSA");
+ }
+
+ if (innerInputElemType.isInteger() &&
+ llvm::isa<Float8E5M2Type, Float8E4M3FNType>(outerOutputElemType)) {
+ return rewriter.notifyMatchFailure(
+ castOp, "avoid introducing integer -> fp8 casts which are not "
+ "legal in TOSA");
+ }
+
+ if (llvm::isa<Float16Type>(innerInputElemType) &&
+ llvm::isa<BFloat16Type>(outerOutputElemType)) {
+ return rewriter.notifyMatchFailure(
+ castOp, "avoid introducing fp16 -> bf16 casts which are not "
+ "legal in TOSA");
+ }
+
+ if (llvm::isa<BFloat16Type>(innerInputElemType) &&
+ llvm::isa<Float16Type>(outerOutputElemType)) {
+ return rewriter.notifyMatchFailure(
+ castOp, "avoid introducing bf16 -> fp16 casts which are not "
+ "legal in TOSA");
+ }
+
+ const auto isIntegerOneOfWidth = [](Type type, size_t bitwidth1,
+ size_t bitwidth2) {
+ return type.isInteger(bitwidth1) || type.isInteger(bitwidth2);
+ };
+
+ if (isIntegerOneOfWidth(innerInputElemType, 8, 16) &&
+ outerOutputElemType.isInteger(64)) {
+ return rewriter.notifyMatchFailure(
+ castOp, "avoid introducing i8/i16 -> i64 casts which are not "
+ "legal in TOSA");
+ }
+
+ if (isIntegerOneOfWidth(innerInputElemType, 1, 64) &&
+ !outerOutputElemType.isInteger()) {
+ return rewriter.notifyMatchFailure(
+ castOp, "avoid introducing bool/i64 to float casts which are not "
+ "supported in all versions of TOSA");
+ }
+
+ if (!innerInputElemType.isInteger() &&
+ isIntegerOneOfWidth(outerOutputElemType, 1, 64)) {
+ return rewriter.notifyMatchFailure(
+ castOp, "avoid introducing float to bool/i64 casts which are not "
+ "supported in all versions of TOSA");
+ }
+
// Check that the cast we're considering for removal is non-narrowing
if (isNarrowingCast(innerInputType, innerOutputType))
return rewriter.notifyMatchFailure(castOp,
diff --git a/mlir/test/Dialect/Tosa/canonicalize.mlir b/mlir/test/Dialect/Tosa/canonicalize.mlir
index 19583e111ebef..d4f3d23fd761e 100644
--- a/mlir/test/Dialect/Tosa/canonicalize.mlir
+++ b/mlir/test/Dialect/Tosa/canonicalize.mlir
@@ -1541,6 +1541,17 @@ func.func @test_canonicalize_non_narrowing_cast_f16_to_f32_to_f8(%arg0: tensor<1
// -----
+// CHECK-LABEL: @test_canonicalize_non_narrowing_cast_i8_to_i32_to_f16
+// CHECK: %[[OUT:.+]] = tosa.cast %arg0 : (tensor<13x21x3xi8>) -> tensor<13x21x3xf16>
+// CHECK: return %[[OUT]] : tensor<13x21x3xf16>
+func.func @test_canonicalize_non_narrowing_cast_i8_to_i32_to_f16(%arg0: tensor<13x21x3xi8>) -> tensor<13x21x3xf16> {
+ %0 = tosa.cast %arg0 : (tensor<13x21x3xi8>) -> tensor<13x21x3xi32>
+ %1 = tosa.cast %0 : (tensor<13x21x3xi32>) -> tensor<13x21x3xf16>
+ return %1 : tensor<13x21x3xf16>
+}
+
+// -----
+
// CHECK-LABEL: @test_canonicalize_non_narrowing_cast_f8E4M3FN_to_f16_to_f8E5M2
// CHECK: tosa.cast
// CHECK: tosa.cast
@@ -1596,6 +1607,50 @@ func.func @test_canonicalize_non_narrowing_cast_f6E3M2FN_to_f8E4M3FN_to_f16_unsu
// -----
+// CHECK-LABEL: @test_canonicalize_non_narrowing_cast_i1_to_f32_unsupported
+// CHECK: tosa.cast
+// CHECK: tosa.cast
+func.func @test_canonicalize_non_narrowing_cast_i1_to_f32_unsupported(%arg0: tensor<13x21x3xi1>) -> tensor<13x21x3xf32> {
+ %0 = tosa.cast %arg0 : (tensor<13x21x3xi1>) -> tensor<13x21x3xi8>
+ %1 = tosa.cast %0 : (tensor<13x21x3xi8>) -> tensor<13x21x3xf32>
+ return %1 : tensor<13x21x3xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @test_canonicalize_non_narrowing_cast_i8_to_i64_unsupported
+// CHECK: tosa.cast
+// CHECK: tosa.cast
+func.func @test_canonicalize_non_narrowing_cast_i8_to_i64_unsupported(%arg0: tensor<13x21x3xi8>) -> tensor<13x21x3xi64> {
+ %0 = tosa.cast %arg0 : (tensor<13x21x3xi8>) -> tensor<13x21x3xi32>
+ %1 = tosa.cast %0 : (tensor<13x21x3xi32>) -> tensor<13x21x3xi64>
+ return %1 : tensor<13x21x3xi64>
+}
+
+// -----
+
+// CHECK-LABEL: @test_canonicalize_non_narrowing_cast_f16_to_bf16_unsupported
+// CHECK: tosa.cast
+// CHECK: tosa.cast
+func.func @test_canonicalize_non_narrowing_cast_f16_to_bf16_unsupported(%arg0: tensor<13x21x3xf16>) -> tensor<13x21x3xbf16> {
+ %0 = tosa.cast %arg0 : (tensor<13x21x3xf16>) -> tensor<13x21x3xf32>
+ %1 = tosa.cast %0 : (tensor<13x21x3xf32>) -> tensor<13x21x3xbf16>
+ return %1 : tensor<13x21x3xbf16>
+}
+
+// -----
+
+// CHECK-LABEL: @test_canonicalize_non_narrowing_cast_i8_to_f8E4M3FN_unsupported
+// CHECK: tosa.cast
+// CHECK: tosa.cast
+func.func @test_canonicalize_non_narrowing_cast_i8_to_f8E4M3FN_unsupported(%arg0: tensor<13x21x3xi8>) -> tensor<13x21x3xf8E4M3FN> {
+ %0 = tosa.cast %arg0 : (tensor<13x21x3xi8>) -> tensor<13x21x3xf32>
+ %1 = tosa.cast %0 : (tensor<13x21x3xf32>) -> tensor<13x21x3xf8E4M3FN>
+ return %1 : tensor<13x21x3xf8E4M3FN>
+}
+
+// -----
+
// 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>) {
More information about the Mlir-commits
mailing list