diff --git a/mlir/include/mlir/Dialect/Arith/IR/ArithOps.td b/mlir/include/mlir/Dialect/Arith/IR/ArithOps.td index 54481d3232483..39d6838ae1ab4 100644 --- a/mlir/include/mlir/Dialect/Arith/IR/ArithOps.td +++ b/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 `:` diff --git a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp index fc50284fc8a94..b0e7dd7721da8 100644 --- a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp +++ b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp @@ -1822,6 +1822,108 @@ LogicalResult arith::ExtFOp::verify() { return verifyExtOp(*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(const APFloat &, const APFloat &)> + calculate) { + // Poison propagates, as it does in the generic constant folders. + if (isa_and_nonnull(inAttr)) + return inAttr; + if (isa_and_nonnull(scaleAttr)) + return scaleAttr; + + if (!inAttr || !scaleAttr || !resultType) + return {}; + + if (auto inFloat = dyn_cast(inAttr)) { + auto scaleFloat = dyn_cast(scaleAttr); + if (!scaleFloat) + return {}; + std::optional result = + calculate(inFloat.getValue(), scaleFloat.getValue()); + if (!result) + return {}; + return FloatAttr::get(resultType, *result); + } + + auto inElements = dyn_cast(inAttr); + auto scaleElements = dyn_cast(scaleAttr); + auto shapedResultType = dyn_cast(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 result = + calculate(inElements.getSplatValue(), + scaleElements.getSplatValue()); + if (!result) + return {}; + return DenseElementsAttr::get(shapedResultType, *result); + } + + SmallVector results; + results.reserve(inElements.getNumElements()); + auto scaleIt = scaleElements.begin(); + for (const APFloat &in : inElements) { + std::optional 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, which for a f8E8M0FNU NaN +/// gives a value encoded as an infinity: that format has no significand bits, +/// so there is nowhere to put the quiet bit. +static std::optional +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(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(getElementTypeOrSelf(getType())); + const llvm::fltSemantics &resSemantics = resElemType.getFloatSemantics(); + return foldScalingCastOp( + adaptor.getIn(), adaptor.getScale(), getType(), + [&resSemantics](const APFloat &in, + const APFloat &scale) -> std::optional { + if (scale.isNaN()) + return getScalingCastNaN(resSemantics); + FailureOr inExt = convertFloatValue(in, resSemantics); + FailureOr 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(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(getElementTypeOrSelf(getIn().getType())); + auto resElemType = cast(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 { + if (scale.isNaN()) + return getScalingCastNaN(resSemantics); + FailureOr scaleExt = convertFloatValue(scale, inSemantics); + if (failed(scaleExt)) + return std::nullopt; + APFloat quotient(in); + quotient.divide(*scaleExt, kDefaultRoundingMode); + FailureOr result = + convertFloatValue(quotient, resSemantics, roundingMode); + if (failed(result)) + return std::nullopt; + return *result; + }); +} + bool arith::ScalingTruncFOp::areCastCompatible(TypeRange inputs, TypeRange outputs) { return checkWidthChangeCast(inputs.front(), outputs); diff --git a/mlir/test/Dialect/Arith/canonicalize.mlir b/mlir/test/Dialect/Arith/canonicalize.mlir index 0c06aa6e861a4..3d4c997dae9aa 100644 --- a/mlir/test/Dialect/Arith/canonicalize.mlir +++ b/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