From 6319069bfd7b65a499180605f15c9a987d528888 Mon Sep 17 00:00:00 2001 From: Hung-Kuan Tseng Date: Mon, 10 Aug 2026 01:38:20 +0800 Subject: [PATCH] [mlir][arith] Fold scaling_extf and scaling_truncf with constant operands Add constant folders for the two MXFP scaling casts, mirroring the expansion in `ExpandOps.cpp`: ``` scaling_extf(in, scale) -> mulf(extf(in), extf(scale)) scaling_truncf(in, scale) -> truncf(in / extf(scale)) ``` Note the asymmetry the expansion already has: `scaling_extf` widens the scale to the result type, `scaling_truncf` to the type of `in`. These were the only two ops in the dialect with neither a folder nor a canonicalizer, so `-canonicalize` left them alone even with both operands constant, while `-arith-expand -canonicalize` folded them away. Pipelines that lower them to hardware instructions rather than expanding them, such as `ArithToAMDGPU` and `XeGPUToXeVM`, never got the fold. As in `arith.extf` and `arith.truncf`, the widening and narrowing steps go through `convertFloatValue` and fold only when lossless, so a `roundingmode` on `scaling_truncf` never changes a folded result. Poison propagates and shaped results are guarded with a static-shape check, as in the generic folders. Scales that are not already `f8E8M0FNU` are left alone, because what such a scale means is unsettled: the tree does not say whether truncating a wider scale to `f8E8M0FNU` rounds or takes its exponent, and `ArithToAMDGPU` passes it to an instruction that reads only the exponent. #215295 asks for a ruling. A NaN scale is built directly rather than by widening it, 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. It stays unfolded when the result type is finite-only, such as `f4E2M1FN`, which cannot represent a NaN at all. The input space of these types is small enough to enumerate, so every combination was checked against `-arith-expand -canonicalize`, against the OCP MXFP semantics computed independently of LLVM, and -- for the decision to fold or not -- against the losslessness rule above: | sweep | cases | folded | disagreements | | :--- | ---: | ---: | ---: | | `scaling_extf` f4E2M1FN x f8E8M0FNU -> f16 | 4096 | 656 | 0 | | `scaling_extf` f4E2M1FN x f8E8M0FNU -> f32 | 4096 | 4096 | 0 | | `scaling_truncf` f16 x f8E8M0FNU -> f4E2M1FN | 16777216 | 33390 | 0 | --- .../include/mlir/Dialect/Arith/IR/ArithOps.td | 2 + mlir/lib/Dialect/Arith/IR/ArithOps.cpp | 134 ++++++++++++++++ mlir/test/Dialect/Arith/canonicalize.mlir | 145 ++++++++++++++++++ 3 files changed, 281 insertions(+) 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..659399c0fc399 100644 --- a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp +++ b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp @@ -1822,6 +1822,109 @@ 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, 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 +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. What a wider scale means is +/// unsettled -- the tree does not say whether truncating one to f8E8M0FNU +/// rounds or takes its exponent -- so a folder should not settle it, see +/// https://github.com/llvm/llvm-project/issues/215295. +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 +2097,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 89632305fa99d..1f00117429c38 100644 --- a/mlir/test/Dialect/Arith/canonicalize.mlir +++ b/mlir/test/Dialect/Arith/canonicalize.mlir @@ -1139,6 +1139,79 @@ func.func @extFPVectorConstantE8M0NaN() -> vector<2xf32> { return %0 : vector<2xf32> } +// 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: what such a +// scale means is unsettled (https://github.com/llvm/llvm-project/issues/215295). +// 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 @@ -1452,6 +1525,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