[MLIR][Math] Move exponent threshold check before IR creation in PowIStrengthReduction

[joker-eph]

PowIStrengthReduction::matchAndRewrite was creating the one constant (using complex::ConstantOp::create for complex::PowiOp) before the threshold check that guards whether the rewrite is profitable. When the exponent exceeds the threshold, the pattern returned failure() after IR was already modified, violating MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS.

Fix: reorder so the abs(exponent) computation and threshold check occur before any IR creation.

Assisted-by: Claude Code
Fix a failure present with MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS=ON.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-math

Author: Mehdi Amini (joker-eph)

Changes

PowIStrengthReduction::matchAndRewrite was creating the one constant (using complex::ConstantOp::create for complex::PowiOp) before the threshold check that guards whether the rewrite is profitable. When the exponent exceeds the threshold, the pattern returned failure() after IR was already modified, violating MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS.

Fix: reorder so the abs(exponent) computation and threshold check occur before any IR creation.

Assisted-by: Claude Code
Fix a failure present with MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS=ON.

---

Full diff: https://github.com/llvm/llvm-project/pull/188955.diff

1 Files Affected:

• (modified) mlir/lib/Dialect/Math/Transforms/AlgebraicSimplification.cpp (+12-10)

diff --git a/mlir/lib/Dialect/Math/Transforms/AlgebraicSimplification.cpp b/mlir/lib/Dialect/Math/Transforms/AlgebraicSimplification.cpp
index ff5f7f685903f..7ce7b2153ffb4 100644
--- a/mlir/lib/Dialect/Math/Transforms/AlgebraicSimplification.cpp
+++ b/mlir/lib/Dialect/Math/Transforms/AlgebraicSimplification.cpp
@@ -166,6 +166,18 @@ PowIStrengthReduction<PowIOpTy, DivOpTy, MulOpTy>::matchAndRewrite(
   else
     return failure();
 
+  // Compute abs(exponent) and check the threshold before creating any IR,
+  // so that returning failure() here does not violate the pattern API contract.
+  bool exponentIsNegative = false;
+  if (exponentValue < 0) {
+    exponentIsNegative = true;
+    exponentValue *= -1;
+  }
+
+  // Bail out if `abs(exponent)` exceeds the threshold (exponent==0 is free).
+  if (exponentValue != 0 && exponentValue > exponentThreshold)
+    return failure();
+
   // Maybe broadcasts scalar value into vector type compatible with `op`.
   auto bcast = [&loc, &op, &rewriter](Value value) -> Value {
     if (auto vec = dyn_cast<VectorType>(op.getType()))
@@ -196,16 +208,6 @@ PowIStrengthReduction<PowIOpTy, DivOpTy, MulOpTy>::matchAndRewrite(
     return success();
   }
 
-  bool exponentIsNegative = false;
-  if (exponentValue < 0) {
-    exponentIsNegative = true;
-    exponentValue *= -1;
-  }
-
-  // Bail out if `abs(exponent)` exceeds the threshold.
-  if (exponentValue > exponentThreshold)
-    return failure();
-
   Value result = base;
   // Transform to naive sequence of multiplications:
   //   * For positive exponent case replace: