[Mlir-commits] [mlir] a43db54 - [mlir][math] Avoid integer division by zero when lowering math.ipowi (#230271)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Oct 19 15:15:11 PDT 2026
Author: Chris Frisz
Date: 2026-10-19T22:15:04Z
New Revision: a43db5458aec686da77142ea42b77df6ba4db283
URL: https://github.com/llvm/llvm-project/commit/a43db5458aec686da77142ea42b77df6ba4db283
DIFF: https://github.com/llvm/llvm-project/commit/a43db5458aec686da77142ea42b77df6ba4db283.diff
LOG: [mlir][math] Avoid integer division by zero when lowering math.ipowi (#230271)
`math.ipowi` with a zero base and a negative exponent is lowered to an
integer division by zero in two places: `MathToFuncs` emits a literal `1
/ 0` for that case, and `PowIStrengthReduction` rewrites a constant
negative exponent into `1 / (b * b * ...)` for a base that may be zero.
`arith.divsi` by zero is undefined behavior, so the result depends on
the target and the optimizer: on AArch64, `0 ** -1` lowered through
`MathToFuncs` returns 0 under `mlir-runner --O0` and 1 under `--O3`.
This patch:
- deletes the `b == 0` branch from the `MathToFuncs` helper, so a zero
base with a negative exponent returns 0 through the same final return as
every `|b| > 1` base;
- stops `PowIStrengthReduction` from reducing `ipowi` on `IntegerType`
with a negative exponent, leaving it to `MathToFuncs`;
- makes `PowIStrengthReduction` bail out for an `INT64_MIN` exponent,
since negating it is signed overflow (this changes `fpowi`,
`complex.powi` and index `ipowi` for that one exponent).
Why 0: `0 ** -n` has no representable value, so any result is a choice.
0 is deterministic, in range for every integer type, needs no new
semantics, and matches the `b == 1` and `b == -1` cases, which already
return a value. Whether the dialect should instead declare this UB is
still open in the Discourse thread "Undefined behavior in the math
dialect"
(https://discourse.llvm.org/t/undefined-behavior-in-the-math-dialect/90522).
This patch only removes lowerings whose answer depends on the optimizer.
Related: PR 188783 takes a different approach. It makes
`PowIStrengthReduction` skip every negative-exponent `ipowi`, but leaves
the `1 / 0` in `MathToFuncs` and the `INT64_MIN` negation as they are.
This patch defines the zero-base result in `MathToFuncs` itself and also
fixes the negation.
Caveat: `MathToFuncs` does not lower `index`, so index `ipowi` keeps the
old reduction, and an index `0 ** -n` still divides by zero.
Assisted-by: Claude (Anthropic)
Co-authored-by: Slava Zakharin <szakharin at nvidia.com>
Added:
mlir/test/Integration/Dialect/Math/CPU/mathtofuncs_ipowi.mlir
Modified:
mlir/lib/Conversion/MathToFuncs/MathToFuncs.cpp
mlir/lib/Dialect/Math/Transforms/AlgebraicSimplification.cpp
mlir/test/Conversion/MathToFuncs/ipowi.mlir
mlir/test/Dialect/Complex/powi-simplify.mlir
mlir/test/Dialect/Math/algebraic-simplification.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Conversion/MathToFuncs/MathToFuncs.cpp b/mlir/lib/Conversion/MathToFuncs/MathToFuncs.cpp
index f653c6b38e002..ead1b56c135de 100644
--- a/mlir/lib/Conversion/MathToFuncs/MathToFuncs.cpp
+++ b/mlir/lib/Conversion/MathToFuncs/MathToFuncs.cpp
@@ -160,8 +160,6 @@ static FunctionType getElementalFuncTypeForOp(Operation *op) {
/// if (p == T(0))
/// return T(1);
/// if (p < T(0)) {
-/// if (b == T(0))
-/// return T(1) / T(0); // trigger div-by-zero
/// if (b == T(1))
/// return T(1);
/// if (b == T(-1)) {
@@ -169,7 +167,7 @@ static FunctionType getElementalFuncTypeForOp(Operation *op) {
/// return T(-1);
/// return T(1);
/// }
-/// return T(0);
+/// return T(0); // |b| > 1 or b == 0
/// }
/// T result = T(1);
/// while (true) {
@@ -232,22 +230,8 @@ static func::FuncOp createElementIPowIFunc(ModuleOp *module, Type elementType) {
builder.setInsertionPointToEnd(fallthroughBlock);
auto pIsNeg = arith::CmpIOp::create(builder, arith::CmpIPredicate::sle, pArg,
zeroValue);
- // if (b == T(0))
- builder.createBlock(funcBody);
- auto bIsZero =
- arith::CmpIOp::create(builder, arith::CmpIPredicate::eq, bArg, zeroValue);
- // return T(1) / T(0);
- thenBlock = builder.createBlock(funcBody);
- func::ReturnOp::create(
- builder,
- arith::DivSIOp::create(builder, oneValue, zeroValue).getResult());
- fallthroughBlock = builder.createBlock(funcBody);
- // Set up conditional branch for (b == T(0)).
- builder.setInsertionPointToEnd(bIsZero->getBlock());
- cf::CondBranchOp::create(builder, bIsZero, thenBlock, fallthroughBlock);
-
// if (b == T(1))
- builder.setInsertionPointToEnd(fallthroughBlock);
+ builder.createBlock(funcBody);
auto bIsOne =
arith::CmpIOp::create(builder, arith::CmpIPredicate::eq, bArg, oneValue);
// return T(1);
@@ -295,7 +279,7 @@ static func::FuncOp createElementIPowIFunc(ModuleOp *module, Type elementType) {
// Set up conditional branch for (p < T(0)).
builder.setInsertionPointToEnd(pIsNeg->getBlock());
// Set initial values of 'result', 'b' and 'p' for the loop.
- cf::CondBranchOp::create(builder, pIsNeg, bIsZero->getBlock(), loopHeader,
+ cf::CondBranchOp::create(builder, pIsNeg, bIsOne->getBlock(), loopHeader,
ValueRange{oneValue, bArg, pArg});
// T result = T(1);
diff --git a/mlir/lib/Dialect/Math/Transforms/AlgebraicSimplification.cpp b/mlir/lib/Dialect/Math/Transforms/AlgebraicSimplification.cpp
index 92be8fe40c163..eecb40b2689b1 100644
--- a/mlir/lib/Dialect/Math/Transforms/AlgebraicSimplification.cpp
+++ b/mlir/lib/Dialect/Math/Transforms/AlgebraicSimplification.cpp
@@ -21,6 +21,7 @@
#include "mlir/IR/Matchers.h"
#include "mlir/IR/TypeUtilities.h"
#include <climits>
+#include <limits>
using namespace mlir;
@@ -171,6 +172,20 @@ PowIStrengthReduction<PowIOpTy, DivOpTy, MulOpTy>::matchAndRewrite(
bool exponentIsNegative = false;
if (exponentValue < 0) {
exponentIsNegative = true;
+ // Integer division by zero is UB (unlike float/complex, where the
+ // division is defined and yields Inf). base isn't necessarily constant
+ // here, so defer negative-exponent IPowIOp to MathToFuncs's
+ // createElementIPowIFunc, which returns 0 for base == 0 without
+ // dividing. MathToFuncs does not lower index, so index keeps the
+ // strength reduction below.
+ if constexpr (std::is_same_v<PowIOpTy, math::IPowIOp>)
+ if (isa<IntegerType>(getElementTypeOrSelf(op.getType())))
+ return failure();
+ // INT64_MIN has no positive representation in int64_t, so negating it
+ // below would overflow. Bail out rather than negate. This is reached by
+ // FPowIOp, complex::PowiOp and index-typed IPowIOp.
+ if (exponentValue == std::numeric_limits<int64_t>::min())
+ return failure();
exponentValue *= -1;
}
diff --git a/mlir/test/Conversion/MathToFuncs/ipowi.mlir b/mlir/test/Conversion/MathToFuncs/ipowi.mlir
index 2702a1e22e621..d3c71b8fc6d5c 100644
--- a/mlir/test/Conversion/MathToFuncs/ipowi.mlir
+++ b/mlir/test/Conversion/MathToFuncs/ipowi.mlir
@@ -1,4 +1,5 @@
// RUN: mlir-opt %s -split-input-file -pass-pipeline="builtin.module(convert-math-to-funcs)" | FileCheck %s
+// RUN: mlir-opt %s -split-input-file -pass-pipeline="builtin.module(convert-math-to-funcs)" | FileCheck %s --check-prefix=NODIV
// -----
@@ -24,47 +25,41 @@ func.func @ipowi(%arg0: i64, %arg1: i64) {
// CHECK: return %[[VAL_3]] : i64
// CHECK: ^bb2:
// CHECK: %[[VAL_6:.*]] = arith.cmpi sle, %[[VAL_1]], %[[VAL_2]] : i64
-// CHECK: cf.cond_br %[[VAL_6]], ^bb3, ^bb12(%[[VAL_3]], %[[VAL_0]], %[[VAL_1]] : i64, i64, i64)
+// CHECK: cf.cond_br %[[VAL_6]], ^bb3, ^bb10(%[[VAL_3]], %[[VAL_0]], %[[VAL_1]] : i64, i64, i64)
// CHECK: ^bb3:
-// CHECK: %[[VAL_7:.*]] = arith.cmpi eq, %[[VAL_0]], %[[VAL_2]] : i64
+// CHECK: %[[VAL_7:.*]] = arith.cmpi eq, %[[VAL_0]], %[[VAL_3]] : i64
// CHECK: cf.cond_br %[[VAL_7]], ^bb4, ^bb5
// CHECK: ^bb4:
-// CHECK: %[[VAL_8:.*]] = arith.divsi %[[VAL_3]], %[[VAL_2]] : i64
-// CHECK: return %[[VAL_8]] : i64
+// CHECK: return %[[VAL_3]] : i64
// CHECK: ^bb5:
-// CHECK: %[[VAL_9:.*]] = arith.cmpi eq, %[[VAL_0]], %[[VAL_3]] : i64
-// CHECK: cf.cond_br %[[VAL_9]], ^bb6, ^bb7
+// CHECK: %[[VAL_8:.*]] = arith.cmpi eq, %[[VAL_0]], %[[VAL_4]] : i64
+// CHECK: cf.cond_br %[[VAL_8]], ^bb6, ^bb9
// CHECK: ^bb6:
-// CHECK: return %[[VAL_3]] : i64
+// CHECK: %[[VAL_9:.*]] = arith.andi %[[VAL_1]], %[[VAL_3]] : i64
+// CHECK: %[[VAL_10:.*]] = arith.cmpi ne, %[[VAL_9]], %[[VAL_2]] : i64
+// CHECK: cf.cond_br %[[VAL_10]], ^bb7, ^bb8
// CHECK: ^bb7:
-// CHECK: %[[VAL_10:.*]] = arith.cmpi eq, %[[VAL_0]], %[[VAL_4]] : i64
-// CHECK: cf.cond_br %[[VAL_10]], ^bb8, ^bb11
-// CHECK: ^bb8:
-// CHECK: %[[VAL_11:.*]] = arith.andi %[[VAL_1]], %[[VAL_3]] : i64
-// CHECK: %[[VAL_12:.*]] = arith.cmpi ne, %[[VAL_11]], %[[VAL_2]] : i64
-// CHECK: cf.cond_br %[[VAL_12]], ^bb9, ^bb10
-// CHECK: ^bb9:
// CHECK: return %[[VAL_4]] : i64
-// CHECK: ^bb10:
+// CHECK: ^bb8:
// CHECK: return %[[VAL_3]] : i64
-// CHECK: ^bb11:
+// CHECK: ^bb9:
// CHECK: return %[[VAL_2]] : i64
-// CHECK: ^bb12(%[[VAL_13:.*]]: i64, %[[VAL_14:.*]]: i64, %[[VAL_15:.*]]: i64):
-// CHECK: %[[VAL_16:.*]] = arith.andi %[[VAL_15]], %[[VAL_3]] : i64
-// CHECK: %[[VAL_17:.*]] = arith.cmpi ne, %[[VAL_16]], %[[VAL_2]] : i64
-// CHECK: cf.cond_br %[[VAL_17]], ^bb13, ^bb14(%[[VAL_13]] : i64)
+// CHECK: ^bb10(%[[VAL_11:.*]]: i64, %[[VAL_12:.*]]: i64, %[[VAL_13:.*]]: i64):
+// CHECK: %[[VAL_14:.*]] = arith.andi %[[VAL_13]], %[[VAL_3]] : i64
+// CHECK: %[[VAL_15:.*]] = arith.cmpi ne, %[[VAL_14]], %[[VAL_2]] : i64
+// CHECK: cf.cond_br %[[VAL_15]], ^bb11, ^bb12(%[[VAL_11]] : i64)
+// CHECK: ^bb11:
+// CHECK: %[[VAL_16:.*]] = arith.muli %[[VAL_11]], %[[VAL_12]] : i64
+// CHECK: cf.br ^bb12(%[[VAL_16]] : i64)
+// CHECK: ^bb12(%[[VAL_17:.*]]: i64):
+// CHECK: %[[VAL_18:.*]] = arith.shrui %[[VAL_13]], %[[VAL_3]] : i64
+// CHECK: %[[VAL_19:.*]] = arith.cmpi eq, %[[VAL_18]], %[[VAL_2]] : i64
+// CHECK: cf.cond_br %[[VAL_19]], ^bb13, ^bb14
// CHECK: ^bb13:
-// CHECK: %[[VAL_18:.*]] = arith.muli %[[VAL_13]], %[[VAL_14]] : i64
-// CHECK: cf.br ^bb14(%[[VAL_18]] : i64)
-// CHECK: ^bb14(%[[VAL_19:.*]]: i64):
-// CHECK: %[[VAL_20:.*]] = arith.shrui %[[VAL_15]], %[[VAL_3]] : i64
-// CHECK: %[[VAL_21:.*]] = arith.cmpi eq, %[[VAL_20]], %[[VAL_2]] : i64
-// CHECK: cf.cond_br %[[VAL_21]], ^bb15, ^bb16
-// CHECK: ^bb15:
-// CHECK: return %[[VAL_19]] : i64
-// CHECK: ^bb16:
-// CHECK: %[[VAL_22:.*]] = arith.muli %[[VAL_14]], %[[VAL_14]] : i64
-// CHECK: cf.br ^bb12(%[[VAL_19]], %[[VAL_22]], %[[VAL_20]] : i64, i64, i64)
+// CHECK: return %[[VAL_17]] : i64
+// CHECK: ^bb14:
+// CHECK: %[[VAL_20:.*]] = arith.muli %[[VAL_12]], %[[VAL_12]] : i64
+// CHECK: cf.br ^bb10(%[[VAL_17]], %[[VAL_20]], %[[VAL_18]] : i64, i64, i64)
// CHECK: }
// -----
@@ -91,47 +86,41 @@ func.func @ipowi(%arg0: i8, %arg1: i8) {
// CHECK: return %[[VAL_3]] : i8
// CHECK: ^bb2:
// CHECK: %[[VAL_6:.*]] = arith.cmpi sle, %[[VAL_1]], %[[VAL_2]] : i8
-// CHECK: cf.cond_br %[[VAL_6]], ^bb3, ^bb12(%[[VAL_3]], %[[VAL_0]], %[[VAL_1]] : i8, i8, i8)
+// CHECK: cf.cond_br %[[VAL_6]], ^bb3, ^bb10(%[[VAL_3]], %[[VAL_0]], %[[VAL_1]] : i8, i8, i8)
// CHECK: ^bb3:
-// CHECK: %[[VAL_7:.*]] = arith.cmpi eq, %[[VAL_0]], %[[VAL_2]] : i8
+// CHECK: %[[VAL_7:.*]] = arith.cmpi eq, %[[VAL_0]], %[[VAL_3]] : i8
// CHECK: cf.cond_br %[[VAL_7]], ^bb4, ^bb5
// CHECK: ^bb4:
-// CHECK: %[[VAL_8:.*]] = arith.divsi %[[VAL_3]], %[[VAL_2]] : i8
-// CHECK: return %[[VAL_8]] : i8
+// CHECK: return %[[VAL_3]] : i8
// CHECK: ^bb5:
-// CHECK: %[[VAL_9:.*]] = arith.cmpi eq, %[[VAL_0]], %[[VAL_3]] : i8
-// CHECK: cf.cond_br %[[VAL_9]], ^bb6, ^bb7
+// CHECK: %[[VAL_8:.*]] = arith.cmpi eq, %[[VAL_0]], %[[VAL_4]] : i8
+// CHECK: cf.cond_br %[[VAL_8]], ^bb6, ^bb9
// CHECK: ^bb6:
-// CHECK: return %[[VAL_3]] : i8
+// CHECK: %[[VAL_9:.*]] = arith.andi %[[VAL_1]], %[[VAL_3]] : i8
+// CHECK: %[[VAL_10:.*]] = arith.cmpi ne, %[[VAL_9]], %[[VAL_2]] : i8
+// CHECK: cf.cond_br %[[VAL_10]], ^bb7, ^bb8
// CHECK: ^bb7:
-// CHECK: %[[VAL_10:.*]] = arith.cmpi eq, %[[VAL_0]], %[[VAL_4]] : i8
-// CHECK: cf.cond_br %[[VAL_10]], ^bb8, ^bb11
-// CHECK: ^bb8:
-// CHECK: %[[VAL_11:.*]] = arith.andi %[[VAL_1]], %[[VAL_3]] : i8
-// CHECK: %[[VAL_12:.*]] = arith.cmpi ne, %[[VAL_11]], %[[VAL_2]] : i8
-// CHECK: cf.cond_br %[[VAL_12]], ^bb9, ^bb10
-// CHECK: ^bb9:
// CHECK: return %[[VAL_4]] : i8
-// CHECK: ^bb10:
+// CHECK: ^bb8:
// CHECK: return %[[VAL_3]] : i8
-// CHECK: ^bb11:
+// CHECK: ^bb9:
// CHECK: return %[[VAL_2]] : i8
-// CHECK: ^bb12(%[[VAL_13:.*]]: i8, %[[VAL_14:.*]]: i8, %[[VAL_15:.*]]: i8):
-// CHECK: %[[VAL_16:.*]] = arith.andi %[[VAL_15]], %[[VAL_3]] : i8
-// CHECK: %[[VAL_17:.*]] = arith.cmpi ne, %[[VAL_16]], %[[VAL_2]] : i8
-// CHECK: cf.cond_br %[[VAL_17]], ^bb13, ^bb14(%[[VAL_13]] : i8)
+// CHECK: ^bb10(%[[VAL_11:.*]]: i8, %[[VAL_12:.*]]: i8, %[[VAL_13:.*]]: i8):
+// CHECK: %[[VAL_14:.*]] = arith.andi %[[VAL_13]], %[[VAL_3]] : i8
+// CHECK: %[[VAL_15:.*]] = arith.cmpi ne, %[[VAL_14]], %[[VAL_2]] : i8
+// CHECK: cf.cond_br %[[VAL_15]], ^bb11, ^bb12(%[[VAL_11]] : i8)
+// CHECK: ^bb11:
+// CHECK: %[[VAL_16:.*]] = arith.muli %[[VAL_11]], %[[VAL_12]] : i8
+// CHECK: cf.br ^bb12(%[[VAL_16]] : i8)
+// CHECK: ^bb12(%[[VAL_17:.*]]: i8):
+// CHECK: %[[VAL_18:.*]] = arith.shrui %[[VAL_13]], %[[VAL_3]] : i8
+// CHECK: %[[VAL_19:.*]] = arith.cmpi eq, %[[VAL_18]], %[[VAL_2]] : i8
+// CHECK: cf.cond_br %[[VAL_19]], ^bb13, ^bb14
// CHECK: ^bb13:
-// CHECK: %[[VAL_18:.*]] = arith.muli %[[VAL_13]], %[[VAL_14]] : i8
-// CHECK: cf.br ^bb14(%[[VAL_18]] : i8)
-// CHECK: ^bb14(%[[VAL_19:.*]]: i8):
-// CHECK: %[[VAL_20:.*]] = arith.shrui %[[VAL_15]], %[[VAL_3]] : i8
-// CHECK: %[[VAL_21:.*]] = arith.cmpi eq, %[[VAL_20]], %[[VAL_2]] : i8
-// CHECK: cf.cond_br %[[VAL_21]], ^bb15, ^bb16
-// CHECK: ^bb15:
-// CHECK: return %[[VAL_19]] : i8
-// CHECK: ^bb16:
-// CHECK: %[[VAL_22:.*]] = arith.muli %[[VAL_14]], %[[VAL_14]] : i8
-// CHECK: cf.br ^bb12(%[[VAL_19]], %[[VAL_22]], %[[VAL_20]] : i8, i8, i8)
+// CHECK: return %[[VAL_17]] : i8
+// CHECK: ^bb14:
+// CHECK: %[[VAL_20:.*]] = arith.muli %[[VAL_12]], %[[VAL_12]] : i8
+// CHECK: cf.br ^bb10(%[[VAL_17]], %[[VAL_20]], %[[VAL_18]] : i8, i8, i8)
// CHECK: }
// -----
@@ -181,3 +170,40 @@ func.func @ipowi_index(%arg0: index, %arg1: index) {
%0 = math.ipowi %arg0, %arg1 : index
func.return
}
+
+// -----
+
+// Zero base raised to a negative exponent is zero. The expansion gets there
+// through the same final return of the zero constant as every |b| > 1 base,
+// and must not get there by dividing by zero: that is undefined behaviour, a
+// trap on targets whose integer division faults and a wrong value on the
+// ones where the optimizer exploits it.
+
+// CHECK-LABEL: func @ipowi_zero_base_negative_exponent(
+// CHECK-SAME: %[[ARG0:.+]]: i32)
+func.func @ipowi_zero_base_negative_exponent(%arg0: i32) -> i32 {
+ %c0_i32 = arith.constant 0 : i32
+ // CHECK: call @__mlir_math_ipowi_i32(%{{.*}}, %[[ARG0]]) : (i32, i32) -> i32
+ %0 = math.ipowi %c0_i32, %arg0 : i32
+ func.return %0 : i32
+}
+
+// CHECK-LABEL: func.func private @__mlir_math_ipowi_i32(
+// CHECK-SAME: %[[VAL_0:.*]]: i32,
+// CHECK-SAME: %[[VAL_1:.*]]: i32) -> i32
+// CHECK: %[[VAL_2:.*]] = arith.constant 0 : i32
+// CHECK: %[[VAL_3:.*]] = arith.constant 1 : i32
+// CHECK: %[[VAL_4:.*]] = arith.constant -1 : i32
+// CHECK: %[[VAL_5:.*]] = arith.cmpi eq, %[[VAL_1]], %[[VAL_2]] : i32
+// CHECK: cf.cond_br %[[VAL_5]], ^bb1, ^bb2
+// CHECK: ^bb2:
+// CHECK: %[[VAL_6:.*]] = arith.cmpi sle, %[[VAL_1]], %[[VAL_2]] : i32
+// CHECK: cf.cond_br %[[VAL_6]], ^bb3, ^bb10
+
+// No block of the helper may divide. The end anchor matches the closing brace
+// of the module, which bounds the NODIV-NOT to the whole helper (a bare `}`
+// would match the brace of the attribute dictionary on the label line).
+
+// NODIV-LABEL: func.func private @__mlir_math_ipowi_i32
+// NODIV-NOT: arith.divsi
+// NODIV: {{^\}$}}
diff --git a/mlir/test/Dialect/Complex/powi-simplify.mlir b/mlir/test/Dialect/Complex/powi-simplify.mlir
index c7bb6a9d81479..9f998d385310b 100644
--- a/mlir/test/Dialect/Complex/powi-simplify.mlir
+++ b/mlir/test/Dialect/Complex/powi-simplify.mlir
@@ -18,3 +18,13 @@ func.func @pow9(%arg0: complex<f32>) -> complex<f32> {
}
// CHECK-LABEL: func.func @pow9(
// CHECK: complex.powi %{{.*}}, %{{.*}} : complex<f32>, i32
+
+// INT64_MIN has no positive int64_t representation; PowIStrengthReduction
+// must bail out rather than negate it (negating would overflow).
+func.func @powi_int64_min(%arg0: complex<f32>) -> complex<f32> {
+ %cmin = arith.constant -9223372036854775808 : i64
+ %0 = complex.powi %arg0, %cmin : complex<f32>, i64
+ return %0 : complex<f32>
+}
+// CHECK-LABEL: func.func @powi_int64_min(
+// CHECK: complex.powi %{{.*}}, %{{.*}} : complex<f32>, i64
diff --git a/mlir/test/Dialect/Math/algebraic-simplification.mlir b/mlir/test/Dialect/Math/algebraic-simplification.mlir
index 897147ee35441..e311fb7d6a2e8 100644
--- a/mlir/test/Dialect/Math/algebraic-simplification.mlir
+++ b/mlir/test/Dialect/Math/algebraic-simplification.mlir
@@ -190,10 +190,16 @@ func.func @ipowi_zero_exp(%arg0: i32, %arg1: vector<4xi32>) -> (i32, vector<4xi3
// CHECK-SAME: %[[ARG1:.+]]: vector<4xi32>
// CHECK-SAME: -> (i32, vector<4xi32>, i32, vector<4xi32>) {
func.func @ipowi_exp_one(%arg0: i32, %arg1: vector<4xi32>) -> (i32, vector<4xi32>, i32, vector<4xi32>) {
- // CHECK-DAG: %[[CST_S:.*]] = arith.constant 1 : i32
- // CHECK-DAG: %[[CST_V:.*]] = arith.constant dense<1> : vector<4xi32>
- // CHECK: %[[SCALAR:.*]] = arith.divsi %[[CST_S]], %[[ARG0]]
- // CHECK: %[[VECTOR:.*]] = arith.divsi %[[CST_V]], %[[ARG1]]
+ // Positive exponent (1) folds to the base directly. Negative exponent
+ // (-1) is intentionally left unconverted: integer negative-exponent
+ // IPowIOp is lowered exactly once, correctly, by ConvertMathToFuncs's
+ // createElementIPowIFunc (base == 0 returns 0 there, see MathToFuncs.cpp
+ // and MathToFuncs/ipowi.mlir). Reducing it here too would divide by zero
+ // for that base.
+ // CHECK-DAG: %[[CM1:.*]] = arith.constant -1 : i32
+ // CHECK-DAG: %[[VM1:.*]] = arith.constant dense<-1> : vector<4xi32>
+ // CHECK: %[[SCALAR:.*]] = math.ipowi %[[ARG0]], %[[CM1]]
+ // CHECK: %[[VECTOR:.*]] = math.ipowi %[[ARG1]], %[[VM1]]
// CHECK: return %[[ARG0]], %[[ARG1]], %[[SCALAR]], %[[VECTOR]]
%c1 = arith.constant 1 : i32
%v1 = arith.constant dense <1> : vector<4xi32>
@@ -211,14 +217,14 @@ func.func @ipowi_exp_one(%arg0: i32, %arg1: vector<4xi32>) -> (i32, vector<4xi32
// CHECK-SAME: %[[ARG1:.+]]: vector<4xi32>
// CHECK-SAME: -> (i32, vector<4xi32>, i32, vector<4xi32>) {
func.func @ipowi_exp_two(%arg0: i32, %arg1: vector<4xi32>) -> (i32, vector<4xi32>, i32, vector<4xi32>) {
- // CHECK-DAG: %[[CST_S:.*]] = arith.constant 1 : i32
- // CHECK-DAG: %[[CST_V:.*]] = arith.constant dense<1> : vector<4xi32>
+ // Positive exponent (2) still strength-reduces to a multiplication.
+ // Negative exponent (-2) is left unconverted (see @ipowi_exp_one).
+ // CHECK-DAG: %[[CM2:.*]] = arith.constant -2 : i32
+ // CHECK-DAG: %[[VM2:.*]] = arith.constant dense<-2> : vector<4xi32>
// CHECK: %[[SCALAR0:.*]] = arith.muli %[[ARG0]], %[[ARG0]]
// CHECK: %[[VECTOR0:.*]] = arith.muli %[[ARG1]], %[[ARG1]]
- // CHECK: %[[SMUL:.*]] = arith.muli %[[ARG0]], %[[ARG0]]
- // CHECK: %[[SCALAR1:.*]] = arith.divsi %[[CST_S]], %[[SMUL]]
- // CHECK: %[[VMUL:.*]] = arith.muli %[[ARG1]], %[[ARG1]]
- // CHECK: %[[VECTOR1:.*]] = arith.divsi %[[CST_V]], %[[VMUL]]
+ // CHECK: %[[SCALAR1:.*]] = math.ipowi %[[ARG0]], %[[CM2]]
+ // CHECK: %[[VECTOR1:.*]] = math.ipowi %[[ARG1]], %[[VM2]]
// CHECK: return %[[SCALAR0]], %[[VECTOR0]], %[[SCALAR1]], %[[VECTOR1]]
%c1 = arith.constant 2 : i32
%v1 = arith.constant dense <2> : vector<4xi32>
@@ -236,18 +242,16 @@ func.func @ipowi_exp_two(%arg0: i32, %arg1: vector<4xi32>) -> (i32, vector<4xi32
// CHECK-SAME: %[[ARG1:.+]]: vector<4xi32>
// CHECK-SAME: -> (i32, vector<4xi32>, i32, vector<4xi32>) {
func.func @ipowi_exp_three(%arg0: i32, %arg1: vector<4xi32>) -> (i32, vector<4xi32>, i32, vector<4xi32>) {
- // CHECK-DAG: %[[CST_S:.*]] = arith.constant 1 : i32
- // CHECK-DAG: %[[CST_V:.*]] = arith.constant dense<1> : vector<4xi32>
+ // Positive exponent (3) still strength-reduces to multiplications.
+ // Negative exponent (-3) is left unconverted (see @ipowi_exp_one).
+ // CHECK-DAG: %[[CM3:.*]] = arith.constant -3 : i32
+ // CHECK-DAG: %[[VM3:.*]] = arith.constant dense<-3> : vector<4xi32>
// CHECK: %[[SMUL0:.*]] = arith.muli %[[ARG0]], %[[ARG0]]
// CHECK: %[[SCALAR0:.*]] = arith.muli %[[SMUL0]], %[[ARG0]]
// CHECK: %[[VMUL0:.*]] = arith.muli %[[ARG1]], %[[ARG1]]
// CHECK: %[[VECTOR0:.*]] = arith.muli %[[VMUL0]], %[[ARG1]]
- // CHECK: %[[SMUL1:.*]] = arith.muli %[[ARG0]], %[[ARG0]]
- // CHECK: %[[SMUL2:.*]] = arith.muli %[[SMUL1]], %[[ARG0]]
- // CHECK: %[[SCALAR1:.*]] = arith.divsi %[[CST_S]], %[[SMUL2]]
- // CHECK: %[[VMUL1:.*]] = arith.muli %[[ARG1]], %[[ARG1]]
- // CHECK: %[[VMUL2:.*]] = arith.muli %[[VMUL1]], %[[ARG1]]
- // CHECK: %[[VECTOR1:.*]] = arith.divsi %[[CST_V]], %[[VMUL2]]
+ // CHECK: %[[SCALAR1:.*]] = math.ipowi %[[ARG0]], %[[CM3]]
+ // CHECK: %[[VECTOR1:.*]] = math.ipowi %[[ARG1]], %[[VM3]]
// CHECK: return %[[SCALAR0]], %[[VECTOR0]], %[[SCALAR1]], %[[VECTOR1]]
%c1 = arith.constant 3 : i32
%v1 = arith.constant dense <3> : vector<4xi32>
@@ -260,6 +264,34 @@ func.func @ipowi_exp_three(%arg0: i32, %arg1: vector<4xi32>) -> (i32, vector<4xi
return %0, %1, %2, %3 : i32, vector<4xi32>, i32, vector<4xi32>
}
+// MathToFuncs does not lower math.ipowi on index, so a negative constant
+// exponent is still strength-reduced here instead of being left unlowered.
+// CHECK-LABEL: @ipowi_index_neg_exp(
+// CHECK-SAME: %[[ARG0:.+]]: index
+// CHECK-DAG: %[[ONE:.*]] = arith.constant 1 : index
+// CHECK: %[[SQ:.*]] = arith.muli %[[ARG0]], %[[ARG0]] : index
+// CHECK: %[[RES:.*]] = arith.divsi %[[ONE]], %[[SQ]] : index
+// CHECK: return %[[RES]]
+func.func @ipowi_index_neg_exp(%arg0: index) -> index {
+ %cm2 = arith.constant -2 : index
+ %0 = math.ipowi %arg0, %cm2 : index
+ return %0 : index
+}
+
+// INT64_MIN has no positive int64_t representation; PowIStrengthReduction
+// must bail out rather than negate it (negating would overflow). Index ipowi
+// is not deferred to MathToFuncs, so it reaches that check.
+// CHECK-LABEL: @ipowi_index_exp_int64_min(
+// CHECK-SAME: %[[ARG0:.+]]: index
+// CHECK: %[[CST:.*]] = arith.constant -9223372036854775808 : index
+// CHECK: %[[RES:.*]] = math.ipowi %[[ARG0]], %[[CST]] : index
+// CHECK: return %[[RES]]
+func.func @ipowi_index_exp_int64_min(%arg0: index) -> index {
+ %cmin = arith.constant -9223372036854775808 : index
+ %0 = math.ipowi %arg0, %cmin : index
+ return %0 : index
+}
+
// CHECK-LABEL: @fpowi_zero_exp(
// CHECK-SAME: %[[ARG0:.+]]: f32
// CHECK-SAME: %[[ARG1:.+]]: vector<4xf32>
@@ -296,6 +328,19 @@ func.func @fpowi_exp_one(%arg0: f32, %arg1: vector<4xf32>) -> (f32, vector<4xf32
return %0, %1, %2, %3 : f32, vector<4xf32>, f32, vector<4xf32>
}
+// INT64_MIN has no positive int64_t representation; PowIStrengthReduction
+// must bail out rather than negate it (negating would overflow).
+// CHECK-LABEL: @fpowi_exp_int64_min(
+// CHECK-SAME: %[[ARG0:.+]]: f32
+// CHECK: %[[CST:.*]] = arith.constant -9223372036854775808 : i64
+// CHECK: %[[RES:.*]] = math.fpowi %[[ARG0]], %[[CST]] : f32, i64
+// CHECK: return %[[RES]]
+func.func @fpowi_exp_int64_min(%arg0: f32) -> f32 {
+ %cmin = arith.constant -9223372036854775808 : i64
+ %0 = math.fpowi %arg0, %cmin : f32, i64
+ return %0 : f32
+}
+
// CHECK-LABEL: @fpowi_exp_two(
// CHECK-SAME: %[[ARG0:.+]]: f32
// CHECK-SAME: %[[ARG1:.+]]: vector<4xf32>
diff --git a/mlir/test/Integration/Dialect/Math/CPU/mathtofuncs_ipowi.mlir b/mlir/test/Integration/Dialect/Math/CPU/mathtofuncs_ipowi.mlir
new file mode 100644
index 0000000000000..48e4dd11f3491
--- /dev/null
+++ b/mlir/test/Integration/Dialect/Math/CPU/mathtofuncs_ipowi.mlir
@@ -0,0 +1,74 @@
+// RUN: mlir-opt %s \
+// RUN: -pass-pipeline="builtin.module( \
+// RUN: convert-math-to-funcs, \
+// RUN: func.func(convert-scf-to-cf,convert-arith-to-llvm), \
+// RUN: convert-func-to-llvm, \
+// RUN: convert-cf-to-llvm, \
+// RUN: reconcile-unrealized-casts)" \
+// RUN: | mlir-runner --O3 -e test_zero_base_negative_exponent_i32 -entry-point-result=i32 | FileCheck %s --check-prefix=CHECK_TEST_ZERO_BASE_NEG_EXP_I32
+
+// 0 ** -1 is 0. The expansion must not divide by zero to get there: on targets
+// where integer division by zero traps, that is a SIGFPE instead of a result,
+// and in any case it is UB that the optimizer is free to turn into the wrong
+// value (hence --O3 on the RUN lines, so that this test exercises it).
+func.func @test_zero_base_negative_exponent_i32() -> i32 {
+ %base = arith.constant 0 : i32
+ %exp = arith.constant -1 : i32
+ %0 = math.ipowi %base, %exp : i32
+ func.return %0 : i32
+}
+// CHECK_TEST_ZERO_BASE_NEG_EXP_I32: 0
+
+// RUN: mlir-opt %s \
+// RUN: -pass-pipeline="builtin.module( \
+// RUN: convert-math-to-funcs, \
+// RUN: func.func(convert-scf-to-cf,convert-arith-to-llvm), \
+// RUN: convert-func-to-llvm, \
+// RUN: convert-cf-to-llvm, \
+// RUN: reconcile-unrealized-casts)" \
+// RUN: | mlir-runner --O3 -e test_zero_base_negative_even_exponent_i64 -entry-point-result=i64 | FileCheck %s --check-prefix=CHECK_TEST_ZERO_BASE_NEG_EVEN_EXP_I64
+
+// 0 ** -2 is 0 as well; the parity of the exponent must not matter.
+func.func @test_zero_base_negative_even_exponent_i64() -> i64 {
+ %base = arith.constant 0 : i64
+ %exp = arith.constant -2 : i64
+ %0 = math.ipowi %base, %exp : i64
+ func.return %0 : i64
+}
+// CHECK_TEST_ZERO_BASE_NEG_EVEN_EXP_I64: 0
+
+// RUN: mlir-opt %s \
+// RUN: -pass-pipeline="builtin.module( \
+// RUN: convert-math-to-funcs, \
+// RUN: func.func(convert-scf-to-cf,convert-arith-to-llvm), \
+// RUN: convert-func-to-llvm, \
+// RUN: convert-cf-to-llvm, \
+// RUN: reconcile-unrealized-casts)" \
+// RUN: | mlir-runner --O3 -e test_minus_one_base_negative_odd_exponent_i32 -entry-point-result=i32 | FileCheck %s --check-prefix=CHECK_TEST_MINUS_ONE_BASE_NEG_ODD_EXP_I32
+
+// The neighbouring negative-exponent cases keep working: (-1) ** -3 is -1.
+func.func @test_minus_one_base_negative_odd_exponent_i32() -> i32 {
+ %base = arith.constant -1 : i32
+ %exp = arith.constant -3 : i32
+ %0 = math.ipowi %base, %exp : i32
+ func.return %0 : i32
+}
+// CHECK_TEST_MINUS_ONE_BASE_NEG_ODD_EXP_I32: -1
+
+// RUN: mlir-opt %s \
+// RUN: -pass-pipeline="builtin.module( \
+// RUN: convert-math-to-funcs, \
+// RUN: func.func(convert-scf-to-cf,convert-arith-to-llvm), \
+// RUN: convert-func-to-llvm, \
+// RUN: convert-cf-to-llvm, \
+// RUN: reconcile-unrealized-casts)" \
+// RUN: | mlir-runner --O3 -e test_positive_exponent_i32 -entry-point-result=i32 | FileCheck %s --check-prefix=CHECK_TEST_POSITIVE_EXP_I32
+
+// Sanity check that the common path still computes: 2 ** 10 is 1024.
+func.func @test_positive_exponent_i32() -> i32 {
+ %base = arith.constant 2 : i32
+ %exp = arith.constant 10 : i32
+ %0 = math.ipowi %base, %exp : i32
+ func.return %0 : i32
+}
+// CHECK_TEST_POSITIVE_EXP_I32: 1024
More information about the Mlir-commits
mailing list