[mlir][arith] Fold scaling_extf and scaling_truncf with constant operands

mlir/include/mlir/Dialect/Arith/IR/ArithOps.td

@@ -1500,6 +1500,7 @@ def Arith_ScalingExtFOp
     %h = arith.scaling_extf %i, %f : vector<32xf4E2M1FN>, vector<32xf8E8M0FNU> to vector<32xbf16>
     ```
   }];
+  let hasFolder = 1;
   let hasVerifier = 1;
   let assemblyFormat =
       [{ $in `,` $scale (`fastmath` `` $fastmath^)? attr-dict `:`
@@ -1689,6 +1690,7 @@ def Arith_ScalingTruncFOp
     %h = arith.scaling_truncf %i, %f : vector<32xbf16>, vector<32xf8E8M0FNU> to vector<32xf4E2M1FN>
     ```
   }];
+  let hasFolder = 1;
   let hasVerifier = 1;
   let assemblyFormat =
       [{ $in `,` $scale ($roundingmode^)? (`fastmath` `` $fastmath^)? attr-dict `:`

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

@@ -1822,6 +1822,108 @@ LogicalResult arith::ExtFOp::verify() { return verifyExtOp<FloatType>(*this); }
 // ScalingExtFOp
 //===----------------------------------------------------------------------===//
 
+/// Fold `calculate` element-wise over the operands of a scaling cast op. The
+/// `constFoldBinaryOp` helpers cannot be used: they bail out unless both
+/// operands have the same type, and `in` and `scale` never do.
+static Attribute foldScalingCastOp(
+    Attribute inAttr, Attribute scaleAttr, Type resultType,
+    function_ref<std::optional<APFloat>(const APFloat &, const APFloat &)>
+        calculate) {
+  // Poison propagates, as it does in the generic constant folders.
+  if (isa_and_nonnull<ub::PoisonAttr>(inAttr))
+    return inAttr;
+  if (isa_and_nonnull<ub::PoisonAttr>(scaleAttr))
+    return scaleAttr;
+
+  if (!inAttr || !scaleAttr || !resultType)
+    return {};
+
+  if (auto inFloat = dyn_cast<FloatAttr>(inAttr)) {
+    auto scaleFloat = dyn_cast<FloatAttr>(scaleAttr);
+    if (!scaleFloat)
+      return {};
+    std::optional<APFloat> result =
+        calculate(inFloat.getValue(), scaleFloat.getValue());
+    if (!result)
+      return {};
+    return FloatAttr::get(resultType, *result);
+  }
+
+  auto inElements = dyn_cast<DenseFPElementsAttr>(inAttr);
+  auto scaleElements = dyn_cast<DenseFPElementsAttr>(scaleAttr);
+  auto shapedResultType = dyn_cast<ShapedType>(resultType);
+  if (!inElements || !scaleElements || !shapedResultType ||
+      !shapedResultType.hasStaticShape() ||
+      inElements.getNumElements() != scaleElements.getNumElements())
+    return {};
+
+  // Both operands are splats, so avoid expanding the elements out.
+  if (inElements.isSplat() && scaleElements.isSplat()) {
+    std::optional<APFloat> result =
+        calculate(inElements.getSplatValue<APFloat>(),
+                  scaleElements.getSplatValue<APFloat>());
+    if (!result)
+      return {};
+    return DenseElementsAttr::get(shapedResultType, *result);
+  }
+
+  SmallVector<APFloat> results;
+  results.reserve(inElements.getNumElements());
+  auto scaleIt = scaleElements.begin();
+  for (const APFloat &in : inElements) {
+    std::optional<APFloat> result = calculate(in, *scaleIt++);
+    if (!result)
+      return {};
+    results.push_back(*result);
+  }
+  return DenseElementsAttr::get(shapedResultType, results);
+}
+
+/// Return the NaN both scaling ops document propagating, or `std::nullopt` if
+/// `semantics` cannot hold one (`f4E2M1FN`, for one, is finite-only).
+///
+/// Built directly rather than by widening the scale, so that the fold still
+/// succeeds when widening `in` is lossy -- which a NaN scale makes irrelevant,
+/// since it forces a NaN result whatever `in` is.
+static std::optional<APFloat>
+getScalingCastNaN(const llvm::fltSemantics &semantics) {
+  if (semantics.nonFiniteBehavior == llvm::fltNonfiniteBehavior::FiniteOnly)
+    return std::nullopt;
+  return APFloat::getNaN(semantics);
+}
+
+/// Only scales that already are f8E8M0FNU fold: `-arith-expand` truncates
+/// wider scales to f8E8M0FNU first, while `ArithToAMDGPU` reads their exponent
+/// field instead, and the two disagree.
+static bool isFoldableScalingScale(Value scale) {
+  return isa<Float8E8M0FNUType>(getElementTypeOrSelf(scale.getType()));
+}
+
+OpFoldResult arith::ScalingExtFOp::fold(FoldAdaptor adaptor) {
+  // scaling_extf(in, scale) -> mulf(extf(in), extf(scale)), matching the
+  // expansion in ExpandOps.cpp. As in arith.extf, the widening steps only fold
+  // when they are lossless.
+  if (!isFoldableScalingScale(getScale()))
+    return {};
+
+  auto resElemType = cast<FloatType>(getElementTypeOrSelf(getType()));
+  const llvm::fltSemantics &resSemantics = resElemType.getFloatSemantics();
+  return foldScalingCastOp(
+      adaptor.getIn(), adaptor.getScale(), getType(),
+      [&resSemantics](const APFloat &in,
+                      const APFloat &scale) -> std::optional<APFloat> {
+        if (scale.isNaN())
+          return getScalingCastNaN(resSemantics);
+        FailureOr<APFloat> inExt = convertFloatValue(in, resSemantics);
+        FailureOr<APFloat> scaleExt = convertFloatValue(scale, resSemantics);
+        if (failed(inExt) || failed(scaleExt))
+          return std::nullopt;
+        APFloat result(*inExt);
+        result.multiply(*scaleExt, kDefaultRoundingMode);
+        return result;
+      });
+}
+
 bool arith::ScalingExtFOp::areCastCompatible(TypeRange inputs,
                                              TypeRange outputs) {
   return checkWidthChangeCast<std::greater, FloatType>(inputs.front(), outputs);
@@ -1994,6 +2096,37 @@ LogicalResult arith::ConvertFOp::verify() {
 // ScalingTruncFOp
 //===----------------------------------------------------------------------===//
 
+OpFoldResult arith::ScalingTruncFOp::fold(FoldAdaptor adaptor) {
+  // scaling_truncf(in, scale) -> truncf(in / extf(scale)), matching the
+  // expansion in ExpandOps.cpp. Unlike scaling_extf, the scale is widened to
+  // the type of `in` rather than to the result type.
+  if (!isFoldableScalingScale(getScale()))
+    return {};
+
+  auto inElemType = cast<FloatType>(getElementTypeOrSelf(getIn().getType()));
+  auto resElemType = cast<FloatType>(getElementTypeOrSelf(getType()));
+  const llvm::fltSemantics &inSemantics = inElemType.getFloatSemantics();
+  const llvm::fltSemantics &resSemantics = resElemType.getFloatSemantics();
+  llvm::RoundingMode roundingMode =
+      convertArithRoundingModeToLLVMIR(getRoundingmode());
+  return foldScalingCastOp(
+      adaptor.getIn(), adaptor.getScale(), getType(),
+      [&](const APFloat &in, const APFloat &scale) -> std::optional<APFloat> {
+        if (scale.isNaN())
+          return getScalingCastNaN(resSemantics);
+        FailureOr<APFloat> scaleExt = convertFloatValue(scale, inSemantics);
+        if (failed(scaleExt))
+          return std::nullopt;
+        APFloat quotient(in);
+        quotient.divide(*scaleExt, kDefaultRoundingMode);
+        FailureOr<APFloat> result =
+            convertFloatValue(quotient, resSemantics, roundingMode);
+        if (failed(result))
+          return std::nullopt;
+        return *result;
+      });
+}
+
 bool arith::ScalingTruncFOp::areCastCompatible(TypeRange inputs,
                                                TypeRange outputs) {
   return checkWidthChangeCast<std::less, FloatType>(inputs.front(), outputs);

mlir/test/Dialect/Arith/canonicalize.mlir

@@ -1120,6 +1120,79 @@ func.func @extFPVectorConstant() -> vector<2xf128> {
   return %0 : vector<2xf128>
 }
 
+// A f8E8M0FNU scale stands for 2^scale, so this is 1.5 * 2^2.
+// CHECK-LABEL: @scalingExtFConstant
+//       CHECK:   %[[cres:.+]] = arith.constant 6.000000e+00 : f32
+//       CHECK:   return %[[cres]]
+func.func @scalingExtFConstant() -> f32 {
+  %in = arith.constant 1.500000e+00 : f4E2M1FN
+  %scale = arith.constant 4.000000e+00 : f8E8M0FNU
+  %0 = arith.scaling_extf %in, %scale : f4E2M1FN, f8E8M0FNU to f32
+  return %0 : f32
+}
+
+// CHECK-LABEL: @scalingExtFVectorConstant
+//       CHECK:   %[[cres:.+]] = arith.constant dense<[2.000000e+00, 8.000000e+00]> : vector<2xf32>
+//       CHECK:   return %[[cres]]
+func.func @scalingExtFVectorConstant() -> vector<2xf32> {
+  %in = arith.constant dense<[1.000000e+00, 2.000000e+00]> : vector<2xf4E2M1FN>
+  %scale = arith.constant dense<[2.000000e+00, 4.000000e+00]> : vector<2xf8E8M0FNU>
+  %0 = arith.scaling_extf %in, %scale : vector<2xf4E2M1FN>, vector<2xf8E8M0FNU> to vector<2xf32>
+  return %0 : vector<2xf32>
+}
+
+// CHECK-LABEL: @scalingExtFSplatConstant
+//       CHECK:   %[[cres:.+]] = arith.constant dense<3.000000e+00> : vector<4xf32>
+//       CHECK:   return %[[cres]]
+func.func @scalingExtFSplatConstant() -> vector<4xf32> {
+  %in = arith.constant dense<1.500000e+00> : vector<4xf4E2M1FN>
+  %scale = arith.constant dense<2.000000e+00> : vector<4xf8E8M0FNU>
+  %0 = arith.scaling_extf %in, %scale : vector<4xf4E2M1FN>, vector<4xf8E8M0FNU> to vector<4xf32>
+  return %0 : vector<4xf32>
+}
+
+// The op propagates NaN from either operand. 0xFF is the only f8E8M0FNU NaN.
+// CHECK-LABEL: @scalingExtFNaNScaleConstant
+//       CHECK:   %[[cres:.+]] = arith.constant 0x7FC00000 : f32
+//       CHECK:   return %[[cres]]
+func.func @scalingExtFNaNScaleConstant() -> f32 {
+  %in = arith.constant 1.500000e+00 : f4E2M1FN
+  %scale = arith.constant 0xFF : f8E8M0FNU
+  %0 = arith.scaling_extf %in, %scale : f4E2M1FN, f8E8M0FNU to f32
+  return %0 : f32
+}
+
+// Test that scales which are not already f8E8M0FNU are NOT folded: -arith-expand
+// truncates them to f8E8M0FNU first, while ArithToAMDGPU reads their exponent.
+// CHECK-LABEL: @scalingExtFNonE8M0ScaleConstant
+//       CHECK:   arith.scaling_extf
+func.func @scalingExtFNonE8M0ScaleConstant() -> f32 {
+  %in = arith.constant 1.500000e+00 : f4E2M1FN
+  %scale = arith.constant 1.000000e+00 : f16
+  %0 = arith.scaling_extf %in, %scale : f4E2M1FN, f16 to f32
+  return %0 : f32
+}
+
+// Test that cases where widening the scale is lossy are NOT folded: 0xFE is
+// 2^127, which overflows f16.
+// CHECK-LABEL: @scalingExtFOverflowingScaleConstant
+//       CHECK:   arith.scaling_extf
+func.func @scalingExtFOverflowingScaleConstant() -> f16 {
+  %in = arith.constant 1.500000e+00 : f4E2M1FN
+  %scale = arith.constant 0xFE : f8E8M0FNU
+  %0 = arith.scaling_extf %in, %scale : f4E2M1FN, f8E8M0FNU to f16
+  return %0 : f16
+}
+
+// CHECK-LABEL: @scalingExtFPoisonScale
+//       CHECK:   %[[cres:.+]] = ub.poison : f32
+//       CHECK:   return %[[cres]]
+func.func @scalingExtFPoisonScale(%in: f4E2M1FN) -> f32 {
+  %scale = ub.poison : f8E8M0FNU
+  %0 = arith.scaling_extf %in, %scale : f4E2M1FN, f8E8M0FNU to f32
+  return %0 : f32
+}
+
 // CHECK-LABEL: @truncExtf
 //       CHECK-NOT:  truncf
 //       CHECK:   return  %arg0
@@ -1433,6 +1506,78 @@ func.func @truncFPConstantRounding() -> bf16 {
   return %0 : bf16
 }
 
+// Unlike arith.scaling_extf, this divides by 2^scale: 6.0 / 2^1.
+// CHECK-LABEL: @scalingTruncFConstant
+//       CHECK:   %[[cres:.+]] = arith.constant 3.000000e+00 : f4E2M1FN
+//       CHECK:   return %[[cres]]
+func.func @scalingTruncFConstant() -> f4E2M1FN {
+  %in = arith.constant 6.000000e+00 : f32
+  %scale = arith.constant 2.000000e+00 : f8E8M0FNU
+  %0 = arith.scaling_truncf %in, %scale : f32, f8E8M0FNU to f4E2M1FN
+  return %0 : f4E2M1FN
+}
+
+// CHECK-LABEL: @scalingTruncFVectorConstant
+//       CHECK:   %[[cres:.+]] = arith.constant dense<[3.000000e+00, 4.000000e+00]> : vector<2xf4E2M1FN>
+//       CHECK:   return %[[cres]]
+func.func @scalingTruncFVectorConstant() -> vector<2xf4E2M1FN> {
+  %in = arith.constant dense<[6.000000e+00, 8.000000e+00]> : vector<2xf32>
+  %scale = arith.constant dense<2.000000e+00> : vector<2xf8E8M0FNU>
+  %0 = arith.scaling_truncf %in, %scale : vector<2xf32>, vector<2xf8E8M0FNU> to vector<2xf4E2M1FN>
+  return %0 : vector<2xf4E2M1FN>
+}
+
+// CHECK-LABEL: @scalingTruncFDownwardConstant
+//       CHECK:   %[[cres:.+]] = arith.constant 3.000000e+00 : f4E2M1FN
+//       CHECK:   return %[[cres]]
+func.func @scalingTruncFDownwardConstant() -> f4E2M1FN {
+  %in = arith.constant 6.000000e+00 : f32
+  %scale = arith.constant 2.000000e+00 : f8E8M0FNU
+  %0 = arith.scaling_truncf %in, %scale downward : f32, f8E8M0FNU to f4E2M1FN
+  return %0 : f4E2M1FN
+}
+
+// Test that cases with rounding are NOT propagated: 5.0 is not representable
+// in f4E2M1FN, whose values step 0, 0.5, 1, 1.5, 2, 3, 4, 6.
+// CHECK-LABEL: @scalingTruncFConstantRounding
+//       CHECK:   arith.scaling_truncf
+func.func @scalingTruncFConstantRounding() -> f4E2M1FN {
+  %in = arith.constant 5.000000e+00 : f32
+  %scale = arith.constant 1.000000e+00 : f8E8M0FNU
+  %0 = arith.scaling_truncf %in, %scale : f32, f8E8M0FNU to f4E2M1FN
+  return %0 : f4E2M1FN
+}
+
+// CHECK-LABEL: @scalingTruncFNaNScaleConstant
+//       CHECK:   %[[cres:.+]] = arith.constant 0x7E00 : f16
+//       CHECK:   return %[[cres]]
+func.func @scalingTruncFNaNScaleConstant() -> f16 {
+  %in = arith.constant 6.000000e+00 : f32
+  %scale = arith.constant 0xFF : f8E8M0FNU
+  %0 = arith.scaling_truncf %in, %scale : f32, f8E8M0FNU to f16
+  return %0 : f16
+}
+
+// Test that a NaN scale is NOT folded when the result type cannot hold a NaN:
+// f4E2M1FN is finite-only.
+// CHECK-LABEL: @scalingTruncFNaNScaleFiniteOnlyResult
+//       CHECK:   arith.scaling_truncf
+func.func @scalingTruncFNaNScaleFiniteOnlyResult() -> f4E2M1FN {
+  %in = arith.constant 6.000000e+00 : f32
+  %scale = arith.constant 0xFF : f8E8M0FNU
+  %0 = arith.scaling_truncf %in, %scale : f32, f8E8M0FNU to f4E2M1FN
+  return %0 : f4E2M1FN
+}
+
+// CHECK-LABEL: @scalingTruncFPoisonInput
+//       CHECK:   %[[cres:.+]] = ub.poison : f4E2M1FN
+//       CHECK:   return %[[cres]]
+func.func @scalingTruncFPoisonInput(%scale: f8E8M0FNU) -> f4E2M1FN {
+  %in = ub.poison : f32
+  %0 = arith.scaling_truncf %in, %scale : f32, f8E8M0FNU to f4E2M1FN
+  return %0 : f4E2M1FN
+}
+
 // CHECK-LABEL: @tripleAddAdd
 //       CHECK:   %[[cres:.+]] = arith.constant 59 : index
 //       CHECK:   %[[add:.+]] = arith.addi %arg0, %[[cres]] : index