[Mlir-commits] [mlir] e1b9d03 - [MLIR][Math] Fix math.ceil expansion to avoid undefined behavior on Inf/NaN (#170028)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Feb 11 09:36:37 PST 2026
Author: Hank
Date: 2026-02-11T12:36:32-05:00
New Revision: e1b9d033cdbd957a6a4adb684340298e98b1b06a
URL: https://github.com/llvm/llvm-project/commit/e1b9d033cdbd957a6a4adb684340298e98b1b06a
DIFF: https://github.com/llvm/llvm-project/commit/e1b9d033cdbd957a6a4adb684340298e98b1b06a.diff
LOG: [MLIR][Math] Fix math.ceil expansion to avoid undefined behavior on Inf/NaN (#170028)
Fixes #151786
The original `ceilf` expansion lowers to `fptosi`, which produces poison
for Inf, and any subsequent use leads to undefined behavior. This patch
adds a safe path, similar to the existing `round` expansion, for large
or special inputs and avoids the UB.
Added:
Modified:
mlir/lib/Dialect/Math/Transforms/ExpandOps.cpp
mlir/test/Dialect/Math/expand-math.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/Math/Transforms/ExpandOps.cpp b/mlir/lib/Dialect/Math/Transforms/ExpandOps.cpp
index 249a95cc7924a..4fc435533a1c5 100644
--- a/mlir/lib/Dialect/Math/Transforms/ExpandOps.cpp
+++ b/mlir/lib/Dialect/Math/Transforms/ExpandOps.cpp
@@ -232,6 +232,52 @@ static LogicalResult convertCeilOp(math::CeilOp op, PatternRewriter &rewriter) {
ImplicitLocOpBuilder b(op->getLoc(), rewriter);
Value operand = op.getOperand();
Type opType = operand.getType();
+ Type operandETy = getElementTypeOrSelf(opType);
+ FloatType floatTy = llvm::dyn_cast<FloatType>(operandETy);
+ const llvm::fltSemantics &semantics = floatTy.getFloatSemantics();
+
+ unsigned bitWidth = floatTy.getWidth();
+ unsigned mantissaWidth = floatTy.getFPMantissaWidth() - 1;
+ const int bias = (&semantics == &APFloat::Float8E8M0FNU())
+ ? -semantics.minExponent
+ : -(semantics.minExponent - 1);
+ bool hasNegativeZeroNaNEncoding =
+ (semantics.nanEncoding == llvm::fltNanEncoding::NegativeZero);
+
+ Type iTy = rewriter.getIntegerType(bitWidth);
+ if (auto shapedTy = dyn_cast<ShapedType>(opType))
+ iTy = shapedTy.clone(iTy);
+
+ // For IEEE-like floating-point formats with an unbiased exponent ≥
+ // `mantissaWidth` falls into one of these categories:
+ // - a large finite value (|x| ≥ 2^mantissaWidth), where all representable
+ // numbers are already integral, or
+ // - a special value (NaN or ±Inf), which also satisfies this exponent
+ // condition.
+ // For all such cases, `ceilf(x)` is defined to return `x` directly.
+ Value operandBitcast = arith::BitcastOp::create(b, iTy, operand);
+ Value cMask = createIntConst(
+ op->getLoc(), iTy, static_cast<int64_t>((1ull << (bitWidth - 1)) - 1), b);
+ Value unsignedBits = arith::AndIOp::create(b, operandBitcast, cMask);
+ Value cThreshold = createIntConst(
+ op->getLoc(), iTy,
+ static_cast<int64_t>((uint64_t(bias + mantissaWidth)) << mantissaWidth),
+ b);
+ Value isLargeExp = arith::CmpIOp::create(b, arith::CmpIPredicate::uge,
+ unsignedBits, cThreshold);
+ Value isSpecialValOrLargeVal = isLargeExp;
+
+ // In FNUZ-suffixed floating point, NaN is represented by a sign bit of 1 and
+ // all 0s in the exponent and mantissa, therefore requires an explicit check.
+ if (hasNegativeZeroNaNEncoding) {
+ Value cNegZeroBits = createIntConst(
+ op->getLoc(), iTy, static_cast<int64_t>(1ull << (bitWidth - 1)), b);
+ Value isNegZeroEncoding = arith::CmpIOp::create(
+ b, arith::CmpIPredicate::eq, operandBitcast, cNegZeroBits);
+ isSpecialValOrLargeVal =
+ arith::OrIOp::create(b, isLargeExp, isNegZeroEncoding);
+ }
+
Value fpFixedConvert = createTruncatedFPValue(operand, b);
// Creating constants for later use.
@@ -243,7 +289,8 @@ static LogicalResult convertCeilOp(math::CeilOp op, PatternRewriter &rewriter) {
Value incrValue =
arith::SelectOp::create(b, op->getLoc(), gtCheck, one, zero);
- Value ret = arith::AddFOp::create(b, opType, fpFixedConvert, incrValue);
+ Value add = arith::AddFOp::create(b, opType, fpFixedConvert, incrValue);
+ Value ret = arith::SelectOp::create(b, isSpecialValOrLargeVal, operand, add);
rewriter.replaceOp(op, ret);
return success();
}
diff --git a/mlir/test/Dialect/Math/expand-math.mlir b/mlir/test/Dialect/Math/expand-math.mlir
index 1d26c826e8d6b..5f2843045b885 100644
--- a/mlir/test/Dialect/Math/expand-math.mlir
+++ b/mlir/test/Dialect/Math/expand-math.mlir
@@ -143,15 +143,21 @@ func.func @fmaf_func(%a: f64, %b: f64, %c: f64) -> f64 {
// CHECK-LABEL: func @ceilf_func
// CHECK-SAME: ([[ARG0:%.+]]: f64) -> f64
func.func @ceilf_func(%a: f64) -> f64 {
- // CHECK-DAG: [[CST:%.+]] = arith.constant 0.000
- // CHECK-DAG: [[CST_0:%.+]] = arith.constant 1.000
+ // CHECK-DAG: [[C_0:%.+]] = arith.constant 0.000
+ // CHECK-DAG: [[C_1:%.+]] = arith.constant 1.000
+ // CHECK-DAG: [[C_4841369599423283200:%.*]] = arith.constant 4841369599423283200
+ // CHECK-DAG: [[C_9223372036854775807:%.*]] = arith.constant 9223372036854775807
+ // CHECK-NEXT: [[ARG_BITCAST:%.*]] = arith.bitcast [[ARG0]] : f64 to i64
+ // CHECK-NEXT: [[ANDI:%.*]] = arith.andi [[ARG_BITCAST]], [[C_9223372036854775807]]
+ // CHECK-NEXT: [[IS_SPECIAL_VAL:%.*]] = arith.cmpi uge, [[ANDI]], [[C_4841369599423283200]]
// CHECK-NEXT: [[CVTI:%.+]] = arith.fptosi [[ARG0]]
// CHECK-NEXT: [[CVTF:%.+]] = arith.sitofp [[CVTI]]
// CHECK-NEXT: [[COPYSIGN:%.+]] = math.copysign [[CVTF]], [[ARG0]]
// CHECK-NEXT: [[COMP:%.+]] = arith.cmpf ogt, [[ARG0]], [[COPYSIGN]]
- // CHECK-NEXT: [[INCR:%.+]] = arith.select [[COMP]], [[CST_0]], [[CST]]
+ // CHECK-NEXT: [[INCR:%.+]] = arith.select [[COMP]], [[C_1]], [[C_0]]
// CHECK-NEXT: [[ADDF:%.+]] = arith.addf [[COPYSIGN]], [[INCR]]
- // CHECK-NEXT: return [[ADDF]]
+ // CHECK-NEXT: [[RESULT:%.*]] = arith.select [[IS_SPECIAL_VAL]], [[ARG0]], [[ADDF]]
+ // CHECK-NEXT: return [[RESULT]]
// CHECK-FILTER: math.ceil
%ret = math.ceil %a : f64
return %ret : f64
@@ -159,6 +165,34 @@ func.func @ceilf_func(%a: f64) -> f64 {
// -----
+// CHECK-LABEL: func @ceilf_fnuz_func
+// CHECK-SAME: ([[ARG0:%.+]]: f8E5M2FNUZ) -> f8E5M2FNUZ
+func.func @ceilf_fnuz_func(%a: f8E5M2FNUZ) -> f8E5M2FNUZ {
+ // CHECK-DAG: [[C_0:%.+]] = arith.constant 0.000
+ // CHECK-DAG: [[C_1:%.+]] = arith.constant 1.000
+ // CHECK-DAG: [[C_NEG_128:%.*]] = arith.constant -128
+ // CHECK-DAG: [[C_72:%.*]] = arith.constant 72
+ // CHECK-DAG: [[C_127:%.*]] = arith.constant 127
+ // CHECK-NEXT: [[ARG_BITCAST:%.*]] = arith.bitcast [[ARG0]] : f8E5M2FNUZ to i8
+ // CHECK-NEXT: [[ANDI:%.*]] = arith.andi [[ARG_BITCAST]], [[C_127]]
+ // CHECK-NEXT: [[IS_LARGE:%.+]] = arith.cmpi uge, [[ANDI]], [[C_72]]
+ // CHECK-NEXT: [[IS_NAN:%.+]] = arith.cmpi eq, [[ARG_BITCAST]], [[C_NEG_128]]
+ // CHECK-NEXT: [[IS_SPECIAL_VAL:%.+]] = arith.ori [[IS_LARGE]], [[IS_NAN]]
+ // CHECK-NEXT: [[CVTI:%.+]] = arith.fptosi [[ARG0]]
+ // CHECK-NEXT: [[CVTF:%.+]] = arith.sitofp [[CVTI]]
+ // CHECK-NEXT: [[COPYSIGN:%.+]] = math.copysign [[CVTF]], [[ARG0]]
+ // CHECK-NEXT: [[COMP:%.+]] = arith.cmpf ogt, [[ARG0]], [[COPYSIGN]]
+ // CHECK-NEXT: [[INCR:%.+]] = arith.select [[COMP]], [[C_1]], [[C_0]]
+ // CHECK-NEXT: [[ADDF:%.+]] = arith.addf [[COPYSIGN]], [[INCR]]
+ // CHECK-NEXT: [[RESULT:%.*]] = arith.select [[IS_SPECIAL_VAL]], [[ARG0]], [[ADDF]]
+ // CHECK-NEXT: return [[RESULT]]
+ // CHECK-FILTER: math.ceil
+ %ret = math.ceil %a : f8E5M2FNUZ
+ return %ret : f8E5M2FNUZ
+}
+
+// -----
+
// CHECK-LABEL: func @exp2f_func
// CHECK-SAME: ([[ARG0:%.+]]: f64) -> f64
func.func @exp2f_func(%a: f64) -> f64 {
More information about the Mlir-commits
mailing list