[Mlir-commits] [mlir] [mlir][ArithToAMDGPU] Only use the scaled conversions for f8E8M0FNU scales (PR #221386)
Hung-Kuan Tseng
llvmlistbot at llvm.org
Sun Nov 1 06:00:20 PST 2026
https://github.com/Tim096 updated https://github.com/llvm/llvm-project/pull/221386
>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 cf44c2a0033ac4..8f8429ec07cf9b 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 fe5b0520e37c16..c64b0ef96e5c07 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 4c707680378158..5518bc8f89bc75 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