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

[Jakub Kuderski]

================
@@ -202,6 +202,47 @@ 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 convertFPowIOp(math::FPowIOp op,
+                                    PatternRewriter &rewriter) {
+  ImplicitLocOpBuilder b(op->getLoc(), rewriter);
+  Value operandA = op.getOperand(0);
+  Value operandB = op.getOperand(1);
+  Type opType = operandA.getType();
+
+  Attribute cstAttr;
+  if (!matchPattern(operandB, m_Constant(&cstAttr)))
+    return failure();
+
+  auto iAttr = dyn_cast<SplatElementsAttr>(cstAttr);
+
+  if (!iAttr)
+    return failure();
+
+  int64_t power = iAttr.getSplatValue<int64_t>();
+  bool neg = power < 0;
+  int64_t absPower = std::abs(power);
+  Value one = createFloatConst(op->getLoc(), opType, 1.00, rewriter);
+  Value res = createFloatConst(op->getLoc(), opType, 1.00, rewriter);
+
+  while (absPower > 0) {
+
+    if (absPower & 1)
+      res = b.create<arith::MulFOp>(opType, operandA, res);
+
+    absPower = absPower >> 1;
+    operandA = b.create<arith::MulFOp>(opType, operandA, operandA);
+  }
+
+  if (neg)
+    res = b.create<arith::DivFOp>(opType, one, res);
----------------
kuhar wrote:

could this lead to UB on division by zero? I think this could happen when lhs is zero. Is pow defined in this case? The docs don't mention this AFAICT: https://mlir.llvm.org/docs/Dialects/MathOps/#mathfpowi-mathfpowiop

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