Restrict to bf16-f32-bf16

Author: mgehre-amd

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

@@ -1008,11 +1008,13 @@ OpFoldResult CastOp::fold(FoldAdaptor adaptor) {
     }
   }
 
-  // cast-to-bf16(cast-to-f32(x)) -> cast-to-bf16(x)
+  // Fold cast from bf16 -> f32 -> bf16 into no-op.
   if (auto cast = getInput().getDefiningOp<CastOp>()) {
+    auto sourceElTy = cast.getInput().getType().getElementType();
     auto intermediateElTy = cast.getType().getElementType();
     auto finalElTy = getType().getElementType();
-    if (isa<Float32Type>(intermediateElTy) && isa<BFloat16Type>(finalElTy)) {
+    if (isa<BFloat16Type>(sourceElTy) && isa<Float32Type>(intermediateElTy) &&
+        isa<BFloat16Type>(finalElTy)) {
       getInputMutable().assign(cast.getInput());
       return getResult();
     }

mlir/test/Dialect/Tosa/canonicalize.mlir

@@ -56,9 +56,9 @@ func.func @cast_fold_double(%arg0: tensor<?x1xf32>) -> 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>
+func.func @cast_fold_double2(%arg0: tensor<?x1xbf16>) -> tensor<?x1xbf16> {
+  // CHECK: return %arg0
+  %0 = tosa.cast %arg0 : (tensor<?x1xbf16>) -> tensor<?x1xf32>
   %1 = tosa.cast %0 : (tensor<?x1xf32>) -> tensor<?x1xbf16>
   return %1 : tensor<?x1xbf16>
 }