[mlir][arith] Fold scaling_extf and scaling_truncf with constant operands - #215123
Open
Tim096 wants to merge 1 commit into
Open
[mlir][arith] Fold scaling_extf and scaling_truncf with constant operands#215123Tim096 wants to merge 1 commit into
Tim096 wants to merge 1 commit into
Conversation
Contributor
Author
|
@umangyadav (you introduced these ops in #141965) and @kuhar — could you take a look? One decision is worth flagging, because it is the part I am least sure about. The folder only fires when the scale already has element type The other thing to know when reading the NaN path: widening a |
|
@llvm/pr-subscribers-mlir-arith @llvm/pr-subscribers-mlir Author: 曾鈜寬 Tseng Hung Kuan (Tim096) ChangesAdd constant folders for the two MXFP scaling casts, mirroring the expansion Note the asymmetry the expansion already has: These were the only two ops in the dialect with neither a folder nor a As in Scales that are not already A NaN scale is built directly rather than by widening it, which for The input space of these types is small enough to enumerate, so every
Full diff: https://github.com/llvm/llvm-project/pull/215123.diff 3 Files Affected:
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 ff6a5d4a0c29a..73abdc68f7f0c 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<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, 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<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);
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
|
…ands 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: `-arith-expand` truncates wider scales to `f8E8M0FNU` first while `ArithToAMDGPU` reads their exponent field, and for `scale = 1.6 : f16` the two disagree (2.0 vs 1.0). A NaN scale is built directly rather than by widening it, which for `f8E8M0FNU` yields a value encoded as an infinity: that format has no significand bits to hold the quiet bit. 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 |
Tim096
force-pushed
the
arith-scaling-fold
branch
from
August 9, 2026 18:06
66b9c0e to
439d460
Compare
kuhar
requested review from
Muzammiluddin-Syed-ECE and
krzysz00
and removed request for
kuhar
August 9, 2026 23:44
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Add constant folders for the two MXFP scaling casts, mirroring the expansion
in
ExpandOps.cpp:Note the asymmetry the expansion already has:
scaling_extfwidens the scaleto the result type,
scaling_truncfto the type ofin.These were the only two ops in the dialect with neither a folder nor a
canonicalizer, so
-canonicalizeleft them alone even with both operandsconstant, while
-arith-expand -canonicalizefolded them away. Pipelines thatlower them to hardware instructions rather than expanding them, such as
ArithToAMDGPUandXeGPUToXeVM, never got the fold.As in
arith.extfandarith.truncf, the widening and narrowing steps gothrough
convertFloatValueand fold only when lossless, so aroundingmodeon
scaling_truncfnever changes a folded result. Poison propagates andshaped results are guarded with a static-shape check, as in the generic
folders.
Scales that are not already
f8E8M0FNUare left alone:-arith-expandtruncates wider scales to
f8E8M0FNUfirst whileArithToAMDGPUreads theirexponent field, and for
scale = 1.6 : f16the two disagree (2.0 vs 1.0).A NaN scale is built directly rather than by widening it, which for
f8E8M0FNUyields a value encoded as an infinity: that format has nosignificand bits to hold the quiet bit. 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 OCPMXFP semantics computed independently of LLVM, and -- for the decision to fold
or not -- against the losslessness rule above:
scaling_extff4E2M1FN x f8E8M0FNU -> f16scaling_extff4E2M1FN x f8E8M0FNU -> f32scaling_truncff16 x f8E8M0FNU -> f4E2M1FN