[Mlir-commits] [mlir] [MLIR][Math] Move exponent threshold check before IR creation in PowIStrengthReduction (PR #188955)

[Mehdi Amini]

https://github.com/joker-eph created https://github.com/llvm/llvm-project/pull/188955

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.

>From 8147dbbd3fee78fec4200aa943ffe38f719f73db Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Thu, 26 Mar 2026 15:57:41 -0700
Subject: [PATCH] [MLIR][Math] Move exponent threshold check before IR creation
 in PowIStrengthReduction

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.
---
 .../Transforms/AlgebraicSimplification.cpp    | 22 ++++++++++---------
 1 file changed, 12 insertions(+), 10 deletions(-)

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: