[Mlir-commits] [mlir] [mlir][arith] Expand ops for F8E4M3FN and F8E5M2 type. (PR #216653)
Arun Thangamani
llvmlistbot at llvm.org
Tue Aug 25 01:34:04 PDT 2026
https://github.com/arun-thmn updated https://github.com/llvm/llvm-project/pull/216653
>From da280d6ceae4877fb722c686c9e13f068f472333 Mon Sep 17 00:00:00 2001
From: Arun Thangamani <arun.thangamani at intel.com>
Date: Mon, 17 Aug 2026 00:54:52 -0700
Subject: [PATCH 1/4] fix the llvm instrincs selection for F8E4M3FN and F8E5M2
---
.../mlir/Dialect/Arith/Transforms/Passes.h | 7 +
.../mlir/Dialect/Arith/Transforms/Passes.td | 4 +
.../Dialect/Arith/Transforms/ExpandOps.cpp | 277 ++++++++++++++++++
3 files changed, 288 insertions(+)
diff --git a/mlir/include/mlir/Dialect/Arith/Transforms/Passes.h b/mlir/include/mlir/Dialect/Arith/Transforms/Passes.h
index 938c77c9dd5ef..26832ba652680 100644
--- a/mlir/include/mlir/Dialect/Arith/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/Arith/Transforms/Passes.h
@@ -62,6 +62,13 @@ void populateExpandBFloat16Patterns(RewritePatternSet &patterns);
/// Add patterns to expand Arith f4e2m1 patterns to lower level bitcasts/shifts.
void populateExpandF4E2M1Patterns(RewritePatternSet &patterns);
+/// Add patterns to expand Arith f8e5m2 patterns to lower level bitcasts/shifts.
+void populateExpandF8E5M2Patterns(RewritePatternSet &patterns);
+
+/// Add patterns to expand Arith f8e4m3fn patterns to lower level
+/// bitcasts/shifts.
+void populateExpandF8E4M3FNPatterns(RewritePatternSet &patterns);
+
/// Add patterns to expand Arith f8e8m0 patterns to lower level bitcasts/shifts.
void populateExpandF8E8M0Patterns(RewritePatternSet &patterns);
diff --git a/mlir/include/mlir/Dialect/Arith/Transforms/Passes.td b/mlir/include/mlir/Dialect/Arith/Transforms/Passes.td
index 34c28db9cb161..98e549f53040a 100644
--- a/mlir/include/mlir/Dialect/Arith/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Arith/Transforms/Passes.td
@@ -21,6 +21,10 @@ def ArithExpandOpsPass : Pass<"arith-expand"> {
"Enable the F8E8M0 expansion patterns">,
Option<"includeF4E2M1", "include-f4e2m1", "bool", /*default=*/"false",
"Enable the F4E2M1 expansion patterns">,
+ Option<"includeF8E5M2", "include-f8e5m2", "bool", /*default=*/"false",
+ "Enable the F8E5M2 expansion patterns">,
+ Option<"includeF8E4M3FN", "include-f8e4m3fn", "bool", /*default=*/"false",
+ "Enable the F8E4M3FN expansion patterns">,
Option<"includeFlushDenormals", "include-flush-denormals", "bool",
/*default=*/"false",
"Enable expansion of `arith.flush_denormals` on IEEE-like "
diff --git a/mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp b/mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp
index 02ed6ccd87a42..a7b4c8125f3cf 100644
--- a/mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp
+++ b/mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp
@@ -662,6 +662,261 @@ struct F8E8M0TruncFOpConverter : public OpRewritePattern<arith::TruncFOp> {
}
};
+/// Expand an ExtF from F8E5M2. F8E5M2 uses a 5-bit exponent with bias 15 and a
+/// 2-bit mantissa, i.e. it is bit-for-bit the high byte of an IEEE F16 value
+/// (same exponent field and bias), including infinities and NaNs. The exact
+/// F16 value is therefore obtained by placing the 8 F8E5M2 bits into the high
+/// byte of an i16. The F16 value is then converted to the requested result
+/// type with the native (LLVM-lowerable) extf/truncf.
+struct F8E5M2ExtFOpConverter : public OpRewritePattern<arith::ExtFOp> {
+ using Base::Base;
+ LogicalResult matchAndRewrite(arith::ExtFOp op,
+ PatternRewriter &rewriter) const final {
+ ImplicitLocOpBuilder b(op.getLoc(), rewriter);
+ Value operand = op.getOperand();
+ Type operandTy = operand.getType();
+ Type resultTy = op.getType();
+ Type operandETy = getElementTypeOrSelf(operandTy);
+ Type resultETy = getElementTypeOrSelf(resultTy);
+
+ if (!llvm::isa<Float8E5M2Type>(operandETy))
+ return rewriter.notifyMatchFailure(op, "not a ext of F8E5M2");
+
+ Type i8Ty = cloneToShapedType(operandTy, b.getI8Type());
+ Type i16Ty = cloneToShapedType(operandTy, b.getI16Type());
+ Type f16Ty = cloneToShapedType(operandTy, b.getF16Type());
+
+ Value bitcast = arith::BitcastOp::create(b, i8Ty, operand);
+ Value exti = arith::ExtUIOp::create(b, i16Ty, bitcast);
+ Value c8 = createConst(op.getLoc(), i16Ty, 8, rewriter);
+ Value f16Bits = arith::ShLIOp::create(b, exti, c8);
+ Value f16 = arith::BitcastOp::create(b, f16Ty, f16Bits);
+
+ Value result = f16;
+ if (!resultETy.isF16()) {
+ if (resultETy.getIntOrFloatBitWidth() < 16)
+ result = arith::TruncFOp::create(b, resultTy, f16, nullptr,
+ op.getFastmathAttr());
+ else
+ result = arith::ExtFOp::create(b, resultTy, f16, op.getFastmathAttr());
+ }
+ rewriter.replaceOp(op, result);
+ return success();
+ }
+};
+
+/// Expand a TruncF to F8E5M2. The input is first reduced to F16 (which shares
+/// the F8E5M2 exponent layout and bias) using the native truncf, then the low
+/// 8 mantissa bits of the F16 value are dropped with round-to-nearest-even.
+/// The rounding-bias trick is borrowed from the BF16 converter: adding the
+/// bias may carry into the exponent field, which is exactly the desired
+/// behavior since F16 and F8E5M2 share the same exponent bias. NaN is handled
+/// separately.
+struct F8E5M2TruncFOpConverter : public OpRewritePattern<arith::TruncFOp> {
+ using Base::Base;
+ LogicalResult matchAndRewrite(arith::TruncFOp op,
+ PatternRewriter &rewriter) const final {
+ ImplicitLocOpBuilder b(op.getLoc(), rewriter);
+ Value operand = op.getOperand();
+ Type operandTy = operand.getType();
+ Type resultTy = op.getType();
+ Type operandETy = getElementTypeOrSelf(operandTy);
+ Type resultETy = getElementTypeOrSelf(resultTy);
+
+ if (!llvm::isa<Float8E5M2Type>(resultETy))
+ return rewriter.notifyMatchFailure(op, "not a trunc to F8E5M2");
+ if (op.getRoundingmodeAttr())
+ return rewriter.notifyMatchFailure(
+ op, "only applicable to default rounding mode.");
+
+ Type i8Ty = cloneToShapedType(operandTy, b.getI8Type());
+ Type i16Ty = cloneToShapedType(operandTy, b.getI16Type());
+ Type f16Ty = cloneToShapedType(operandTy, b.getF16Type());
+
+ Value h16 = operand;
+ if (!operandETy.isF16())
+ h16 = arith::TruncFOp::create(b, f16Ty, operand, nullptr,
+ op.getFastmathAttr());
+
+ Value isNan =
+ arith::CmpFOp::create(b, arith::CmpFPredicate::UNE, h16, h16);
+ Value h16Bits = arith::BitcastOp::create(b, i16Ty, h16);
+ // Rounding bias constants for dropping the low 8 mantissa bits.
+ Value c7F = createConst(op.getLoc(), i16Ty, 0x7f, rewriter);
+ Value c8 = createConst(op.getLoc(), i16Ty, 8, rewriter);
+ Value c1 = createConst(op.getLoc(), i16Ty, 1, rewriter);
+ Value bit8 = arith::AndIOp::create(
+ b, arith::ShRUIOp::create(b, h16Bits, c8), c1);
+ Value roundingBias = arith::AddIOp::create(b, bit8, c7F);
+ Value biased = arith::AddIOp::create(b, h16Bits, roundingBias);
+ Value biasedAndShifted = arith::ShRUIOp::create(b, biased, c8);
+ Value normalCaseResult = arith::TruncIOp::create(b, i8Ty, biasedAndShifted);
+ // Quiet NaN for F8E5M2 (exponent all ones, mantissa MSB set).
+ Value cNan = createConst(op.getLoc(), i8Ty, 0x7e, rewriter);
+ Value select = arith::SelectOp::create(b, isNan, cNan, normalCaseResult);
+ Value result = arith::BitcastOp::create(b, resultTy, select);
+ rewriter.replaceOp(op, result);
+ return success();
+ }
+};
+
+/// Expand an ExtF from F8E4M3FN. F8E4M3FN uses a 4-bit exponent with bias 7, a
+/// 3-bit mantissa, no infinities and a single NaN encoding (S.1111.111). The
+/// 7 magnitude bits (EEEE.MMM) are placed into the high mantissa/exponent bits
+/// of an F16 by shifting left by 7, producing an F16 whose value equals the
+/// desired magnitude scaled by 2^-8 (the F16 bias is 15 while F8E4M3FN's is
+/// 7). Multiplying by 256 in F32 recovers the true magnitude for both normal
+/// and subnormal inputs. The sign bit and the NaN encoding are re-applied
+/// explicitly.
+struct F8E4M3FNExtFOpConverter : public OpRewritePattern<arith::ExtFOp> {
+ using Base::Base;
+ LogicalResult matchAndRewrite(arith::ExtFOp op,
+ PatternRewriter &rewriter) const final {
+ ImplicitLocOpBuilder b(op.getLoc(), rewriter);
+ Value operand = op.getOperand();
+ Type operandTy = operand.getType();
+ Type resultTy = op.getType();
+ Type operandETy = getElementTypeOrSelf(operandTy);
+ Type resultETy = getElementTypeOrSelf(resultTy);
+
+ if (!llvm::isa<Float8E4M3FNType>(operandETy))
+ return rewriter.notifyMatchFailure(op, "not a ext of F8E4M3FN");
+
+ Type i8Ty = cloneToShapedType(operandTy, b.getI8Type());
+ Type i16Ty = cloneToShapedType(operandTy, b.getI16Type());
+ Type i32Ty = cloneToShapedType(operandTy, b.getI32Type());
+ Type f16Ty = cloneToShapedType(operandTy, b.getF16Type());
+ Type f32Ty = cloneToShapedType(operandTy, b.getF32Type());
+
+ Value bits = arith::BitcastOp::create(b, i8Ty, operand);
+ Value c7F8 = createConst(op.getLoc(), i8Ty, 0x7f, rewriter);
+ Value mag8 = arith::AndIOp::create(b, bits, c7F8);
+ // Build an F16 equal to the magnitude times 2^-8.
+ Value mag16 = arith::ExtUIOp::create(b, i16Ty, mag8);
+ Value c7 = createConst(op.getLoc(), i16Ty, 7, rewriter);
+ Value g16Bits = arith::ShLIOp::create(b, mag16, c7);
+ Value g16 = arith::BitcastOp::create(b, f16Ty, g16Bits);
+ Value gF32 = arith::ExtFOp::create(b, f32Ty, g16, op.getFastmathAttr());
+ Value c256 =
+ createFloatConst(op.getLoc(), f32Ty, APFloat(256.0f), rewriter);
+ Value magF32 = arith::MulFOp::create(b, gF32, c256, op.getFastmathAttr());
+ // Re-apply the sign bit into the F32 result.
+ Value magI32 = arith::BitcastOp::create(b, i32Ty, magF32);
+ Value c80I8 = createConst(op.getLoc(), i8Ty, 0x80, rewriter);
+ Value sign8 = arith::AndIOp::create(b, bits, c80I8);
+ Value sign32 = arith::ExtUIOp::create(b, i32Ty, sign8);
+ Value c24 = createConst(op.getLoc(), i32Ty, 24, rewriter);
+ Value signBit = arith::ShLIOp::create(b, sign32, c24);
+ Value signedI32 = arith::OrIOp::create(b, magI32, signBit);
+ Value signedF32 = arith::BitcastOp::create(b, f32Ty, signedI32);
+ // NaN encoding is magnitude == 0x7f.
+ Value isNan =
+ arith::CmpIOp::create(b, arith::CmpIPredicate::eq, mag8, c7F8);
+ Value cNan32 = createConst(op.getLoc(), i32Ty, 0x7fc00000, rewriter);
+ Value nanSigned = arith::OrIOp::create(b, cNan32, signBit);
+ Value nanF32 = arith::BitcastOp::create(b, f32Ty, nanSigned);
+ Value resultF32 = arith::SelectOp::create(b, isNan, nanF32, signedF32);
+
+ Value result = resultF32;
+ if (!resultETy.isF32()) {
+ if (resultETy.getIntOrFloatBitWidth() < 32)
+ result = arith::TruncFOp::create(b, resultTy, resultF32, nullptr,
+ op.getFastmathAttr());
+ else
+ result =
+ arith::ExtFOp::create(b, resultTy, resultF32, op.getFastmathAttr());
+ }
+ rewriter.replaceOp(op, result);
+ return success();
+ }
+};
+
+/// Expand a TruncF to F8E4M3FN. The input magnitude is clamped to the F8E4M3FN
+/// maximum (448), scaled by 2^-8, and reduced to F16 (undoing the bias
+/// difference so the F16 value equals the magnitude times 2^-8). The low 7
+/// bits of the F16 encoding are then dropped with round-to-nearest-even to
+/// recover the 7 magnitude bits. Overflow into the NaN encoding is prevented
+/// by clamping, the sign is re-applied, and a NaN input maps to the F8E4M3FN
+/// NaN encoding.
+struct F8E4M3FNTruncFOpConverter : public OpRewritePattern<arith::TruncFOp> {
+ using Base::Base;
+ LogicalResult matchAndRewrite(arith::TruncFOp op,
+ PatternRewriter &rewriter) const final {
+ ImplicitLocOpBuilder b(op.getLoc(), rewriter);
+ Value operand = op.getOperand();
+ Type operandTy = operand.getType();
+ Type resultTy = op.getType();
+ Type operandETy = getElementTypeOrSelf(operandTy);
+ Type resultETy = getElementTypeOrSelf(resultTy);
+
+ if (!llvm::isa<Float8E4M3FNType>(resultETy))
+ return rewriter.notifyMatchFailure(op, "not a trunc to F8E4M3FN");
+ if (op.getRoundingmodeAttr())
+ return rewriter.notifyMatchFailure(
+ op, "only applicable to default rounding mode.");
+
+ Type i8Ty = cloneToShapedType(operandTy, b.getI8Type());
+ Type i16Ty = cloneToShapedType(operandTy, b.getI16Type());
+ Type i32Ty = cloneToShapedType(operandTy, b.getI32Type());
+ Type f16Ty = cloneToShapedType(operandTy, b.getF16Type());
+ Type f32Ty = cloneToShapedType(operandTy, b.getF32Type());
+
+ Value f32 = operand;
+ if (!operandETy.isF32()) {
+ if (operandETy.getIntOrFloatBitWidth() < 32)
+ f32 = arith::ExtFOp::create(b, f32Ty, operand, op.getFastmathAttr());
+ else
+ f32 = arith::TruncFOp::create(b, f32Ty, operand, nullptr,
+ op.getFastmathAttr());
+ }
+
+ Value isNan = arith::CmpFOp::create(b, arith::CmpFPredicate::UNE, f32, f32);
+ // Split sign and magnitude.
+ Value f32Bits = arith::BitcastOp::create(b, i32Ty, f32);
+ Value cSignMask = createConst(op.getLoc(), i32Ty, 0x80000000, rewriter);
+ Value cAbsMask = createConst(op.getLoc(), i32Ty, 0x7fffffff, rewriter);
+ Value signBits = arith::AndIOp::create(b, f32Bits, cSignMask);
+ Value absBits = arith::AndIOp::create(b, f32Bits, cAbsMask);
+ Value absF32 = arith::BitcastOp::create(b, f32Ty, absBits);
+ // Clamp to the F8E4M3FN maximum magnitude (448) then scale by 2^-8.
+ Value cMax = createFloatConst(op.getLoc(), f32Ty, APFloat(448.0f), rewriter);
+ absF32 = arith::MinNumFOp::create(b, absF32, cMax);
+ Value cInv256 =
+ createFloatConst(op.getLoc(), f32Ty, APFloat(0.00390625f), rewriter);
+ Value scaled = arith::MulFOp::create(b, absF32, cInv256, nullptr);
+ Value h16 = arith::TruncFOp::create(b, f16Ty, scaled, nullptr,
+ op.getFastmathAttr());
+ Value h16Bits = arith::BitcastOp::create(b, i16Ty, h16);
+ // Drop the low 7 bits with round-to-nearest-even.
+ Value c3F = createConst(op.getLoc(), i16Ty, 0x3f, rewriter);
+ Value c7 = createConst(op.getLoc(), i16Ty, 7, rewriter);
+ Value c1 = createConst(op.getLoc(), i16Ty, 1, rewriter);
+ Value bit7 =
+ arith::AndIOp::create(b, arith::ShRUIOp::create(b, h16Bits, c7), c1);
+ Value roundingBias = arith::AddIOp::create(b, bit7, c3F);
+ Value biased = arith::AddIOp::create(b, h16Bits, roundingBias);
+ Value shifted = arith::ShRUIOp::create(b, biased, c7);
+ Value mag8 = arith::TruncIOp::create(b, i8Ty, shifted);
+ Value c7F8 = createConst(op.getLoc(), i8Ty, 0x7f, rewriter);
+ mag8 = arith::AndIOp::create(b, mag8, c7F8);
+ // Never emit the NaN encoding (0x7f) for a finite input.
+ Value c7E8 = createConst(op.getLoc(), i8Ty, 0x7e, rewriter);
+ Value isOverflow =
+ arith::CmpIOp::create(b, arith::CmpIPredicate::ugt, mag8, c7E8);
+ mag8 = arith::SelectOp::create(b, isOverflow, c7E8, mag8);
+ // Re-apply the sign.
+ Value c24 = createConst(op.getLoc(), i32Ty, 24, rewriter);
+ Value sign8 =
+ arith::TruncIOp::create(b, i8Ty, arith::ShRUIOp::create(b, signBits, c24));
+ Value res8 = arith::OrIOp::create(b, mag8, sign8);
+ Value cNan8 = createConst(op.getLoc(), i8Ty, 0x7f, rewriter);
+ Value res = arith::SelectOp::create(b, isNan, cNan8, res8);
+ Value result = arith::BitcastOp::create(b, resultTy, res);
+ rewriter.replaceOp(op, result);
+ return success();
+ }
+};
+
struct ScalingExtFOpConverter : public OpRewritePattern<arith::ScalingExtFOp> {
using Base::Base;
LogicalResult matchAndRewrite(arith::ScalingExtFOp op,
@@ -874,6 +1129,10 @@ struct ArithExpandOpsPass
arith::populateExpandF8E8M0Patterns(patterns);
if (includeF4E2M1)
arith::populateExpandF4E2M1Patterns(patterns);
+ if (includeF8E5M2)
+ arith::populateExpandF8E5M2Patterns(patterns);
+ if (includeF8E4M3FN)
+ arith::populateExpandF8E4M3FNPatterns(patterns);
if (includeFlushDenormals) {
arith::populateExpandFlushDenormalsPatterns(patterns);
// Only IEEE-like floating-point types are expanded by the pattern;
@@ -900,6 +1159,10 @@ struct ArithExpandOpsPass
legalTypes &= !llvm::isa<Float8E8M0FNUType>(inETy);
if (includeF4E2M1)
legalTypes &= !llvm::isa<Float4E2M1FNType>(inETy);
+ if (includeF8E5M2)
+ legalTypes &= !llvm::isa<Float8E5M2Type>(inETy);
+ if (includeF8E4M3FN)
+ legalTypes &= !llvm::isa<Float8E4M3FNType>(inETy);
return legalTypes;
});
@@ -914,6 +1177,10 @@ struct ArithExpandOpsPass
legalTypes &= !(llvm::isa<Float8E8M0FNUType>(outETy));
if (includeF4E2M1)
legalTypes &= !llvm::isa<Float4E2M1FNType>(outETy);
+ if (includeF8E5M2)
+ legalTypes &= !llvm::isa<Float8E5M2Type>(outETy);
+ if (includeF8E4M3FN)
+ legalTypes &= !llvm::isa<Float8E4M3FNType>(outETy);
return legalTypes;
});
@@ -943,6 +1210,16 @@ void mlir::arith::populateExpandF4E2M1Patterns(RewritePatternSet &patterns) {
patterns.getContext());
}
+void mlir::arith::populateExpandF8E5M2Patterns(RewritePatternSet &patterns) {
+ patterns.add<F8E5M2ExtFOpConverter, F8E5M2TruncFOpConverter>(
+ patterns.getContext());
+}
+
+void mlir::arith::populateExpandF8E4M3FNPatterns(RewritePatternSet &patterns) {
+ patterns.add<F8E4M3FNExtFOpConverter, F8E4M3FNTruncFOpConverter>(
+ patterns.getContext());
+}
+
void mlir::arith::populateExpandF8E8M0Patterns(RewritePatternSet &patterns) {
patterns.add<F8E8M0ExtFOpConverter, F8E8M0TruncFOpConverter>(
patterns.getContext());
>From 0609ac294314c544d9cf45d5837c0ff407b62a72 Mon Sep 17 00:00:00 2001
From: Arun Thangamani <arun.thangamani at intel.com>
Date: Mon, 17 Aug 2026 01:15:30 -0700
Subject: [PATCH 2/4] fix the clang error
---
.../Dialect/Arith/Transforms/ExpandOps.cpp | 82 +++++++++----------
1 file changed, 40 insertions(+), 42 deletions(-)
diff --git a/mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp b/mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp
index a7b4c8125f3cf..d2b670576d474 100644
--- a/mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp
+++ b/mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp
@@ -738,15 +738,14 @@ struct F8E5M2TruncFOpConverter : public OpRewritePattern<arith::TruncFOp> {
h16 = arith::TruncFOp::create(b, f16Ty, operand, nullptr,
op.getFastmathAttr());
- Value isNan =
- arith::CmpFOp::create(b, arith::CmpFPredicate::UNE, h16, h16);
+ Value isNan = arith::CmpFOp::create(b, arith::CmpFPredicate::UNE, h16, h16);
Value h16Bits = arith::BitcastOp::create(b, i16Ty, h16);
// Rounding bias constants for dropping the low 8 mantissa bits.
Value c7F = createConst(op.getLoc(), i16Ty, 0x7f, rewriter);
Value c8 = createConst(op.getLoc(), i16Ty, 8, rewriter);
Value c1 = createConst(op.getLoc(), i16Ty, 1, rewriter);
- Value bit8 = arith::AndIOp::create(
- b, arith::ShRUIOp::create(b, h16Bits, c8), c1);
+ Value bit8 =
+ arith::AndIOp::create(b, arith::ShRUIOp::create(b, h16Bits, c8), c1);
Value roundingBias = arith::AddIOp::create(b, bit8, c7F);
Value biased = arith::AddIOp::create(b, h16Bits, roundingBias);
Value biasedAndShifted = arith::ShRUIOp::create(b, biased, c8);
@@ -879,7 +878,8 @@ struct F8E4M3FNTruncFOpConverter : public OpRewritePattern<arith::TruncFOp> {
Value absBits = arith::AndIOp::create(b, f32Bits, cAbsMask);
Value absF32 = arith::BitcastOp::create(b, f32Ty, absBits);
// Clamp to the F8E4M3FN maximum magnitude (448) then scale by 2^-8.
- Value cMax = createFloatConst(op.getLoc(), f32Ty, APFloat(448.0f), rewriter);
+ Value cMax =
+ createFloatConst(op.getLoc(), f32Ty, APFloat(448.0f), rewriter);
absF32 = arith::MinNumFOp::create(b, absF32, cMax);
Value cInv256 =
createFloatConst(op.getLoc(), f32Ty, APFloat(0.00390625f), rewriter);
@@ -906,8 +906,8 @@ struct F8E4M3FNTruncFOpConverter : public OpRewritePattern<arith::TruncFOp> {
mag8 = arith::SelectOp::create(b, isOverflow, c7E8, mag8);
// Re-apply the sign.
Value c24 = createConst(op.getLoc(), i32Ty, 24, rewriter);
- Value sign8 =
- arith::TruncIOp::create(b, i8Ty, arith::ShRUIOp::create(b, signBits, c24));
+ Value sign8 = arith::TruncIOp::create(
+ b, i8Ty, arith::ShRUIOp::create(b, signBits, c24));
Value res8 = arith::OrIOp::create(b, mag8, sign8);
Value cNan8 = createConst(op.getLoc(), i8Ty, 0x7f, rewriter);
Value res = arith::SelectOp::create(b, isNan, cNan8, res8);
@@ -1148,41 +1148,39 @@ struct ArithExpandOpsPass
});
}
- target.addDynamicallyLegalOp<arith::ExtFOp>(
- [=](arith::ExtFOp op) {
- Type inETy = getElementTypeOrSelf(op.getOperand().getType());
- Type outETy = getElementTypeOrSelf(op.getType());
- bool legalTypes = true;
- if (includeBf16)
- legalTypes &= !(inETy.isBF16() && outETy.isF32());
- if (includeF8E8M0)
- legalTypes &= !llvm::isa<Float8E8M0FNUType>(inETy);
- if (includeF4E2M1)
- legalTypes &= !llvm::isa<Float4E2M1FNType>(inETy);
- if (includeF8E5M2)
- legalTypes &= !llvm::isa<Float8E5M2Type>(inETy);
- if (includeF8E4M3FN)
- legalTypes &= !llvm::isa<Float8E4M3FNType>(inETy);
- return legalTypes;
- });
-
- target.addDynamicallyLegalOp<arith::TruncFOp>(
- [=](arith::TruncFOp op) {
- Type inETy = getElementTypeOrSelf(op.getOperand().getType());
- Type outETy = getElementTypeOrSelf(op.getType());
- bool legalTypes = true;
- if (includeBf16)
- legalTypes &= !(inETy.isF32() && outETy.isBF16());
- if (includeF8E8M0)
- legalTypes &= !(llvm::isa<Float8E8M0FNUType>(outETy));
- if (includeF4E2M1)
- legalTypes &= !llvm::isa<Float4E2M1FNType>(outETy);
- if (includeF8E5M2)
- legalTypes &= !llvm::isa<Float8E5M2Type>(outETy);
- if (includeF8E4M3FN)
- legalTypes &= !llvm::isa<Float8E4M3FNType>(outETy);
- return legalTypes;
- });
+ target.addDynamicallyLegalOp<arith::ExtFOp>([=](arith::ExtFOp op) {
+ Type inETy = getElementTypeOrSelf(op.getOperand().getType());
+ Type outETy = getElementTypeOrSelf(op.getType());
+ bool legalTypes = true;
+ if (includeBf16)
+ legalTypes &= !(inETy.isBF16() && outETy.isF32());
+ if (includeF8E8M0)
+ legalTypes &= !llvm::isa<Float8E8M0FNUType>(inETy);
+ if (includeF4E2M1)
+ legalTypes &= !llvm::isa<Float4E2M1FNType>(inETy);
+ if (includeF8E5M2)
+ legalTypes &= !llvm::isa<Float8E5M2Type>(inETy);
+ if (includeF8E4M3FN)
+ legalTypes &= !llvm::isa<Float8E4M3FNType>(inETy);
+ return legalTypes;
+ });
+
+ target.addDynamicallyLegalOp<arith::TruncFOp>([=](arith::TruncFOp op) {
+ Type inETy = getElementTypeOrSelf(op.getOperand().getType());
+ Type outETy = getElementTypeOrSelf(op.getType());
+ bool legalTypes = true;
+ if (includeBf16)
+ legalTypes &= !(inETy.isF32() && outETy.isBF16());
+ if (includeF8E8M0)
+ legalTypes &= !(llvm::isa<Float8E8M0FNUType>(outETy));
+ if (includeF4E2M1)
+ legalTypes &= !llvm::isa<Float4E2M1FNType>(outETy);
+ if (includeF8E5M2)
+ legalTypes &= !llvm::isa<Float8E5M2Type>(outETy);
+ if (includeF8E4M3FN)
+ legalTypes &= !llvm::isa<Float8E4M3FNType>(outETy);
+ return legalTypes;
+ });
// clang-format on
if (failed(applyPartialConversion(getOperation(), target,
>From c14bc54ae62e47c7d3fc9d7b276ec249833c104c Mon Sep 17 00:00:00 2001
From: Arun Thangamani <arun.thangamani at intel.com>
Date: Tue, 18 Aug 2026 07:33:21 -0700
Subject: [PATCH 3/4] test-cases inclusion
---
mlir/test/Dialect/Arith/expand-ops.mlir | 223 +++++++++++++++++++++++-
1 file changed, 222 insertions(+), 1 deletion(-)
diff --git a/mlir/test/Dialect/Arith/expand-ops.mlir b/mlir/test/Dialect/Arith/expand-ops.mlir
index 20f00b82505e7..4fe7ba75c429f 100644
--- a/mlir/test/Dialect/Arith/expand-ops.mlir
+++ b/mlir/test/Dialect/Arith/expand-ops.mlir
@@ -1,4 +1,4 @@
-// RUN: mlir-opt %s -arith-expand="include-bf16=true include-f8e8m0=true include-f4e2m1=true include-min-max-f=true include-min-max-i=true" -verify-diagnostics -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -arith-expand="include-bf16=true include-f8e8m0=true include-f4e2m1=true include-f8e5m2=true include-f8e4m3fn=true include-min-max-f=true include-min-max-i=true" -verify-diagnostics -split-input-file | FileCheck %s
// RUN: mlir-opt %s -arith-expand -split-input-file -verify-diagnostics | FileCheck %s --check-prefix=SCHECK
// Test ceil divide with signed integer
@@ -648,3 +648,224 @@ func.func @extf_vector_f4E2M1FN_to_f32(%arg0 : vector<4xf4E2M1FN>) -> vector<4xf
// CHECK-LABEL: @extf_vector_f4E2M1FN_to_f32
// CHECK-NOT: arith.extf
+
+// -----
+
+func.func @extf_f8E5M2_to_f32(%arg0 : f8E5M2) -> f32 {
+ %0 = arith.extf %arg0 : f8E5M2 to f32
+ return %0 : f32
+}
+
+// CHECK-LABEL: @extf_f8E5M2_to_f32
+// CHECK-SAME: %[[ARG0:.+]]: f8E5M2
+// CHECK: %[[BITCAST:.+]] = arith.bitcast %[[ARG0]] : f8E5M2 to i8
+// CHECK: %[[EXTUI:.+]] = arith.extui %[[BITCAST]] : i8 to i16
+// CHECK: %[[C8:.+]] = arith.constant 8 : i16
+// CHECK: %[[SHL:.+]] = arith.shli %[[EXTUI]], %[[C8]] : i16
+// CHECK: %[[F16:.+]] = arith.bitcast %[[SHL]] : i16 to f16
+// CHECK: %[[RESULT:.+]] = arith.extf %[[F16]] : f16 to f32
+// CHECK: return %[[RESULT]]
+
+// -----
+
+// The F16 result matches the F8E5M2 pivot type, so no trailing extf/truncf is
+// emitted.
+func.func @extf_f8E5M2_to_f16(%arg0 : f8E5M2) -> f16 {
+ %0 = arith.extf %arg0 : f8E5M2 to f16
+ return %0 : f16
+}
+
+// CHECK-LABEL: @extf_f8E5M2_to_f16
+// CHECK: %[[BITCAST:.+]] = arith.bitcast %arg0 : f8E5M2 to i8
+// CHECK: %[[EXTUI:.+]] = arith.extui %[[BITCAST]] : i8 to i16
+// CHECK: %[[C8:.+]] = arith.constant 8 : i16
+// CHECK: %[[SHL:.+]] = arith.shli %[[EXTUI]], %[[C8]] : i16
+// CHECK: %[[RESULT:.+]] = arith.bitcast %[[SHL]] : i16 to f16
+// CHECK-NOT: arith.extf
+// CHECK: return %[[RESULT]]
+
+// -----
+
+func.func @extf_vector_f8E5M2_to_f32(%arg0 : vector<4xf8E5M2>) -> vector<4xf32> {
+ %0 = arith.extf %arg0 : vector<4xf8E5M2> to vector<4xf32>
+ return %0 : vector<4xf32>
+}
+
+// CHECK-LABEL: @extf_vector_f8E5M2_to_f32
+// CHECK: arith.bitcast %arg0 : vector<4xf8E5M2> to vector<4xi8>
+// CHECK: arith.extf %{{.+}} : vector<4xf16> to vector<4xf32>
+// CHECK: return
+
+// -----
+
+func.func @truncf_f32_to_f8E5M2(%arg0 : f32) -> f8E5M2 {
+ %0 = arith.truncf %arg0 : f32 to f8E5M2
+ return %0 : f8E5M2
+}
+
+// CHECK-LABEL: @truncf_f32_to_f8E5M2
+// CHECK: %[[H16:.+]] = arith.truncf %arg0 : f32 to f16
+// CHECK: %[[ISNAN:.+]] = arith.cmpf une, %[[H16]], %[[H16]] : f16
+// CHECK: %[[BITS:.+]] = arith.bitcast %[[H16]] : f16 to i16
+// CHECK: %[[C7F:.+]] = arith.constant 127 : i16
+// CHECK: %[[C8:.+]] = arith.constant 8 : i16
+// CHECK: %[[C1:.+]] = arith.constant 1 : i16
+// CHECK: %[[SHR:.+]] = arith.shrui %[[BITS]], %[[C8]] : i16
+// CHECK: %[[BIT8:.+]] = arith.andi %[[SHR]], %[[C1]] : i16
+// CHECK: %[[BIAS:.+]] = arith.addi %[[BIT8]], %[[C7F]] : i16
+// CHECK: %[[BIASED:.+]] = arith.addi %[[BITS]], %[[BIAS]] : i16
+// CHECK: %[[BSHIFT:.+]] = arith.shrui %[[BIASED]], %[[C8]] : i16
+// CHECK: %[[NORMAL:.+]] = arith.trunci %[[BSHIFT]] : i16 to i8
+// CHECK: %[[CNAN:.+]] = arith.constant 126 : i8
+// CHECK: %[[SEL:.+]] = arith.select %[[ISNAN]], %[[CNAN]], %[[NORMAL]] : i8
+// CHECK: %[[RESULT:.+]] = arith.bitcast %[[SEL]] : i8 to f8E5M2
+// CHECK: return %[[RESULT]]
+
+// -----
+
+// The F16 operand matches the F8E5M2 pivot type, so no leading truncf to F16 is
+// emitted.
+func.func @truncf_f16_to_f8E5M2(%arg0 : f16) -> f8E5M2 {
+ %0 = arith.truncf %arg0 : f16 to f8E5M2
+ return %0 : f8E5M2
+}
+
+// CHECK-LABEL: @truncf_f16_to_f8E5M2
+// CHECK-NOT: arith.truncf
+// CHECK: %[[ISNAN:.+]] = arith.cmpf une, %arg0, %arg0 : f16
+// CHECK: %[[BITS:.+]] = arith.bitcast %arg0 : f16 to i16
+// CHECK: %[[RESULT:.+]] = arith.bitcast %{{.+}} : i8 to f8E5M2
+// CHECK: return %[[RESULT]]
+
+// -----
+
+func.func @truncf_vector_f32_to_f8E5M2(%arg0 : vector<4xf32>) -> vector<4xf8E5M2> {
+ %0 = arith.truncf %arg0 : vector<4xf32> to vector<4xf8E5M2>
+ return %0 : vector<4xf8E5M2>
+}
+
+// CHECK-LABEL: @truncf_vector_f32_to_f8E5M2
+// CHECK: arith.truncf %arg0 : vector<4xf32> to vector<4xf16>
+// CHECK: arith.bitcast %{{.+}} : vector<4xi8> to vector<4xf8E5M2>
+// CHECK: return
+
+// -----
+
+func.func @extf_f8E4M3FN_to_f32(%arg0 : f8E4M3FN) -> f32 {
+ %0 = arith.extf %arg0 : f8E4M3FN to f32
+ return %0 : f32
+}
+
+// CHECK-LABEL: @extf_f8E4M3FN_to_f32
+// CHECK: %[[BITS:.+]] = arith.bitcast %arg0 : f8E4M3FN to i8
+// CHECK: %[[C7F:.+]] = arith.constant 127 : i8
+// CHECK: %[[MAG8:.+]] = arith.andi %[[BITS]], %[[C7F]] : i8
+// CHECK: %[[MAG16:.+]] = arith.extui %[[MAG8]] : i8 to i16
+// CHECK: %[[C7:.+]] = arith.constant 7 : i16
+// CHECK: %[[G16BITS:.+]] = arith.shli %[[MAG16]], %[[C7]] : i16
+// CHECK: %[[G16:.+]] = arith.bitcast %[[G16BITS]] : i16 to f16
+// CHECK: %[[GF32:.+]] = arith.extf %[[G16]] : f16 to f32
+// CHECK: %[[C256:.+]] = arith.constant 2.560000e+02 : f32
+// CHECK: %[[MAGF32:.+]] = arith.mulf %[[GF32]], %[[C256]] : f32
+// CHECK: %[[MAGI32:.+]] = arith.bitcast %[[MAGF32]] : f32 to i32
+// CHECK: %[[C80:.+]] = arith.constant -128 : i8
+// CHECK: %[[SIGN8:.+]] = arith.andi %[[BITS]], %[[C80]] : i8
+// CHECK: %[[SIGN32:.+]] = arith.extui %[[SIGN8]] : i8 to i32
+// CHECK: %[[C24:.+]] = arith.constant 24 : i32
+// CHECK: %[[SIGNBIT:.+]] = arith.shli %[[SIGN32]], %[[C24]] : i32
+// CHECK: %[[SIGNED:.+]] = arith.ori %[[MAGI32]], %[[SIGNBIT]] : i32
+// CHECK: %[[SIGNEDF32:.+]] = arith.bitcast %[[SIGNED]] : i32 to f32
+// CHECK: %[[ISNAN:.+]] = arith.cmpi eq, %[[MAG8]], %[[C7F]] : i8
+// CHECK: %[[CNAN:.+]] = arith.constant 2143289344 : i32
+// CHECK: %[[NANSIGNED:.+]] = arith.ori %[[CNAN]], %[[SIGNBIT]] : i32
+// CHECK: %[[NANF32:.+]] = arith.bitcast %[[NANSIGNED]] : i32 to f32
+// CHECK: %[[RESULT:.+]] = arith.select %[[ISNAN]], %[[NANF32]], %[[SIGNEDF32]] : f32
+// CHECK: return %[[RESULT]]
+
+// -----
+
+// The F16 result is narrower than the F32 pivot, so a trailing truncf to F16 is
+// emitted.
+func.func @extf_f8E4M3FN_to_f16(%arg0 : f8E4M3FN) -> f16 {
+ %0 = arith.extf %arg0 : f8E4M3FN to f16
+ return %0 : f16
+}
+
+// CHECK-LABEL: @extf_f8E4M3FN_to_f16
+// CHECK: %[[BITS:.+]] = arith.bitcast %arg0 : f8E4M3FN to i8
+// CHECK: %[[SEL:.+]] = arith.select %{{.+}}, %{{.+}}, %{{.+}} : f32
+// CHECK: %[[RESULT:.+]] = arith.truncf %[[SEL]] : f32 to f16
+// CHECK: return %[[RESULT]]
+
+// -----
+
+func.func @extf_vector_f8E4M3FN_to_f32(%arg0 : vector<4xf8E4M3FN>) -> vector<4xf32> {
+ %0 = arith.extf %arg0 : vector<4xf8E4M3FN> to vector<4xf32>
+ return %0 : vector<4xf32>
+}
+
+// CHECK-LABEL: @extf_vector_f8E4M3FN_to_f32
+// CHECK: arith.bitcast %arg0 : vector<4xf8E4M3FN> to vector<4xi8>
+// CHECK: arith.mulf
+// CHECK: return
+
+// -----
+
+func.func @truncf_f32_to_f8E4M3FN(%arg0 : f32) -> f8E4M3FN {
+ %0 = arith.truncf %arg0 : f32 to f8E4M3FN
+ return %0 : f8E4M3FN
+}
+
+// CHECK-LABEL: @truncf_f32_to_f8E4M3FN
+// CHECK: %[[ISNAN:.+]] = arith.cmpf une, %arg0, %arg0 : f32
+// CHECK: %[[BITS:.+]] = arith.bitcast %arg0 : f32 to i32
+// CHECK: %[[CSIGN:.+]] = arith.constant -2147483648 : i32
+// CHECK: %[[CABS:.+]] = arith.constant 2147483647 : i32
+// CHECK: %[[SIGNBITS:.+]] = arith.andi %[[BITS]], %[[CSIGN]] : i32
+// CHECK: %[[ABSBITS:.+]] = arith.andi %[[BITS]], %[[CABS]] : i32
+// CHECK: %[[ABSF32:.+]] = arith.bitcast %[[ABSBITS]] : i32 to f32
+// CHECK: %[[CMAX:.+]] = arith.constant 4.480000e+02 : f32
+// The clamp to the maximum magnitude is emitted as arith.minnumf, which the
+// include-min-max-f flag further lowers to cmpf/select.
+// CHECK: %[[LT:.+]] = arith.cmpf ult, %[[ABSF32]], %[[CMAX]] : f32
+// CHECK: %[[MIN:.+]] = arith.select %[[LT]], %[[ABSF32]], %[[CMAX]] : f32
+// CHECK: %[[UNO:.+]] = arith.cmpf uno, %[[ABSF32]], %[[ABSF32]] : f32
+// CHECK: %[[CLAMPED:.+]] = arith.select %[[UNO]], %[[CMAX]], %[[MIN]] : f32
+// CHECK: %[[CINV:.+]] = arith.constant 3.906250e-03 : f32
+// CHECK: %[[SCALED:.+]] = arith.mulf %[[CLAMPED]], %[[CINV]] : f32
+// CHECK: %[[H16:.+]] = arith.truncf %[[SCALED]] : f32 to f16
+// CHECK: %[[H16BITS:.+]] = arith.bitcast %[[H16]] : f16 to i16
+// CHECK: %[[C3F:.+]] = arith.constant 63 : i16
+// CHECK: %[[C7:.+]] = arith.constant 7 : i16
+// CHECK: %[[C1:.+]] = arith.constant 1 : i16
+// CHECK: %[[SHR:.+]] = arith.shrui %[[H16BITS]], %[[C7]] : i16
+// CHECK: %[[BIT7:.+]] = arith.andi %[[SHR]], %[[C1]] : i16
+// CHECK: %[[BIAS:.+]] = arith.addi %[[BIT7]], %[[C3F]] : i16
+// CHECK: %[[BIASED:.+]] = arith.addi %[[H16BITS]], %[[BIAS]] : i16
+// CHECK: %[[SHIFTED:.+]] = arith.shrui %[[BIASED]], %[[C7]] : i16
+// CHECK: %[[MAG8:.+]] = arith.trunci %[[SHIFTED]] : i16 to i8
+// CHECK: %[[C7F8:.+]] = arith.constant 127 : i8
+// CHECK: %[[MAGMASK:.+]] = arith.andi %[[MAG8]], %[[C7F8]] : i8
+// CHECK: %[[C7E:.+]] = arith.constant 126 : i8
+// CHECK: %[[OVF:.+]] = arith.cmpi ugt, %[[MAGMASK]], %[[C7E]] : i8
+// CHECK: %[[CLAMPMAG:.+]] = arith.select %[[OVF]], %[[C7E]], %[[MAGMASK]] : i8
+// CHECK: %[[C24:.+]] = arith.constant 24 : i32
+// CHECK: %[[SIGNSHR:.+]] = arith.shrui %[[SIGNBITS]], %[[C24]] : i32
+// CHECK: %[[SIGN8:.+]] = arith.trunci %[[SIGNSHR]] : i32 to i8
+// CHECK: %[[RES8:.+]] = arith.ori %[[CLAMPMAG]], %[[SIGN8]] : i8
+// CHECK: %[[CNAN:.+]] = arith.constant 127 : i8
+// CHECK: %[[RES:.+]] = arith.select %[[ISNAN]], %[[CNAN]], %[[RES8]] : i8
+// CHECK: %[[RESULT:.+]] = arith.bitcast %[[RES]] : i8 to f8E4M3FN
+// CHECK: return %[[RESULT]]
+
+// -----
+
+func.func @truncf_vector_f32_to_f8E4M3FN(%arg0 : vector<4xf32>) -> vector<4xf8E4M3FN> {
+ %0 = arith.truncf %arg0 : vector<4xf32> to vector<4xf8E4M3FN>
+ return %0 : vector<4xf8E4M3FN>
+}
+
+// CHECK-LABEL: @truncf_vector_f32_to_f8E4M3FN
+// CHECK: arith.constant dense<4.480000e+02> : vector<4xf32>
+// CHECK: arith.bitcast %{{.+}} : vector<4xi8> to vector<4xf8E4M3FN>
+// CHECK: return
>From 3836531cc62c95b767e276b65fe07bb66469a5fa Mon Sep 17 00:00:00 2001
From: Arun Thangamani <arun.thangamani at intel.com>
Date: Tue, 25 Aug 2026 01:33:29 -0700
Subject: [PATCH 4/4] fix f8E4M3FN to NaN on overflow
---
.../Dialect/Arith/Transforms/ExpandOps.cpp | 32 +++++++++++--------
mlir/test/Dialect/Arith/expand-ops.mlir | 11 ++++---
2 files changed, 24 insertions(+), 19 deletions(-)
diff --git a/mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp b/mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp
index d2b670576d474..b2482662fa3cb 100644
--- a/mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp
+++ b/mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp
@@ -830,13 +830,13 @@ struct F8E4M3FNExtFOpConverter : public OpRewritePattern<arith::ExtFOp> {
}
};
-/// Expand a TruncF to F8E4M3FN. The input magnitude is clamped to the F8E4M3FN
-/// maximum (448), scaled by 2^-8, and reduced to F16 (undoing the bias
-/// difference so the F16 value equals the magnitude times 2^-8). The low 7
-/// bits of the F16 encoding are then dropped with round-to-nearest-even to
-/// recover the 7 magnitude bits. Overflow into the NaN encoding is prevented
-/// by clamping, the sign is re-applied, and a NaN input maps to the F8E4M3FN
-/// NaN encoding.
+/// Expand a TruncF to F8E4M3FN. The magnitude is scaled by 2^-8 and reduced to
+/// F16 (undoing the bias difference so the F16 value equals the magnitude times
+/// 2^-8). The low 7 bits of the F16 encoding are then dropped with
+/// round-to-nearest-even to recover the 7 magnitude bits. F8E4M3FN has no
+/// infinity, so any input that overflows the maximum representable magnitude
+/// (448), as well as infinities and NaNs, maps to the F8E4M3FN NaN encoding to
+/// match the LLVM APFloat NanOnly overflow behavior.
struct F8E4M3FNTruncFOpConverter : public OpRewritePattern<arith::TruncFOp> {
using Base::Base;
LogicalResult matchAndRewrite(arith::TruncFOp op,
@@ -877,7 +877,14 @@ struct F8E4M3FNTruncFOpConverter : public OpRewritePattern<arith::TruncFOp> {
Value signBits = arith::AndIOp::create(b, f32Bits, cSignMask);
Value absBits = arith::AndIOp::create(b, f32Bits, cAbsMask);
Value absF32 = arith::BitcastOp::create(b, f32Ty, absBits);
- // Clamp to the F8E4M3FN maximum magnitude (448) then scale by 2^-8.
+ // F8E4M3FN has no infinity: a magnitude above the round-to-nearest-even
+ // overflow boundary (464 = 448 + half an ulp) or an infinity maps to NaN.
+ Value cOverflow =
+ createFloatConst(op.getLoc(), f32Ty, APFloat(464.0f), rewriter);
+ Value isOverflow =
+ arith::CmpFOp::create(b, arith::CmpFPredicate::OGT, absF32, cOverflow);
+ // Clamp to the F8E4M3FN maximum magnitude (448) so the finite path stays
+ // well-defined; overflowing inputs are replaced by NaN below.
Value cMax =
createFloatConst(op.getLoc(), f32Ty, APFloat(448.0f), rewriter);
absF32 = arith::MinNumFOp::create(b, absF32, cMax);
@@ -899,18 +906,15 @@ struct F8E4M3FNTruncFOpConverter : public OpRewritePattern<arith::TruncFOp> {
Value mag8 = arith::TruncIOp::create(b, i8Ty, shifted);
Value c7F8 = createConst(op.getLoc(), i8Ty, 0x7f, rewriter);
mag8 = arith::AndIOp::create(b, mag8, c7F8);
- // Never emit the NaN encoding (0x7f) for a finite input.
- Value c7E8 = createConst(op.getLoc(), i8Ty, 0x7e, rewriter);
- Value isOverflow =
- arith::CmpIOp::create(b, arith::CmpIPredicate::ugt, mag8, c7E8);
- mag8 = arith::SelectOp::create(b, isOverflow, c7E8, mag8);
// Re-apply the sign.
Value c24 = createConst(op.getLoc(), i32Ty, 24, rewriter);
Value sign8 = arith::TruncIOp::create(
b, i8Ty, arith::ShRUIOp::create(b, signBits, c24));
Value res8 = arith::OrIOp::create(b, mag8, sign8);
+ // NaN input or an overflowing/infinite magnitude maps to the NaN encoding.
+ Value isNanOrOverflow = arith::OrIOp::create(b, isNan, isOverflow);
Value cNan8 = createConst(op.getLoc(), i8Ty, 0x7f, rewriter);
- Value res = arith::SelectOp::create(b, isNan, cNan8, res8);
+ Value res = arith::SelectOp::create(b, isNanOrOverflow, cNan8, res8);
Value result = arith::BitcastOp::create(b, resultTy, res);
rewriter.replaceOp(op, result);
return success();
diff --git a/mlir/test/Dialect/Arith/expand-ops.mlir b/mlir/test/Dialect/Arith/expand-ops.mlir
index 4fe7ba75c429f..0342d6a7036cf 100644
--- a/mlir/test/Dialect/Arith/expand-ops.mlir
+++ b/mlir/test/Dialect/Arith/expand-ops.mlir
@@ -824,6 +824,9 @@ func.func @truncf_f32_to_f8E4M3FN(%arg0 : f32) -> f8E4M3FN {
// CHECK: %[[SIGNBITS:.+]] = arith.andi %[[BITS]], %[[CSIGN]] : i32
// CHECK: %[[ABSBITS:.+]] = arith.andi %[[BITS]], %[[CABS]] : i32
// CHECK: %[[ABSF32:.+]] = arith.bitcast %[[ABSBITS]] : i32 to f32
+// F8E4M3FN has no infinity, so an overflowing/infinite magnitude maps to NaN.
+// CHECK: %[[COVF:.+]] = arith.constant 4.640000e+02 : f32
+// CHECK: %[[ISOVF:.+]] = arith.cmpf ogt, %[[ABSF32]], %[[COVF]] : f32
// CHECK: %[[CMAX:.+]] = arith.constant 4.480000e+02 : f32
// The clamp to the maximum magnitude is emitted as arith.minnumf, which the
// include-min-max-f flag further lowers to cmpf/select.
@@ -846,15 +849,13 @@ func.func @truncf_f32_to_f8E4M3FN(%arg0 : f32) -> f8E4M3FN {
// CHECK: %[[MAG8:.+]] = arith.trunci %[[SHIFTED]] : i16 to i8
// CHECK: %[[C7F8:.+]] = arith.constant 127 : i8
// CHECK: %[[MAGMASK:.+]] = arith.andi %[[MAG8]], %[[C7F8]] : i8
-// CHECK: %[[C7E:.+]] = arith.constant 126 : i8
-// CHECK: %[[OVF:.+]] = arith.cmpi ugt, %[[MAGMASK]], %[[C7E]] : i8
-// CHECK: %[[CLAMPMAG:.+]] = arith.select %[[OVF]], %[[C7E]], %[[MAGMASK]] : i8
// CHECK: %[[C24:.+]] = arith.constant 24 : i32
// CHECK: %[[SIGNSHR:.+]] = arith.shrui %[[SIGNBITS]], %[[C24]] : i32
// CHECK: %[[SIGN8:.+]] = arith.trunci %[[SIGNSHR]] : i32 to i8
-// CHECK: %[[RES8:.+]] = arith.ori %[[CLAMPMAG]], %[[SIGN8]] : i8
+// CHECK: %[[RES8:.+]] = arith.ori %[[MAGMASK]], %[[SIGN8]] : i8
+// CHECK: %[[NANOVF:.+]] = arith.ori %[[ISNAN]], %[[ISOVF]] : i1
// CHECK: %[[CNAN:.+]] = arith.constant 127 : i8
-// CHECK: %[[RES:.+]] = arith.select %[[ISNAN]], %[[CNAN]], %[[RES8]] : i8
+// CHECK: %[[RES:.+]] = arith.select %[[NANOVF]], %[[CNAN]], %[[RES8]] : i8
// CHECK: %[[RESULT:.+]] = arith.bitcast %[[RES]] : i8 to f8E4M3FN
// CHECK: return %[[RESULT]]
More information about the Mlir-commits
mailing list