[Mlir-commits] [mlir] [mlir][arith] Fold scaling_extf and scaling_truncf with constant operands (PR #215123)
曾鈜寬 Tseng Hung Kuan
llvmlistbot at llvm.org
Sun Aug 9 11:01:51 PDT 2026
https://github.com/Tim096 created https://github.com/llvm/llvm-project/pull/215123
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 |
>From 66b9c0eef2196b9634ca1c38691aef92921ca30b Mon Sep 17 00:00:00 2001
From: Hung-Kuan Tseng <p76091014 at gs.ncku.edu.tw>
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: `-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 |
---
.../include/mlir/Dialect/Arith/IR/ArithOps.td | 2 +
mlir/lib/Dialect/Arith/IR/ArithOps.cpp | 133 ++++++++++++++++
mlir/test/Dialect/Arith/canonicalize.mlir | 145 ++++++++++++++++++
3 files changed, 280 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 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
More information about the Mlir-commits
mailing list