TOSA: fold cast-to-bf16(cast-to-f32(x)) -> cast-to-bf16(x)

Author: mgehre-amd

mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp

@@ -1008,6 +1008,16 @@ OpFoldResult CastOp::fold(FoldAdaptor adaptor) {
     }
   }
 
+  // cast-to-bf16(cast-to-f32(x)) -> cast-to-bf16(x)
+  if (auto cast = getInput().getDefiningOp<CastOp>()) {
+    auto intermediateElTy = cast.getType().getElementType();
+    auto finalElTy = getType().getElementType();
+    if (isa<Float32Type>(intermediateElTy) && isa<BFloat16Type>(finalElTy)) {
+      getInputMutable().assign(cast.getInput());
+      return getResult();
+    }
+  }
+
   auto operand = llvm::dyn_cast_if_present<ElementsAttr>(adaptor.getInput());
   if (!operand)
     return {};

mlir/test/Dialect/Tosa/canonicalize.mlir

@@ -55,6 +55,14 @@ func.func @cast_fold_double(%arg0: tensor<?x1xf32>) -> tensor<?x1xi8> {
   return %1 : tensor<?x1xi8>
 }
 
+// CHECK-LABEL: @cast_fold_double
+func.func @cast_fold_double2(%arg0: tensor<?x1xf16>) -> tensor<?x1xbf16> {
+  // CHECK: tosa.cast{{.*}} (tensor<?x1xf16>) -> tensor<?x1xbf16>
+  %0 = tosa.cast %arg0 : (tensor<?x1xf16>) -> tensor<?x1xf32>
+  %1 = tosa.cast %0 : (tensor<?x1xf32>) -> tensor<?x1xbf16>
+  return %1 : tensor<?x1xbf16>
+}
+
 // CHECK-LABEL: @cast_no_fold_double1
 func.func @cast_no_fold_double1(%arg0: tensor<?x1xf32>) -> tensor<?x1xi8> {
   // CHECK: tosa.cast{{.*}} (tensor<?x1xf32>) -> tensor<?x1xui16>