Some fixes

Author: zahimoud

mlir/lib/Dialect/Arith/IR/ArithOps.cpp

@@ -1517,23 +1517,25 @@ LogicalResult arith::TruncIOp::verify() {
 /// Perform safe const propagation for truncf, i.e., only propagate if FP value
 /// can be represented without precision loss.
 OpFoldResult arith::TruncFOp::fold(FoldAdaptor adaptor) {
+  auto resElemType = cast<FloatType>(getElementTypeOrSelf(getType()));
   if (auto extOp = getOperand().getDefiningOp<arith::ExtFOp>()) {
     Value src = extOp.getIn();
-    Type srcType = getElementTypeOrSelf(src.getType());
-    Type dstType = getElementTypeOrSelf(getType());
-    // truncf(extf(a)) -> truncf(a)
-    if (llvm::cast<FloatType>(srcType).getWidth() >
-        llvm::cast<FloatType>(dstType).getWidth()) {
-      setOperand(src);
-      return getResult();
-    }
+    auto srcType = cast<FloatType>(getElementTypeOrSelf(src.getType()));
+    auto intermediateType = cast<FloatType>(getElementTypeOrSelf(extOp.getType()));
+    // Check if the srcType is representable in the intermediateType
+    if(llvm::APFloatBase::isRepresentableBy(srcType.getFloatSemantics(), intermediateType.getFloatSemantics())) {
+      // truncf(extf(a)) -> truncf(a)
+      if (srcType.getWidth() > resElemType.getWidth()) {
+        setOperand(src);
+        return getResult();
+      }
 
-    // truncf(extf(a)) -> a
-    if (srcType == dstType)
-      return src;
+      // truncf(extf(a)) -> a
+      if (srcType == resElemType)
+        return src;
+    }
   }
 
-  auto resElemType = cast<FloatType>(getElementTypeOrSelf(getType()));
   const llvm::fltSemantics &targetSemantics = resElemType.getFloatSemantics();
   return constFoldCastOp<FloatAttr, FloatAttr>(
       adaptor.getOperands(), getType(),

mlir/test/Dialect/Arith/canonicalize.mlir

@@ -723,11 +723,31 @@ func.func @truncExtf(%arg0: f32) -> f32 {
   return %trunc : f32
 }
 
+// CHECK-LABEL: @truncExtf1
+//       CHECK-NOT:  truncf
+//       CHECK:   return  %arg0
+func.func @truncExtf1(%arg0: bf16) -> bf16 {
+  %extf = arith.extf %arg0 : bf16 to f32
+  %trunc = arith.truncf %extf : f32 to bf16
+  return %trunc : bf16
+}
+
 // CHECK-LABEL: @truncExtf2
+//       CHECK:  %[[ARG0:.+]]: bf16
+//       CHECK:  %[[EXTF:.*]] = arith.extf %[[ARG0:.+]] : bf16 to f32
+//       CHECK:  %[[TRUNCF:.*]] = arith.truncf %[[EXTF:.*]] : f32 to f16
+//       CHECK:   return  %[[TRUNCF:.*]]
+func.func @truncExtf2(%arg0: bf16) -> f16 {
+  %extf = arith.extf %arg0 : bf16 to f32
+  %trunc = arith.truncf %extf : f32 to f16
+  return %trunc : f16
+}
+
+// CHECK-LABEL: @truncExtf3
 //       CHECK:  %[[ARG0:.+]]: f32
 //       CHECK:  %[[CST:.*]] = arith.truncf %[[ARG0:.+]] : f32 to f16
 //       CHECK:   return  %[[CST:.*]]
-func.func @truncExtf2(%arg0: f32) -> f16 {
+func.func @truncExtf3(%arg0: f32) -> f16 {
   %extf = arith.extf %arg0 : f32 to f64
   %truncf = arith.truncf %extf : f64 to f16
   return %truncf : f16