[Mlir-commits] [mlir] [mlir][Math] Handle wide exponents in powi strength reduction (PR #237513)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sat Oct 31 00:39:27 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-math
Author: woruyu (woruyu)
<details>
<summary>Changes</summary>
This fixes https://github.com/llvm/llvm-project/issues/237494. PowIStrengthReduction converts constant exponents to int64_t before checking the reduction threshold, triggering assertions for out-of-range values. Match scalar and splat constants as APInt and check their signed range before conversion, preserving simplification of small values in wide types.
Assisted-by: OpenAI Codex
---
Full diff: https://github.com/llvm/llvm-project/pull/237513.diff
3 Files Affected:
- (modified) mlir/lib/Dialect/Math/Transforms/AlgebraicSimplification.cpp (+7-12)
- (modified) mlir/test/Dialect/Complex/powi-simplify.mlir (+12)
- (modified) mlir/test/Dialect/Math/algebraic-simplification.mlir (+80)
``````````diff
diff --git a/mlir/lib/Dialect/Math/Transforms/AlgebraicSimplification.cpp b/mlir/lib/Dialect/Math/Transforms/AlgebraicSimplification.cpp
index eecb40b2689b15..547ccdbefc4b79 100644
--- a/mlir/lib/Dialect/Math/Transforms/AlgebraicSimplification.cpp
+++ b/mlir/lib/Dialect/Math/Transforms/AlgebraicSimplification.cpp
@@ -152,20 +152,15 @@ PowIStrengthReduction<PowIOpTy, DivOpTy, MulOpTy>::matchAndRewrite(
Location loc = op.getLoc();
Value base = op.getLhs();
- IntegerAttr scalarExponent;
- DenseIntElementsAttr vectorExponent;
-
- bool isScalar = matchPattern(op.getRhs(), m_Constant(&scalarExponent));
- bool isVector = matchPattern(op.getRhs(), m_Constant(&vectorExponent));
-
// Simplify cases with known exponent value.
- int64_t exponentValue = 0;
- if (isScalar)
- exponentValue = scalarExponent.getInt();
- else if (isVector && vectorExponent.isSplat())
- exponentValue = vectorExponent.getSplatValue<IntegerAttr>().getInt();
- else
+ APInt exponent;
+ if (!matchPattern(op.getRhs(), m_ConstantInt(&exponent)))
+ return failure();
+
+ // Values outside int64_t exceed the strength reduction threshold.
+ if (!exponent.isSignedIntN(64))
return failure();
+ int64_t exponentValue = exponent.getSExtValue();
// Compute abs(exponent) and check the threshold before creating any IR,
// so that returning failure() here does not violate the pattern API contract.
diff --git a/mlir/test/Dialect/Complex/powi-simplify.mlir b/mlir/test/Dialect/Complex/powi-simplify.mlir
index 9f998d385310b1..8e5a02d32988e0 100644
--- a/mlir/test/Dialect/Complex/powi-simplify.mlir
+++ b/mlir/test/Dialect/Complex/powi-simplify.mlir
@@ -28,3 +28,15 @@ func.func @powi_int64_min(%arg0: complex<f32>) -> complex<f32> {
}
// CHECK-LABEL: func.func @powi_int64_min(
// CHECK: complex.powi %{{.*}}, %{{.*}} : complex<f32>, i64
+
+// Exponents outside the int64_t range must not be truncated.
+func.func @powi_wide_exp(%arg0: complex<f32>) -> complex<f32> {
+ %c = arith.constant 18446744073709551618 : i128
+ %0 = complex.powi %arg0, %c : complex<f32>, i128
+ return %0 : complex<f32>
+}
+// CHECK-LABEL: func.func @powi_wide_exp(
+// CHECK-SAME: %[[ARG:.+]]: complex<f32>
+// CHECK: %[[CST:.*]] = arith.constant 18446744073709551618 : i128
+// CHECK: %[[RES:.*]] = complex.powi %[[ARG]], %[[CST]] : complex<f32>, i128
+// CHECK: return %[[RES]] : complex<f32>
diff --git a/mlir/test/Dialect/Math/algebraic-simplification.mlir b/mlir/test/Dialect/Math/algebraic-simplification.mlir
index e311fb7d6a2e87..85578ce81e820f 100644
--- a/mlir/test/Dialect/Math/algebraic-simplification.mlir
+++ b/mlir/test/Dialect/Math/algebraic-simplification.mlir
@@ -292,6 +292,23 @@ func.func @ipowi_index_exp_int64_min(%arg0: index) -> index {
return %0 : index
}
+// The same pattern handles integer bases. Truncating this exponent to 64 bits
+// would incorrectly simplify the operation to its base.
+// CHECK-LABEL: @ipowi_exp_wide(
+// CHECK-SAME: %[[ARG0:.+]]: i128, %[[ARG1:.+]]: vector<4xi128>
+func.func @ipowi_exp_wide(%arg0: i128, %arg1: vector<4xi128>) -> (i128, vector<4xi128>) {
+ // CHECK-DAG: %[[CST_S:.*]] = arith.constant 18446744073709551617 : i128
+ // CHECK-DAG: %[[CST_V:.*]] = arith.constant dense<18446744073709551617> : vector<4xi128>
+ // CHECK: %[[SCALAR:.*]] = math.ipowi %[[ARG0]], %[[CST_S]] : i128
+ // CHECK: %[[VECTOR:.*]] = math.ipowi %[[ARG1]], %[[CST_V]] : vector<4xi128>
+ // CHECK: return %[[SCALAR]], %[[VECTOR]]
+ %c = arith.constant 18446744073709551617 : i128
+ %v = arith.constant dense<18446744073709551617> : vector<4xi128>
+ %0 = math.ipowi %arg0, %c : i128
+ %1 = math.ipowi %arg1, %v : vector<4xi128>
+ return %0, %1 : i128, vector<4xi128>
+}
+
// CHECK-LABEL: @fpowi_zero_exp(
// CHECK-SAME: %[[ARG0:.+]]: f32
// CHECK-SAME: %[[ARG1:.+]]: vector<4xf32>
@@ -341,6 +358,69 @@ func.func @fpowi_exp_int64_min(%arg0: f32) -> f32 {
return %0 : f32
}
+// Wide exponents outside the int64_t range must not be converted to int64_t.
+// CHECK-LABEL: @fpowi_exp_i104(
+// CHECK-SAME: %[[ARG0:.+]]: f32, %[[ARG1:.+]]: vector<4xf32>
+func.func @fpowi_exp_i104(%arg0: f32, %arg1: vector<4xf32>) -> (f32, vector<4xf32>) {
+ // CHECK-DAG: %[[CST_S:.*]] = arith.constant 18446744073709551578 : i104
+ // CHECK-DAG: %[[CST_V:.*]] = arith.constant dense<18446744073709551578> : vector<4xi104>
+ // CHECK: %[[SCALAR:.*]] = math.fpowi %[[ARG0]], %[[CST_S]] : f32, i104
+ // CHECK: %[[VECTOR:.*]] = math.fpowi %[[ARG1]], %[[CST_V]] : vector<4xf32>, vector<4xi104>
+ // CHECK: return %[[SCALAR]], %[[VECTOR]]
+ %c = arith.constant 18446744073709551578 : i104
+ %v = arith.constant dense<18446744073709551578> : vector<4xi104>
+ %0 = math.fpowi %arg0, %c : f32, i104
+ %1 = math.fpowi %arg1, %v : vector<4xf32>, vector<4xi104>
+ return %0, %1 : f32, vector<4xf32>
+}
+
+// CHECK-LABEL: @fpowi_exp_i112(
+// CHECK-SAME: %[[ARG0:.+]]: f32, %[[ARG1:.+]]: vector<4xf32>
+func.func @fpowi_exp_i112(%arg0: f32, %arg1: vector<4xf32>) -> (f32, vector<4xf32>) {
+ // CHECK-DAG: %[[CST_S:.*]] = arith.constant 18446744073709551558 : i112
+ // CHECK-DAG: %[[CST_V:.*]] = arith.constant dense<18446744073709551558> : vector<4xi112>
+ // CHECK: %[[SCALAR:.*]] = math.fpowi %[[ARG0]], %[[CST_S]] : f32, i112
+ // CHECK: %[[VECTOR:.*]] = math.fpowi %[[ARG1]], %[[CST_V]] : vector<4xf32>, vector<4xi112>
+ // CHECK: return %[[SCALAR]], %[[VECTOR]]
+ %c = arith.constant 18446744073709551558 : i112
+ %v = arith.constant dense<18446744073709551558> : vector<4xi112>
+ %0 = math.fpowi %arg0, %c : f32, i112
+ %1 = math.fpowi %arg1, %v : vector<4xf32>, vector<4xi112>
+ return %0, %1 : f32, vector<4xf32>
+}
+
+// This negative exponent is also the minimum signed value for its bit width.
+// CHECK-LABEL: @fpowi_exp_wide_negative(
+// CHECK-SAME: %[[ARG0:.+]]: f32, %[[ARG1:.+]]: vector<4xf32>
+func.func @fpowi_exp_wide_negative(%arg0: f32, %arg1: vector<4xf32>) -> (f32, vector<4xf32>) {
+ // CHECK-DAG: %[[CST_S:.*]] = arith.constant -18446744073709551616 : i65
+ // CHECK-DAG: %[[CST_V:.*]] = arith.constant dense<-18446744073709551616> : vector<4xi65>
+ // CHECK: %[[SCALAR:.*]] = math.fpowi %[[ARG0]], %[[CST_S]] : f32, i65
+ // CHECK: %[[VECTOR:.*]] = math.fpowi %[[ARG1]], %[[CST_V]] : vector<4xf32>, vector<4xi65>
+ // CHECK: return %[[SCALAR]], %[[VECTOR]]
+ %c = arith.constant -18446744073709551616 : i65
+ %v = arith.constant dense<-18446744073709551616> : vector<4xi65>
+ %0 = math.fpowi %arg0, %c : f32, i65
+ %1 = math.fpowi %arg1, %v : vector<4xf32>, vector<4xi65>
+ return %0, %1 : f32, vector<4xf32>
+}
+
+// Small values in wide integer types should still be simplified.
+// CHECK-LABEL: @fpowi_exp_wide_small(
+// CHECK-SAME: %[[ARG0:.+]]: f32, %[[ARG1:.+]]: vector<4xf32>
+func.func @fpowi_exp_wide_small(%arg0: f32, %arg1: vector<4xf32>) -> (f32, vector<4xf32>) {
+ // CHECK: %[[ONE:.*]] = arith.constant dense<1.000000e+00> : vector<4xf32>
+ // CHECK: %[[SCALAR:.*]] = arith.mulf %[[ARG0]], %[[ARG0]] : f32
+ // CHECK: %[[SQUARE:.*]] = arith.mulf %[[ARG1]], %[[ARG1]] : vector<4xf32>
+ // CHECK: %[[VECTOR:.*]] = arith.divf %[[ONE]], %[[SQUARE]] : vector<4xf32>
+ // CHECK: return %[[SCALAR]], %[[VECTOR]]
+ %c = arith.constant 2 : i104
+ %v = arith.constant dense<-2> : vector<4xi112>
+ %0 = math.fpowi %arg0, %c : f32, i104
+ %1 = math.fpowi %arg1, %v : vector<4xf32>, vector<4xi112>
+ return %0, %1 : f32, vector<4xf32>
+}
+
// CHECK-LABEL: @fpowi_exp_two(
// CHECK-SAME: %[[ARG0:.+]]: f32
// CHECK-SAME: %[[ARG1:.+]]: vector<4xf32>
``````````
</details>
https://github.com/llvm/llvm-project/pull/237513
More information about the Mlir-commits
mailing list