[MLIR][STD] Add safe scalar constant propagation for FPTruncOp

[max-kudr]

Summary

Perform scalar constant propagation for FPTruncOp only if the resulting value can be represented without precision loss or rounding.

Example:
%cst = constant 1.000000e+00 : f32
%0 = fptrunc %cst : f32 to bf16
-->
%cst = constant 1.000000e+00 : bf16

[max-kudr]

If there is desire to always perform constant propagation, even if rounding or precision loss happens, let me know. Please also specify if warning message would be appropriate in these cases.

[mehdi_amini]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

1356

I think safe would be more accurate here?

mlir/test/Dialect/Standard/canonicalize.mlir

96

You can simplify the checks:

// CHECK: constant 1.444000e+25 : f32
// CHECK: fptrunc

That's enough to check that fptrunc isn't folded and the constant type unchanged.

[max-kudr]

Updated according to comments

[mehdi_amini]

LGTM

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

     If the value cannot be exactly represented, it is rounded using the default
     rounding mode. When operating on vectors, casts elementwise.
   }];
+
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

   return areVectorCastSimpleCompatible(a, b, areCastCompatible);
 }
 
+/// Perform safe const propagation for fptrunc, i.e. only propagate
+/// if FP value can be represented without precision loss or rounding.
+OpFoldResult FPTruncOp::fold(ArrayRef<Attribute> operands) {
+  assert(operands.size() == 1 && "unary operation takes one operand");
+
+  auto constOperand = operands.front();
+  if (!constOperand || !constOperand.isa<FloatAttr>())
+    return {};
+
+  // Convert to target type via 'double'.
+  double sourceValue =
+      constOperand.dyn_cast<FloatAttr>().getValue().convertToDouble();
+  auto targetAttr = FloatAttr::get(getType(), sourceValue);
+
+  // Propagate if constant's value does not change after truncation.
+  if (sourceValue == targetAttr.getValue().convertToDouble())
+    return targetAttr;
+
+  return {};
+}
+
 //===----------------------------------------------------------------------===//
 // IndexCastOp
 //===----------------------------------------------------------------------===//

mlir/test/Dialect/Linalg/generalize-named-polymorphic-ops.mlir

 }
 
 // CHECK-LABEL: @generalize_soft_plus_2d_f32
-//      CHECK: %[[C1:.+]] = constant 1.000000e+00 : f64
+//      CHECK: %[[C1:.+]] = constant 1.000000e+00 : f32
 //      CHECK: ^{{.*}}(%[[IN:.+]]: f32, %[[OUT:.+]]: f32
-// CHECK-NEXT:   %[[C1_CAST:.+]] = fptrunc %[[C1]] : f64 to f32
 // CHECK-NEXT:   %[[EXP:.+]] = math.exp %[[IN]] : f32
-// CHECK-NEXT:   %[[SUM:.+]] = addf %[[C1_CAST]], %[[EXP]] : f32
+// CHECK-NEXT:   %[[SUM:.+]] = addf %[[C1]], %[[EXP]] : f32
 // CHECK-NEXT:   %[[LOG:.+]] = math.log %[[SUM]] : f32
 // CHECK-NEXT:   linalg.yield %[[LOG]] : f32
 // CHECK-NEXT: -> tensor<16x32xf32>

mlir/test/Dialect/Standard/canonicalize.mlir

   return %tr : i16
 }
 
+// CHECK-LABEL: @truncFPConstant
+//       CHECK:   %[[cres:.+]] = constant 1.000000e+00 : bf16
+//       CHECK:   return %[[cres]]
+func @truncFPConstant() -> bf16 {
+  %cst = constant 1.000000e+00 : f32
+  %0 = fptrunc %cst : f32 to bf16
+  return %0 : bf16
+}
+
+// Test that cases with rounding are NOT propagated
+// CHECK-LABEL: @truncFPConstantRounding
+//       CHECK:   constant 1.444000e+25 : f32
+//       CHECK:   fptrunc
+func @truncFPConstantRounding() -> bf16 {
+  %cst = constant 1.444000e+25 : f32
+  %0 = fptrunc %cst : f32 to bf16
+  return %0 : bf16
+}
+
 // CHECK-LABEL: @tripleAddAdd
 //       CHECK:   %[[cres:.+]] = constant 59 : index 
 //       CHECK:   %[[add:.+]] = addi %arg0, %[[cres]] : index