[Mlir-commits] [mlir] [mlir][math] Expand powfI operation for constant power operand. (PR #87081)

[Jakub Kuderski]

================
@@ -202,6 +210,70 @@ static LogicalResult convertCeilOp(math::CeilOp op, PatternRewriter &rewriter) {
   rewriter.replaceOp(op, ret);
   return success();
 }
+
+// Convert `math.fpowi` to a series of `arith.mulf` operations.
+// If the power is negative, we divide one by the result.
+static LogicalResult convertFPowICstOp(math::FPowIOp op,
+                                       PatternRewriter &rewriter) {
+  ImplicitLocOpBuilder b(op->getLoc(), rewriter);
+  Value base = op.getOperand(0);
+  Value power = op.getOperand(1);
+  Type baseType = base.getType();
+  Value tempBase = op.getOperand(0);
+
+  Attribute cstAttr;
+  if (!matchPattern(power, m_Constant(&cstAttr)))
+    return failure();
+
+  APInt value;
+  if (!matchPattern(cstAttr, m_ConstantInt(&value)))
+    return failure();
+
+  int64_t powerInt = value.getSExtValue();
+  bool isNegative = powerInt < 0;
+  int64_t absPower = std::abs(powerInt);
+  Value one = createFloatConst(op->getLoc(), baseType, 1.00, rewriter);
+  Value res = createFloatConst(op->getLoc(), baseType, 1.00, rewriter);
+
+  auto &sem = dyn_cast<mlir::FloatType>(getElementTypeOrSelf(baseType))
+                  .getFloatSemantics();
+  Value zero =
+      createFloatConst(op->getLoc(), baseType,
+                       APFloat::getZero(sem, /*Negative=*/false), rewriter);
+  Value negZero =
+      createFloatConst(op->getLoc(), baseType,
+                       APFloat::getZero(sem, /*Negative=*/true), rewriter);
+  Value posInfinity =
+      createFloatConst(op->getLoc(), baseType,
+                       APFloat::getInf(sem, /*Negative=*/false), rewriter);
+  Value negInfinity =
+      createFloatConst(op->getLoc(), baseType,
+                       APFloat::getInf(sem, /*Negative=*/true), rewriter);
+
+  while (absPower > 0) {
+    if (absPower & 1)
+      res = b.create<arith::MulFOp>(baseType, tempBase, res);
+    absPower >>= 1;
+    tempBase = b.create<arith::MulFOp>(baseType, tempBase, tempBase);
+  }
+
+  // Take care of UB in case of negative power.
----------------
kuhar wrote:

```suggestion
  // Make sure not to introduce UB in case of negative power.
```

https://github.com/llvm/llvm-project/pull/87081