[Mlir-commits] [mlir] [mlir][ArithToAMDGPU] Only use the scaled conversions for f8E8M0FNU scales (PR #221386)
Hung-Kuan Tseng
llvmlistbot at llvm.org
Fri Sep 4 20:53:27 PDT 2026
https://github.com/Tim096 created https://github.com/llvm/llvm-project/pull/221386
The gfx950 scaled conversion instructions (`v_cvt_scalef32_*`) take
their scale as an f32 and use only its exponent (bits 31:23). That
implements `arith.scaling_extf` and `arith.scaling_truncf` only when the
scale is `f8E8M0FNU`, whose value is its exponent. The lowering accepted
any float scale and extended or truncated it to f32:
```mlir
%0 = arith.scaling_extf %in, %scale : vector<4xf4E2M1FN>, vector<4xf32> to vector<4xf32>
// -convert-arith-to-amdgpu=chipset=gfx950 today: the f32 scale goes straight
// into amdgpu.scaled_ext_packed, which reads its exponent only
```
For a scale that is not `f8E8M0FNU` the exponent-only read is one
particular rounding of the scale (toward zero), and what such a scale
means is still open in #215295. The folder (#215123) already only
handles `f8E8M0FNU` scales for that reason. This makes the conversion
take the same line: it matches only `f8E8M0FNU` scales and leaves the
others to `arith-expand`, as @krzysz00 suggested on the issue.
Also from that discussion: when the scale is an f32 truncated to
`f8E8M0FNU` with `toward_zero`, the truncation is exactly what the
instruction does itself, so the f32 is fed to it directly and the
truncf/extf round trip goes away.
Tests: the four `long_*_broadcast` tests used an f32 scale and are the
only coverage of the multi-slice loop, so they keep their shape and take
an `f8E8M0FNU` scale. New tests cover f32, f16 and bf16 scales being left
alone, the `toward_zero` fold, and a default-rounding truncation not
being folded.
No in-tree pipeline runs this pass. A pipeline that feeds it scales that
are not `f8E8M0FNU` needs `arith-expand` after it, as the XeVM pipeline
already does.
This is the AMDGPU side of #217892: once the expansion uses the scale's
value, the exponent-only instruction path is the one that has to be
restricted to the scale type it is exact for.
>From 3cc500a59ae02d31a204d7ca351dfc4cf6f1f3b3 Mon Sep 17 00:00:00 2001
From: Hung-Kuan Tseng <tseng.tim096 at gmail.com>
Date: Sat, 5 Sep 2026 11:53:15 +0800
Subject: [PATCH] [mlir][ArithToAMDGPU] Only use the scaled conversions for
f8E8M0FNU scales
The gfx950 scaled conversion instructions take their scale as an f32 and
use only its exponent. That implements `arith.scaling_extf` and
`arith.scaling_truncf` only when the scale is f8E8M0FNU, whose value is
its exponent; the lowering accepted any float scale and extended or
truncated it to f32. For any other scale type the exponent-only read is
one particular rounding of the scale, and what such a scale means is
still open (#215295); the folder already only handles f8E8M0FNU scales.
Leave the other scale types to `arith-expand`.
When the scale is an f32 truncated to f8E8M0FNU toward zero, the
truncation is what the instruction does itself, so feed it the f32 and
skip the truncf/extf round trip.
The four tests with an f32 scale are the only coverage of the
multi-slice loop, so they keep their shape and take an f8E8M0FNU scale;
f32, f16 and bf16 scales become negative tests.
---
.../ArithToAMDGPU/ArithToAMDGPU.cpp | 49 ++++++++------
.../ArithToAMDGPU/scaling-extf.mlir | 64 +++++++++++++++++--
.../ArithToAMDGPU/scaling-truncf.mlir | 51 +++++++++++++--
3 files changed, 134 insertions(+), 30 deletions(-)
diff --git a/mlir/lib/Conversion/ArithToAMDGPU/ArithToAMDGPU.cpp b/mlir/lib/Conversion/ArithToAMDGPU/ArithToAMDGPU.cpp
index cf44c2a0033ac..8f8429ec07cf9 100644
--- a/mlir/lib/Conversion/ArithToAMDGPU/ArithToAMDGPU.cpp
+++ b/mlir/lib/Conversion/ArithToAMDGPU/ArithToAMDGPU.cpp
@@ -441,6 +441,29 @@ static Value getOriginalVectorValue(Value value) {
return current;
}
+/// The scaled conversion instructions take their scale as an f32 and use only
+/// its exponent, so they implement `arith.scaling_extf`/`arith.scaling_truncf`
+/// only for an f8E8M0FNU scale, whose value is its exponent. Returns that scale
+/// extended to f32, or failure for any other scale type. When the scale is an
+/// f32 truncated to f8E8M0FNU toward zero, the truncation is what the
+/// instruction does itself, so the f32 is used directly.
+static FailureOr<Value> getF32Scale(PatternRewriter &rewriter, Location loc,
+ Value scale) {
+ if (!isa<Float8E8M0FNUType>(getElementTypeOrSelf(scale)))
+ return failure();
+
+ if (auto truncOp = scale.getDefiningOp<arith::TruncFOp>())
+ if (getElementTypeOrSelf(truncOp.getIn()).isF32() &&
+ truncOp.getRoundingmode() == arith::RoundingMode::toward_zero)
+ return truncOp.getIn();
+
+ Type f32 = rewriter.getF32Type();
+ Type scaleF32Type = f32;
+ if (auto scaleVecType = dyn_cast<VectorType>(scale.getType()))
+ scaleF32Type = VectorType::get(scaleVecType.getShape(), f32);
+ return arith::ExtFOp::create(rewriter, loc, scaleF32Type, scale).getResult();
+}
+
LogicalResult
ScalingExtFRewritePattern::matchAndRewrite(arith::ScalingExtFOp op,
PatternRewriter &rewriter) const {
@@ -451,15 +474,12 @@ ScalingExtFRewritePattern::matchAndRewrite(arith::ScalingExtFOp op,
Value scale = op.getScale();
Value out = op.getOut();
- Type f32 = rewriter.getF32Type();
Type inType = getElementTypeOrSelf(in);
- Type scaleType = getElementTypeOrSelf(scale);
Type outType = getElementTypeOrSelf(out);
int64_t opInWidth = 32 / inType.getIntOrFloatBitWidth();
VectorType outVecType = dyn_cast<VectorType>(out.getType());
- VectorType scaleVecType = dyn_cast<VectorType>(scale.getType());
if (outVecType && outVecType.isScalable())
return failure();
@@ -469,12 +489,10 @@ ScalingExtFRewritePattern::matchAndRewrite(arith::ScalingExtFOp op,
isa<RankedTensorType>(scale.getType()))
return failure();
- Type scaleF32Type =
- scaleVecType ? VectorType::get(scaleVecType.getShape(), f32) : f32;
- if (scaleType.getIntOrFloatBitWidth() < 32)
- scale = arith::ExtFOp::create(rewriter, loc, scaleF32Type, scale);
- else if (scaleType.getIntOrFloatBitWidth() > 32)
- scale = arith::TruncFOp::create(rewriter, loc, scaleF32Type, scale);
+ FailureOr<Value> scaleF32 = getF32Scale(rewriter, loc, scale);
+ if (failed(scaleF32))
+ return rewriter.notifyMatchFailure(op, "scale is not f8E8M0FNU");
+ scale = *scaleF32;
VectorType extScaleResultType = VectorType::get(opOutWidth, outType);
@@ -572,13 +590,10 @@ ScalingTruncFRewritePattern::matchAndRewrite(arith::ScalingTruncFOp op,
Value scale = op.getScale();
Value out = op.getOut();
- Type f32 = rewriter.getF32Type();
Type inType = getElementTypeOrSelf(in);
- Type scaleType = getElementTypeOrSelf(scale);
Type outType = getElementTypeOrSelf(out);
VectorType outVecType = dyn_cast<VectorType>(out.getType());
- VectorType scaleVecType = dyn_cast<VectorType>(scale.getType());
if (outVecType && outVecType.isScalable())
return failure();
@@ -587,12 +602,10 @@ ScalingTruncFRewritePattern::matchAndRewrite(arith::ScalingTruncFOp op,
isa<RankedTensorType>(scale.getType()))
return failure();
- Type scaleF32Type =
- scaleVecType ? VectorType::get(scaleVecType.getShape(), f32) : f32;
- if (scaleType.getIntOrFloatBitWidth() < 32)
- scale = arith::ExtFOp::create(rewriter, loc, scaleF32Type, scale);
- else if (scaleType.getIntOrFloatBitWidth() > 32)
- scale = arith::TruncFOp::create(rewriter, loc, scaleF32Type, scale);
+ FailureOr<Value> scaleF32 = getF32Scale(rewriter, loc, scale);
+ if (failed(scaleF32))
+ return rewriter.notifyMatchFailure(op, "scale is not f8E8M0FNU");
+ scale = *scaleF32;
Value zero = arith::ConstantOp::create(rewriter, loc, outType,
rewriter.getFloatAttr(outType, 0.0));
diff --git a/mlir/test/Conversion/ArithToAMDGPU/scaling-extf.mlir b/mlir/test/Conversion/ArithToAMDGPU/scaling-extf.mlir
index fe5b0520e37c1..c64b0ef96e5c0 100644
--- a/mlir/test/Conversion/ArithToAMDGPU/scaling-extf.mlir
+++ b/mlir/test/Conversion/ArithToAMDGPU/scaling-extf.mlir
@@ -262,9 +262,9 @@ func.func @conversion_scalar(%in: f8E5M2, %scale: f8E8M0FNU) -> f32 {
// CHECK-COUNT-4: amdgpu.scaled_ext_packed %{{.+}}[3]
// CHECK-NOT: amdgpu.scaled_ext_packed
// CHECK: return
-func.func @long_fp4_broadcast(%in: vector<32xf4E2M1FN>, %scale: f32) -> vector<32xf32> {
- %splat = vector.broadcast %scale : f32 to vector<32xf32>
- %ext = arith.scaling_extf %in, %splat : vector<32xf4E2M1FN>, vector<32xf32> to vector<32xf32>
+func.func @long_fp4_broadcast(%in: vector<32xf4E2M1FN>, %scale: f8E8M0FNU) -> vector<32xf32> {
+ %splat = vector.broadcast %scale : f8E8M0FNU to vector<32xf8E8M0FNU>
+ %ext = arith.scaling_extf %in, %splat : vector<32xf4E2M1FN>, vector<32xf8E8M0FNU> to vector<32xf32>
return %ext : vector<32xf32>
}
@@ -274,8 +274,60 @@ func.func @long_fp4_broadcast(%in: vector<32xf4E2M1FN>, %scale: f32) -> vector<3
// CHECK-COUNT-8: amdgpu.scaled_ext_packed %{{.+}}[1]
// CHECK-NOT: amdgpu.scaled_ext_packed
// CHECK: return
-func.func @long_fp8_broadcast(%in: vector<32xf8E4M3FN>, %scale: f32) -> vector<32xf32> {
- %splat = vector.broadcast %scale : f32 to vector<32xf32>
- %ext = arith.scaling_extf %in, %splat : vector<32xf8E4M3FN>, vector<32xf32> to vector<32xf32>
+func.func @long_fp8_broadcast(%in: vector<32xf8E4M3FN>, %scale: f8E8M0FNU) -> vector<32xf32> {
+ %splat = vector.broadcast %scale : f8E8M0FNU to vector<32xf8E8M0FNU>
+ %ext = arith.scaling_extf %in, %splat : vector<32xf8E4M3FN>, vector<32xf8E8M0FNU> to vector<32xf32>
return %ext : vector<32xf32>
}
+
+// -----
+
+// The instruction only reads the exponent of its scale, so a scale that is not
+// f8E8M0FNU is left for arith-expand.
+// CHECK-LABEL: @scale_f32_not_converted
+// CHECK-NOT: amdgpu.scaled_ext_packed
+// CHECK: arith.scaling_extf
+// CHECK-NOT: amdgpu.scaled_ext_packed
+func.func @scale_f32_not_converted(%in: vector<4xf4E2M1FN>, %scale: vector<4xf32>) -> vector<4xf32> {
+ %ext = arith.scaling_extf %in, %scale : vector<4xf4E2M1FN>, vector<4xf32> to vector<4xf32>
+ return %ext : vector<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @scale_f16_not_converted
+// CHECK-NOT: amdgpu.scaled_ext_packed
+// CHECK: arith.scaling_extf
+// CHECK-NOT: amdgpu.scaled_ext_packed
+func.func @scale_f16_not_converted(%in: f8E5M2, %scale: f16) -> f32 {
+ %ext = arith.scaling_extf %in, %scale : f8E5M2, f16 to f32
+ return %ext : f32
+}
+
+// -----
+
+// Truncating an f32 to f8E8M0FNU toward zero keeps its exponent, which is what
+// the instruction reads, so the f32 is used directly.
+// CHECK-LABEL: @scale_truncf_toward_zero_folded
+// CHECK-NOT: arith.truncf
+// CHECK-NOT: arith.extf
+// CHECK: %[[SPLAT_IN:.+]] = vector.broadcast %arg0 : f8E5M2 to vector<1xf8E5M2>
+// CHECK-NEXT: amdgpu.scaled_ext_packed %[[SPLAT_IN]][0], %arg1 : vector<1xf8E5M2> to vector<2xf32>
+func.func @scale_truncf_toward_zero_folded(%in: f8E5M2, %scale: f32) -> f32 {
+ %scale_e8m0 = arith.truncf %scale toward_zero : f32 to f8E8M0FNU
+ %ext = arith.scaling_extf %in, %scale_e8m0 : f8E5M2, f8E8M0FNU to f32
+ return %ext : f32
+}
+
+// -----
+
+// Without an explicit rounding mode the truncation is not folded.
+// CHECK-LABEL: @scale_truncf_default_rounding_not_folded
+// CHECK: %[[SCALE:.+]] = arith.truncf %arg1 : f32 to f8E8M0FNU
+// CHECK: %[[SCALE_F32:.+]] = arith.extf %[[SCALE]] : f8E8M0FNU to f32
+// CHECK: amdgpu.scaled_ext_packed %{{.+}}[0], %[[SCALE_F32]]
+func.func @scale_truncf_default_rounding_not_folded(%in: f8E5M2, %scale: f32) -> f32 {
+ %scale_e8m0 = arith.truncf %scale : f32 to f8E8M0FNU
+ %ext = arith.scaling_extf %in, %scale_e8m0 : f8E5M2, f8E8M0FNU to f32
+ return %ext : f32
+}
diff --git a/mlir/test/Conversion/ArithToAMDGPU/scaling-truncf.mlir b/mlir/test/Conversion/ArithToAMDGPU/scaling-truncf.mlir
index 4c70768037815..5518bc8f89bc7 100644
--- a/mlir/test/Conversion/ArithToAMDGPU/scaling-truncf.mlir
+++ b/mlir/test/Conversion/ArithToAMDGPU/scaling-truncf.mlir
@@ -184,9 +184,9 @@ func.func @conversion_scalar(%in: f32, %scale: f8E8M0FNU) -> f8E5M2 {
// CHECK-COUNT-4: amdgpu.packed_scaled_trunc %{{.*}} into %{{.+}}[3]
// CHECK-NOT: amdgpu.packed_scaled_trunc
// CHECK: return
-func.func @long_fp4_broadcast(%in: vector<32xf32>, %scale: f32) -> vector<32xf4E2M1FN> {
- %splat = vector.broadcast %scale : f32 to vector<32xf32>
- %trunc = arith.scaling_truncf %in, %splat : vector<32xf32>, vector<32xf32> to vector<32xf4E2M1FN>
+func.func @long_fp4_broadcast(%in: vector<32xf32>, %scale: f8E8M0FNU) -> vector<32xf4E2M1FN> {
+ %splat = vector.broadcast %scale : f8E8M0FNU to vector<32xf8E8M0FNU>
+ %trunc = arith.scaling_truncf %in, %splat : vector<32xf32>, vector<32xf8E8M0FNU> to vector<32xf4E2M1FN>
return %trunc : vector<32xf4E2M1FN>
}
@@ -196,8 +196,47 @@ func.func @long_fp4_broadcast(%in: vector<32xf32>, %scale: f32) -> vector<32xf4E
// CHECK-COUNT-8: amdgpu.packed_scaled_trunc %{{.*}} into %{{.+}}[1]
// CHECK-NOT: amdgpu.packed_scaled_trunc
// CHECK: return
-func.func @long_fp8_broadcast(%in: vector<32xf32>, %scale: f32) -> vector<32xf8E4M3FN> {
- %splat = vector.broadcast %scale : f32 to vector<32xf32>
- %trunc = arith.scaling_truncf %in, %splat : vector<32xf32>, vector<32xf32> to vector<32xf8E4M3FN>
+func.func @long_fp8_broadcast(%in: vector<32xf32>, %scale: f8E8M0FNU) -> vector<32xf8E4M3FN> {
+ %splat = vector.broadcast %scale : f8E8M0FNU to vector<32xf8E8M0FNU>
+ %trunc = arith.scaling_truncf %in, %splat : vector<32xf32>, vector<32xf8E8M0FNU> to vector<32xf8E4M3FN>
return %trunc : vector<32xf8E4M3FN>
}
+
+// -----
+
+// The instruction only reads the exponent of its scale, so a scale that is not
+// f8E8M0FNU is left for arith-expand.
+// CHECK-LABEL: @scale_f32_not_converted
+// CHECK-NOT: amdgpu.packed_scaled_trunc
+// CHECK: arith.scaling_truncf
+// CHECK-NOT: amdgpu.packed_scaled_trunc
+func.func @scale_f32_not_converted(%in: vector<4xf32>, %scale: vector<4xf32>) -> vector<4xf4E2M1FN> {
+ %trunc = arith.scaling_truncf %in, %scale : vector<4xf32>, vector<4xf32> to vector<4xf4E2M1FN>
+ return %trunc : vector<4xf4E2M1FN>
+}
+
+// -----
+
+// CHECK-LABEL: @scale_bf16_not_converted
+// CHECK-NOT: amdgpu.packed_scaled_trunc
+// CHECK: arith.scaling_truncf
+// CHECK-NOT: amdgpu.packed_scaled_trunc
+func.func @scale_bf16_not_converted(%in: f32, %scale: bf16) -> f8E5M2 {
+ %trunc = arith.scaling_truncf %in, %scale : f32, bf16 to f8E5M2
+ return %trunc : f8E5M2
+}
+
+// -----
+
+// Truncating an f32 to f8E8M0FNU toward zero keeps its exponent, which is what
+// the instruction reads, so the f32 is used directly.
+// CHECK-LABEL: @scale_truncf_toward_zero_folded
+// CHECK-NOT: arith.truncf
+// CHECK-NOT: arith.extf
+// CHECK: %[[SPLAT_IN:.+]] = vector.broadcast %arg0 : f32 to vector<1xf32>
+// CHECK-NEXT: amdgpu.packed_scaled_trunc %[[SPLAT_IN]] into undef[0], %arg1
+func.func @scale_truncf_toward_zero_folded(%in: f32, %scale: f32) -> f8E5M2 {
+ %scale_e8m0 = arith.truncf %scale toward_zero : f32 to f8E8M0FNU
+ %trunc = arith.scaling_truncf %in, %scale_e8m0 : f32, f8E8M0FNU to f8E5M2
+ return %trunc : f8E5M2
+}
More information about the Mlir-commits
mailing list