[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.
+  if (isNegative) {
+    Value zeroEqCheck =
+        b.create<arith::CmpFOp>(arith::CmpFPredicate::OEQ, res, zero);
+    Value negZeroEqCheck =
+        b.create<arith::CmpFOp>(arith::CmpFPredicate::OEQ, res, negZero);
+    res = b.create<arith::DivFOp>(baseType, one, res);
+    res =
+        b.create<arith::SelectOp>(op->getLoc(), zeroEqCheck, posInfinity, res);
+    res = b.create<arith::SelectOp>(op->getLoc(), negZeroEqCheck, negInfinity,
+                                    res);
+  }
----------------
kuhar wrote:

What happens when both the base and the power are zero? I checked the llvm langref and it doesn't mention this case: https://llvm.org/docs/LangRef.html#llvm-powi-intrinsic . The c standard library is more informative here: https://en.cppreference.com/w/cpp/numeric/math/pow#:~:text=in%20math_errhandling.-,If,is%20negative%2C%20a%20domain%20error%20or%20a%20pole%20error%20may%20occur.,-If%20the%20implementation .

Because neither `math` nor `llvm` define this, I don't think we have to worry about this, but it would be nice to have a comment that explains these corner cases.


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