[mlir][math] Added basic support for IPowI operation.

[vzakhari]

Summary

The operation computes pow(b, p), where 'b' and 'p' are signed integers
of the same width. The result's type matches the operands' type.

[kiranchandramohan]

Would math.powi or math.powsi be a better name? Is it OK to move the folder and strength reduction to dependent patches?

[vzakhari]

In D129809#3663012, @kiranchandramohan wrote:

Would math.powi or math.powsi be a better name? Is it OK to move the folder and strength reduction to dependent patches?

Thank you for looking into this!

There is D129811, where I try to add FPowSI, so I want to distinguish power operations with integer and floating bases. They are somewhat different in behavior, e.g. FPowSI may have some floating point specific effects (like using rounding mode, trapping FP exceptions), whereas IPowSI does not have any FP-related effects.

Sure, I can split it into three patches: basic, folder and strength reduction. Please let me know if you have major concerns about this patch and D129811. If you do not have any, I will proceed with the splitting.

[kiranchandramohan]

In D129809#3663033, @vzakhari wrote:

In D129809#3663012, @kiranchandramohan wrote:

Would math.powi or math.powsi be a better name? Is it OK to move the folder and strength reduction to dependent patches?

Thank you for looking into this!

There is D129811, where I try to add FPowSI, so I want to distinguish power operations with integer and floating bases. They are somewhat different in behavior, e.g. FPowSI may have some floating point specific effects (like using rounding mode, trapping FP exceptions), whereas IPowSI does not have any FP-related effects.

Ahh OK.

From what I understand, looking at other operations in the arithmetic dialect, we specify the signed semantics as part of the operation name only if we want to distinguish the signed and unsigned behaviour (like arith.divsi vs. arith.divui). In this case, we always want a signed semantics so can't we just call it fpowi, ipowi, just like arith.addi? Or did I miss a point?

I see that LLVM IR does not have an equivalent to the ipowsi instruction. Could this be due to C/C++ not having the exponentiation operation in the language or in the math library? I wonder what other languages that lower to llvm IR do.

You might want to use the Depends on <phabricator-patch-number> syntax to mark dependencies between the fpowsi and ipowsi patch.

[vzakhari]

In D129809#3664012, @kiranchandramohan wrote:

In D129809#3663033, @vzakhari wrote:

In D129809#3663012, @kiranchandramohan wrote:

Would math.powi or math.powsi be a better name? Is it OK to move the folder and strength reduction to dependent patches?

Thank you for looking into this!

There is D129811, where I try to add FPowSI, so I want to distinguish power operations with integer and floating bases. They are somewhat different in behavior, e.g. FPowSI may have some floating point specific effects (like using rounding mode, trapping FP exceptions), whereas IPowSI does not have any FP-related effects.

Ahh OK.

From what I understand, looking at other operations in the arithmetic dialect, we specify the signed semantics as part of the operation name only if we want to distinguish the signed and unsigned behaviour (like arith.divsi vs. arith.divui). In this case, we always want a signed semantics so can't we just call it fpowi, ipowi, just like arith.addi? Or did I miss a point?

Yes, I just wanted to make clear that the power operand is signed in the operation name. I can use just 'i' instead of 'si', since I do not expect that we will need a version with unsigned power operand.

I see that LLVM IR does not have an equivalent to the ipowsi instruction. Could this be due to C/C++ not having the exponentiation operation in the language or in the math library? I wonder what other languages that lower to llvm IR do.

Yes, math.h does not define pow functions with integer power argument as well as integer base argument. C++ does have std::pow overloads with int32 power argument, but they are implemented via math.h's double precision pow (so I think their performance is not as good as it could have been if there was a specialized version as in libgcc). Not sure what the other languages do.

You might want to use the Depends on <phabricator-patch-number> syntax to mark dependencies between the fpowsi and ipowsi patch.

Thanks for the hint! I usually use web interface to add parent revisions. In this case I have a chain of 4 revisions that I was writing one after another, so I just used linear chain: D129809->D129810->D129811->D129812

[ftynse]

I think this stack of patches deserves a quick RFC on the forum.

[vzakhari]

In D129809#3665397, @ftynse wrote:

I think this stack of patches deserves a quick RFC on the forum.

Thanks! Created RFC: https://discourse.llvm.org/t/rfc-adding-more-power-operations-into-math-dialect/63975

[vzakhari]

@kiranchandramohan, @ftynse thank you for your comments!

I split this patch into 3 pieces. New parts are in D130389 and D130390.

I will be unavailable next week, so I will not be able to reply. Sorry for the inconvenience.

[vzakhari]

rebase

[vzakhari]

rebase

Please review.

[vzakhari]

Thank you, @jfurtek!

mlir/include/mlir/Dialect/Math/IR/MathOps.td

 class Math_FloatUnaryOp<string mnemonic, list<Trait> traits = []> :
     Math_Op<mnemonic, traits # [SameOperandsAndResultType]> {
   let arguments = (ins FloatLike:$operand);
   let results = (outs FloatLike:$result);
 
   let assemblyFormat = "$operand attr-dict `:` type($result)";
 }
 
+// Base class for binary math operations on integer types. Require two
+// operands and one result of the same type. This type can be an integer
+// type, vector or tensor thereof.
+class Math_IntegerBinaryOp<string mnemonic, list<Trait> traits = []> :
+    Math_Op<mnemonic, traits # [SameOperandsAndResultType]> {
+  let arguments = (ins SignlessIntegerLike:$lhs, SignlessIntegerLike:$rhs);
+  let results = (outs SignlessIntegerLike:$result);
+
+  let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($result)";
+}
+
 // Base class for binary math operations on floating point types. Require two
 // operands and one result of the same type. This type can be a floating point
 // type, vector or tensor thereof.
 class Math_FloatBinaryOp<string mnemonic, list<Trait> traits = []> :
     Math_Op<mnemonic, traits # [SameOperandsAndResultType]> {
   let arguments = (ins FloatLike:$lhs, FloatLike:$rhs);
   let results = (outs FloatLike:$result);
 
let description = [{
     The semantics of the operation correspond to those of the `llvm.fma`
     [intrinsic](https://llvm.org/docs/LangRef.html#llvm-fma-intrinsic). In the
     particular case of lowering to LLVM, this is guaranteed to lower
     to the `llvm.fma.*` intrinsic.
   }];
 }
 
 //===----------------------------------------------------------------------===//
+// IPowSIOp
+//===----------------------------------------------------------------------===//
+
+def Math_IPowSIOp : Math_IntegerBinaryOp<"ipowsi"> {
+  let summary = "signed integer raised to the power of operation";
+  let description = [{
+    Syntax:
+
+    ```
+    operation ::= ssa-id `=` `math.ipowsi` ssa-use `,` ssa-use `:` type
+    ```
+
+    The `ipowsi` operation takes two operands of integer type (i.e., scalar,
+    tensor or vector) and returns one result of the same type. Operands
+    must have the same type.
+
+    Example:
+
+    ```mlir
+    // Scalar signed integer exponentiation.
+    %a = math.ipowsi %b, %c : i32
+    ```
+  }];
+  let hasFolder = 1;
+}
+
+//===----------------------------------------------------------------------===//
 // LogOp
 //===----------------------------------------------------------------------===//
 
 def Math_LogOp : Math_FloatUnaryOp<"log"> {
   let summary = "base-e logarithm of the specified value";
 
   let description = [{
     Computes the base-e logarithm of the given value. It takes one operand of

mlir/lib/Dialect/Math/IR/MathOps.cpp

 
 OpFoldResult math::CtPopOp::fold(ArrayRef<Attribute> operands) {
   return constFoldUnaryOp<IntegerAttr>(operands, [](const APInt &a) {
     return APInt(a.getBitWidth(), a.countPopulation());
   });
 }
 
 //===----------------------------------------------------------------------===//
+// IPowSIOp folder
+//===----------------------------------------------------------------------===//
+
+OpFoldResult math::IPowSIOp::fold(ArrayRef<Attribute> operands) {
+  // TODO: implement.
+  return constFoldBinaryOpConditional<IntegerAttr>(
+      operands, [](const APInt &base, const APInt &power) -> Optional<APInt> {
+        unsigned width = base.getBitWidth();
+        auto zeroValue = APInt::getZero(width);
+        APInt oneValue{width, 1ULL, /*isSigned=*/true};
+        APInt minusOneValue{width, -1ULL, /*isSigned=*/true};
+
+        if (power.isZero())
+          return oneValue;
+
+        if (power.isNegative()) {
+          // Leave 0 raised to negative power not folded.
+          if (base.isZero())
+            return {};
+          if (base.eq(oneValue))
+            return oneValue;
+          // If abs(base) > 1, then the result is zero.
+          if (base.ne(minusOneValue))
+            return zeroValue;
+          // base == -1:
+          //   -1: power is odd
+          //    1: power is even
+          if (power[0] == 1)
+            return minusOneValue;
+
+          return oneValue;
+        }
+
+        // power is positive.
+        APInt result = oneValue;
+        APInt curBase = base;
+        APInt curPower = power;
+        while (true) {
+          if (curPower[0] == 1)
+            result *= curBase;
+          curPower.lshrInPlace(1);
+          if (curPower.isZero())
+            return result;
+          curBase *= curBase;
+        }
+      });
+
+  return Attribute();
+}
+
+//===----------------------------------------------------------------------===//
 // Log2Op folder
 //===----------------------------------------------------------------------===//
 
 OpFoldResult math::Log2Op::fold(ArrayRef<Attribute> operands) {
   return constFoldUnaryOpConditional<FloatAttr>(
       operands, [](const APFloat &a) -> Optional<APFloat> {
         if (a.isNegative())
           return {};

mlir/lib/Dialect/Math/Transforms/AlgebraicSimplification.cpp

if (isExponentValue(-0.5)) {
     rewriter.replaceOpWithNewOp<math::RsqrtOp>(op, x);
     return success();
   }
 
   return failure();
 }
 
 //----------------------------------------------------------------------------//
+// IPowSIOp strength reduction.
+//----------------------------------------------------------------------------//
+
+namespace {
+struct IPowSIStrengthReduction : public OpRewritePattern<math::IPowSIOp> {
+  unsigned exponentThreshold;
+
+public:
+  using OpRewritePattern::OpRewritePattern;
+
+  IPowSIStrengthReduction(MLIRContext *context, unsigned exponentThreshold = 3,
+                          PatternBenefit benefit = 1,
+                          ArrayRef<StringRef> generatedNames = {})
+      : OpRewritePattern<math::IPowSIOp>(context, benefit, generatedNames),
+        exponentThreshold(exponentThreshold) {}
+  LogicalResult matchAndRewrite(math::IPowSIOp op,
+                                PatternRewriter &rewriter) const final;
+};
+} // namespace
+
+LogicalResult
+IPowSIStrengthReduction::matchAndRewrite(math::IPowSIOp op,
+                                         PatternRewriter &rewriter) const {
+  Location loc = op.getLoc();
+  Value base = op.getLhs();
+
+  IntegerAttr scalarExponent;
+  DenseIntElementsAttr vectorExponent;
+
+  bool isScalar = matchPattern(op.getRhs(), m_Constant(&scalarExponent));
+  bool isVector = matchPattern(op.getRhs(), m_Constant(&vectorExponent));
+
+  int64_t exponentValue = 0;
+  if (isScalar)
+    exponentValue = scalarExponent.getInt();
+  else if (isVector && vectorExponent.isSplat())
+    exponentValue = vectorExponent.getSplatValue<IntegerAttr>().getInt();
+  else
+    return failure();
+
+  // Maybe broadcasts scalar value into vector type compatible with `op`.
+  auto bcast = [&](Value value) -> Value {
+    if (auto vec = op.getType().dyn_cast<VectorType>())
+      return rewriter.create<vector::BroadcastOp>(loc, vec, value);
+    return value;
+  };
+
+  if (exponentValue == 0) {
+    Value one = rewriter.create<arith::ConstantOp>(
+        loc, rewriter.getIntegerAttr(getElementTypeOrSelf(op.getType()), 1));
+    rewriter.replaceOp(op, bcast(one));
+    return success();
+  }
+
+  bool exponentIsNegative = false;
+  if (exponentValue < 0) {
+    exponentIsNegative = true;
+    exponentValue *= -1;
+  }
+
+  if (exponentValue > exponentThreshold)
+    return failure();
+
+  // Inverse the base for negative exponent.
+  if (exponentIsNegative) {
+    Value one = rewriter.create<arith::ConstantOp>(
+        loc, rewriter.getIntegerAttr(getElementTypeOrSelf(op.getType()), 1));
+    base = rewriter.create<arith::DivSIOp>(loc, bcast(one), base);
+  }
+
+  Value result = base;
+  // Transform to naive sequence of multiplications.
+  for (unsigned i = 1; i < exponentValue; ++i)
+    result = rewriter.create<arith::MulIOp>(loc, result, base);
+
+  rewriter.replaceOp(op, result);
+  return success();
+}
+
+//----------------------------------------------------------------------------//
 
 void mlir::populateMathAlgebraicSimplificationPatterns(
     RewritePatternSet &patterns) {
   patterns.add<PowFStrengthReduction>(patterns.getContext());
+  patterns.add<IPowSIStrengthReduction>(patterns.getContext());
 }

mlir/test/Dialect/Math/algebraic-simplification.mlir

func.func @pow_rsqrt(%arg0: f32, %arg1 : vector<4xf32>) -> (f32, vector<4xf32>) {
   // CHECK: %[[VECTOR:.*]] = math.rsqrt %arg1
   // CHECK: return %[[SCALAR]], %[[VECTOR]]
   %c = arith.constant -0.5 : f32
   %v = arith.constant dense <-0.5> : vector<4xf32>
   %0 = math.powf %arg0, %c : f32
   %1 = math.powf %arg1, %v : vector<4xf32>
   return %0, %1 : f32, vector<4xf32>
 }
+
+// CHECK-LABEL: @ipowi_zero_exp(
+// CHECK-SAME: %[[ARG0:.+]]: i32
+// CHECK-SAME: %[[ARG1:.+]]: vector<4xi32>
+// CHECK-SAME: -> (i32, vector<4xi32>) {
+func.func @ipowi_zero_exp(%arg0: i32, %arg1: vector<4xi32>) -> (i32, vector<4xi32>) {
+  // CHECK: %[[CST_S:.*]] = arith.constant 1 : i32
+  // CHECK: %[[CST_V:.*]] = arith.constant dense<1> : vector<4xi32>
+  // CHECK: return %[[CST_S]], %[[CST_V]]
+  %c = arith.constant 0 : i32
+  %v = arith.constant dense <0> : vector<4xi32>
+  %0 = math.ipowsi %arg0, %c : i32
+  %1 = math.ipowsi %arg1, %v : vector<4xi32>
+  return %0, %1 : i32, vector<4xi32>
+}
+
+// CHECK-LABEL: @ipowi_exp_one(
+// CHECK-SAME: %[[ARG0:.+]]: i32
+// CHECK-SAME: %[[ARG1:.+]]: vector<4xi32>
+// CHECK-SAME: -> (i32, vector<4xi32>, i32, vector<4xi32>) {
+func.func @ipowi_exp_one(%arg0: i32, %arg1: vector<4xi32>) -> (i32, vector<4xi32>, i32, vector<4xi32>) {
+  // CHECK: %[[CST_S:.*]] = arith.constant 1 : i32
+  // CHECK: %[[CST_V:.*]] = arith.constant dense<1> : vector<4xi32>
+  // CHECK: %[[SCALAR:.*]] = arith.divsi %[[CST_S]], %[[ARG0]]
+  // CHECK: %[[VECTOR:.*]] = arith.divsi %[[CST_V]], %[[ARG1]]
+  // CHECK: return %[[ARG0]], %[[ARG1]], %[[SCALAR]], %[[VECTOR]]
+  %c1 = arith.constant 1 : i32
+  %v1 = arith.constant dense <1> : vector<4xi32>
+  %0 = math.ipowsi %arg0, %c1 : i32
+  %1 = math.ipowsi %arg1, %v1 : vector<4xi32>
+  %cm1 = arith.constant -1 : i32
+  %vm1 = arith.constant dense <-1> : vector<4xi32>
+  %2 = math.ipowsi %arg0, %cm1 : i32
+  %3 = math.ipowsi %arg1, %vm1 : vector<4xi32>
+  return %0, %1, %2, %3 : i32, vector<4xi32>, i32, vector<4xi32>
+}
+
+// CHECK-LABEL: @ipowi_exp_two(
+// CHECK-SAME: %[[ARG0:.+]]: i32
+// CHECK-SAME: %[[ARG1:.+]]: vector<4xi32>
+// CHECK-SAME: -> (i32, vector<4xi32>, i32, vector<4xi32>) {
+func.func @ipowi_exp_two(%arg0: i32, %arg1: vector<4xi32>) -> (i32, vector<4xi32>, i32, vector<4xi32>) {
+  // CHECK: %[[CST_S:.*]] = arith.constant 1 : i32
+  // CHECK: %[[CST_V:.*]] = arith.constant dense<1> : vector<4xi32>
+  // CHECK: %[[SCALAR0:.*]] = arith.muli %[[ARG0]], %[[ARG0]]
+  // CHECK: %[[VECTOR0:.*]] = arith.muli %[[ARG1]], %[[ARG1]]
+  // CHECK: %[[SCALAR1:.*]] = arith.divsi %[[CST_S]], %[[ARG0]]
+  // CHECK: %[[SMUL:.*]] = arith.muli %[[SCALAR1]], %[[SCALAR1]]
+  // CHECK: %[[VECTOR1:.*]] = arith.divsi %[[CST_V]], %[[ARG1]]
+  // CHECK: %[[VMUL:.*]] = arith.muli %[[VECTOR1]], %[[VECTOR1]]
+  // CHECK: return %[[SCALAR0]], %[[VECTOR0]], %[[SMUL]], %[[VMUL]]
+  %c1 = arith.constant 2 : i32
+  %v1 = arith.constant dense <2> : vector<4xi32>
+  %0 = math.ipowsi %arg0, %c1 : i32
+  %1 = math.ipowsi %arg1, %v1 : vector<4xi32>
+  %cm1 = arith.constant -2 : i32
+  %vm1 = arith.constant dense <-2> : vector<4xi32>
+  %2 = math.ipowsi %arg0, %cm1 : i32
+  %3 = math.ipowsi %arg1, %vm1 : vector<4xi32>
+  return %0, %1, %2, %3 : i32, vector<4xi32>, i32, vector<4xi32>
+}
+
+// CHECK-LABEL: @ipowi_exp_three(
+// CHECK-SAME: %[[ARG0:.+]]: i32
+// CHECK-SAME: %[[ARG1:.+]]: vector<4xi32>
+// CHECK-SAME: -> (i32, vector<4xi32>, i32, vector<4xi32>) {
+func.func @ipowi_exp_three(%arg0: i32, %arg1: vector<4xi32>) -> (i32, vector<4xi32>, i32, vector<4xi32>) {
+  // CHECK: %[[CST_S:.*]] = arith.constant 1 : i32
+  // CHECK: %[[CST_V:.*]] = arith.constant dense<1> : vector<4xi32>
+  // CHECK: %[[SMUL0:.*]] = arith.muli %[[ARG0]], %[[ARG0]]
+  // CHECK: %[[SCALAR0:.*]] = arith.muli %[[SMUL0]], %[[ARG0]]
+  // CHECK: %[[VMUL0:.*]] = arith.muli %[[ARG1]], %[[ARG1]]
+  // CHECK: %[[VECTOR0:.*]] = arith.muli %[[VMUL0]], %[[ARG1]]
+  // CHECK: %[[SCALAR1:.*]] = arith.divsi %[[CST_S]], %[[ARG0]]
+  // CHECK: %[[SMUL1:.*]] = arith.muli %[[SCALAR1]], %[[SCALAR1]]
+  // CHECK: %[[SMUL2:.*]] = arith.muli %[[SMUL1]], %[[SCALAR1]]
+  // CHECK: %[[VECTOR1:.*]] = arith.divsi %[[CST_V]], %[[ARG1]]
+  // CHECK: %[[VMUL1:.*]] = arith.muli %[[VECTOR1]], %[[VECTOR1]]
+  // CHECK: %[[VMUL2:.*]] = arith.muli %[[VMUL1]], %[[VECTOR1]]
+  // CHECK: return %[[SCALAR0]], %[[VECTOR0]], %[[SMUL2]], %[[VMUL2]]
+  %c1 = arith.constant 3 : i32
+  %v1 = arith.constant dense <3> : vector<4xi32>
+  %0 = math.ipowsi %arg0, %c1 : i32
+  %1 = math.ipowsi %arg1, %v1 : vector<4xi32>
+  %cm1 = arith.constant -3 : i32
+  %vm1 = arith.constant dense <-3> : vector<4xi32>
+  %2 = math.ipowsi %arg0, %cm1 : i32
+  %3 = math.ipowsi %arg1, %vm1 : vector<4xi32>
+  return %0, %1, %2, %3 : i32, vector<4xi32>, i32, vector<4xi32>
+}

mlir/test/Dialect/Math/canonicalize_ipowsi.mlir

+// RUN: mlir-opt %s -canonicalize | FileCheck %s 
+
+// CHECK-LABEL: @ipowsi32_fold(
+// CHECK-SAME: %[[result:.+]]: memref<?xi32>
+func.func @ipowsi32_fold(%result : memref<?xi32>) {
+// CHECK-DAG: %[[cst0:.+]] = arith.constant 0 : i32
+// CHECK-DAG: %[[cst1:.+]] = arith.constant 1 : i32
+// CHECK-DAG: %[[cst1073741824:.+]] = arith.constant 1073741824 : i32
+// CHECK-DAG: %[[cst_m1:.+]] = arith.constant -1 : i32
+// CHECK-DAG: %[[cst_m27:.+]] = arith.constant -27 : i32
+// CHECK-DAG: %[[i0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[i1:.+]] = arith.constant 1 : index
+// CHECK-DAG: %[[i2:.+]] = arith.constant 2 : index
+// CHECK-DAG: %[[i3:.+]] = arith.constant 3 : index
+// CHECK-DAG: %[[i4:.+]] = arith.constant 4 : index
+// CHECK-DAG: %[[i5:.+]] = arith.constant 5 : index
+// CHECK-DAG: %[[i6:.+]] = arith.constant 6 : index
+// CHECK-DAG: %[[i7:.+]] = arith.constant 7 : index
+// CHECK-DAG: %[[i8:.+]] = arith.constant 8 : index
+// CHECK-DAG: %[[i9:.+]] = arith.constant 9 : index
+// CHECK-DAG: %[[i10:.+]] = arith.constant 10 : index
+// CHECK-DAG: %[[i11:.+]] = arith.constant 11 : index
+
+// --- Test power == 0 ---
+  %arg0_base = arith.constant 0 : i32
+  %arg0_power = arith.constant 0 : i32
+  %res0 = math.ipowsi %arg0_base, %arg0_power : i32
+  %i0 = arith.constant 0 : index
+  memref.store %res0, %result[%i0] : memref<?xi32>
+// CHECK: memref.store %[[cst1]], %[[result]][%[[i0]]] : memref<?xi32>
+
+  %arg1_base = arith.constant 10 : i32
+  %arg1_power = arith.constant 0 : i32
+  %res1 = math.ipowsi %arg1_base, %arg1_power : i32
+  %i1 = arith.constant 1 : index
+  memref.store %res1, %result[%i1] : memref<?xi32>
+// CHECK: memref.store %[[cst1]], %[[result]][%[[i1]]] : memref<?xi32>
+
+  %arg2_base = arith.constant -10 : i32
+  %arg2_power = arith.constant 0 : i32
+  %res2 = math.ipowsi %arg2_base, %arg2_power : i32
+  %i2 = arith.constant 2 : index
+  memref.store %res2, %result[%i2] : memref<?xi32>
+// CHECK: memref.store %[[cst1]], %[[result]][%[[i2]]] : memref<?xi32>
+
+// --- Test negative powers ---
+  %arg3_base = arith.constant 0 : i32
+  %arg3_power = arith.constant -1 : i32
+  %res3 = math.ipowsi %arg3_base, %arg3_power : i32
+  %i3 = arith.constant 3 : index
+  memref.store %res3, %result[%i3] : memref<?xi32>
+// No folding for ipowsi(0, x) for x < 0:
+// CHECK: %[[res3:.+]] = math.ipowsi %[[cst0]], %[[cst_m1]] : i32
+// CHECK: memref.store %[[res3]], %[[result]][%[[i3]]] : memref<?xi32>
+
+  %arg4_base = arith.constant 1 : i32
+  %arg4_power = arith.constant -10 : i32
+  %res4 = math.ipowsi %arg4_base, %arg4_power : i32
+  %i4 = arith.constant 4 : index
+  memref.store %res4, %result[%i4] : memref<?xi32>
+// CHECK: memref.store %[[cst1]], %[[result]][%[[i4]]] : memref<?xi32>
+
+  %arg5_base = arith.constant 2 : i32
+  %arg5_power = arith.constant -1 : i32
+  %res5 = math.ipowsi %arg5_base, %arg5_power : i32
+  %i5 = arith.constant 5 : index
+  memref.store %res5, %result[%i5] : memref<?xi32>
+// CHECK: memref.store %[[cst0]], %[[result]][%[[i5]]] : memref<?xi32>
+
+  %arg6_base = arith.constant -2 : i32
+  %arg6_power = arith.constant -1 : i32
+  %res6 = math.ipowsi %arg6_base, %arg6_power : i32
+  %i6 = arith.constant 6 : index
+  memref.store %res6, %result[%i6] : memref<?xi32>
+// CHECK: memref.store %[[cst0]], %[[result]][%[[i6]]] : memref<?xi32>
+
+  %arg7_base = arith.constant -1 : i32
+  %arg7_power = arith.constant -10 : i32
+  %res7 = math.ipowsi %arg7_base, %arg7_power : i32
+  %i7 = arith.constant 7 : index
+  memref.store %res7, %result[%i7] : memref<?xi32>
+// CHECK: memref.store %[[cst1]], %[[result]][%[[i7]]] : memref<?xi32>
+
+  %arg8_base = arith.constant -1 : i32
+  %arg8_power = arith.constant -11 : i32
+  %res8 = math.ipowsi %arg8_base, %arg8_power : i32
+  %i8 = arith.constant 8 : index
+  memref.store %res8, %result[%i8] : memref<?xi32>
+// CHECK: memref.store %[[cst_m1]], %[[result]][%[[i8]]] : memref<?xi32>
+
+// --- Test positive powers ---
+  %arg9_base = arith.constant -3 : i32
+  %arg9_power = arith.constant 3 : i32
+  %res9 = math.ipowsi %arg9_base, %arg9_power : i32
+  %i9 = arith.constant 9 : index
+  memref.store %res9, %result[%i9] : memref<?xi32>
+// CHECK: memref.store %[[cst_m27]], %[[result]][%[[i9]]] : memref<?xi32>
+
+  %arg10_base = arith.constant 2 : i32
+  %arg10_power = arith.constant 30 : i32
+  %res10 = math.ipowsi %arg10_base, %arg10_power : i32
+  %i10 = arith.constant 10 : index
+  memref.store %res10, %result[%i10] : memref<?xi32>
+// CHECK: memref.store %[[cst1073741824]], %[[result]][%[[i10]]] : memref<?xi32>
+
+// --- Test vector folding ---
+  %arg11_base = arith.constant 2 : i32
+  %arg11_base_vec = vector.splat %arg11_base : vector<2x2xi32>
+  %arg11_power = arith.constant 30 : i32
+  %arg11_power_vec = vector.splat %arg11_power : vector<2x2xi32>
+  %res11_vec = math.ipowsi %arg11_base_vec, %arg11_power_vec : vector<2x2xi32>
+  %i11 = arith.constant 11 : index
+  %res11 = vector.extract %res11_vec[1, 1] : vector<2x2xi32>
+  memref.store %res11, %result[%i11] : memref<?xi32>
+// CHECK: memref.store %[[cst1073741824]], %[[result]][%[[i11]]] : memref<?xi32>
+
+  return
+}
+
+// CHECK-LABEL: @ipowsi64_fold(
+// CHECK-SAME: %[[result:.+]]: memref<?xi64>
+func.func @ipowsi64_fold(%result : memref<?xi64>) {
+// CHECK-DAG: %[[cst0:.+]] = arith.constant 0 : i64
+// CHECK-DAG: %[[cst1:.+]] = arith.constant 1 : i64
+// CHECK-DAG: %[[cst1073741824:.+]] = arith.constant 1073741824 : i64
+// CHECK-DAG: %[[cst281474976710656:.+]] = arith.constant 281474976710656 : i64
+// CHECK-DAG: %[[cst_m1:.+]] = arith.constant -1 : i64
+// CHECK-DAG: %[[cst_m27:.+]] = arith.constant -27 : i64
+// CHECK-DAG: %[[i0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[i1:.+]] = arith.constant 1 : index
+// CHECK-DAG: %[[i2:.+]] = arith.constant 2 : index
+// CHECK-DAG: %[[i3:.+]] = arith.constant 3 : index
+// CHECK-DAG: %[[i4:.+]] = arith.constant 4 : index
+// CHECK-DAG: %[[i5:.+]] = arith.constant 5 : index
+// CHECK-DAG: %[[i6:.+]] = arith.constant 6 : index
+// CHECK-DAG: %[[i7:.+]] = arith.constant 7 : index
+// CHECK-DAG: %[[i8:.+]] = arith.constant 8 : index
+// CHECK-DAG: %[[i9:.+]] = arith.constant 9 : index
+// CHECK-DAG: %[[i10:.+]] = arith.constant 10 : index
+// CHECK-DAG: %[[i11:.+]] = arith.constant 11 : index
+
+// --- Test power == 0 ---
+  %arg0_base = arith.constant 0 : i64
+  %arg0_power = arith.constant 0 : i64
+  %res0 = math.ipowsi %arg0_base, %arg0_power : i64
+  %i0 = arith.constant 0 : index
+  memref.store %res0, %result[%i0] : memref<?xi64>
+// CHECK: memref.store %[[cst1]], %[[result]][%[[i0]]] : memref<?xi64>
+
+  %arg1_base = arith.constant 10 : i64
+  %arg1_power = arith.constant 0 : i64
+  %res1 = math.ipowsi %arg1_base, %arg1_power : i64
+  %i1 = arith.constant 1 : index
+  memref.store %res1, %result[%i1] : memref<?xi64>
+// CHECK: memref.store %[[cst1]], %[[result]][%[[i1]]] : memref<?xi64>
+
+  %arg2_base = arith.constant -10 : i64
+  %arg2_power = arith.constant 0 : i64
+  %res2 = math.ipowsi %arg2_base, %arg2_power : i64
+  %i2 = arith.constant 2 : index
+  memref.store %res2, %result[%i2] : memref<?xi64>
+// CHECK: memref.store %[[cst1]], %[[result]][%[[i2]]] : memref<?xi64>
+
+// --- Test negative powers ---
+  %arg3_base = arith.constant 0 : i64
+  %arg3_power = arith.constant -1 : i64
+  %res3 = math.ipowsi %arg3_base, %arg3_power : i64
+  %i3 = arith.constant 3 : index
+  memref.store %res3, %result[%i3] : memref<?xi64>
+// No folding for ipowsi(0, x) for x < 0:
+// CHECK: %[[res3:.+]] = math.ipowsi %[[cst0]], %[[cst_m1]] : i64
+// CHECK: memref.store %[[res3]], %[[result]][%[[i3]]] : memref<?xi64>
+
+  %arg4_base = arith.constant 1 : i64
+  %arg4_power = arith.constant -10 : i64
+  %res4 = math.ipowsi %arg4_base, %arg4_power : i64
+  %i4 = arith.constant 4 : index
+  memref.store %res4, %result[%i4] : memref<?xi64>
+// CHECK: memref.store %[[cst1]], %[[result]][%[[i4]]] : memref<?xi64>
+
+  %arg5_base = arith.constant 2 : i64
+  %arg5_power = arith.constant -1 : i64
+  %res5 = math.ipowsi %arg5_base, %arg5_power : i64
+  %i5 = arith.constant 5 : index
+  memref.store %res5, %result[%i5] : memref<?xi64>
+// CHECK: memref.store %[[cst0]], %[[result]][%[[i5]]] : memref<?xi64>
+
+  %arg6_base = arith.constant -2 : i64
+  %arg6_power = arith.constant -1 : i64
+  %res6 = math.ipowsi %arg6_base, %arg6_power : i64
+  %i6 = arith.constant 6 : index
+  memref.store %res6, %result[%i6] : memref<?xi64>
+// CHECK: memref.store %[[cst0]], %[[result]][%[[i6]]] : memref<?xi64>
+
+  %arg7_base = arith.constant -1 : i64
+  %arg7_power = arith.constant -10 : i64
+  %res7 = math.ipowsi %arg7_base, %arg7_power : i64
+  %i7 = arith.constant 7 : index
+  memref.store %res7, %result[%i7] : memref<?xi64>
+// CHECK: memref.store %[[cst1]], %[[result]][%[[i7]]] : memref<?xi64>
+
+  %arg8_base = arith.constant -1 : i64
+  %arg8_power = arith.constant -11 : i64
+  %res8 = math.ipowsi %arg8_base, %arg8_power : i64
+  %i8 = arith.constant 8 : index
+  memref.store %res8, %result[%i8] : memref<?xi64>
+// CHECK: memref.store %[[cst_m1]], %[[result]][%[[i8]]] : memref<?xi64>
+
+// --- Test positive powers ---
+  %arg9_base = arith.constant -3 : i64
+  %arg9_power = arith.constant 3 : i64
+  %res9 = math.ipowsi %arg9_base, %arg9_power : i64
+  %i9 = arith.constant 9 : index
+  memref.store %res9, %result[%i9] : memref<?xi64>
+// CHECK: memref.store %[[cst_m27]], %[[result]][%[[i9]]] : memref<?xi64>
+
+  %arg10_base = arith.constant 2 : i64
+  %arg10_power = arith.constant 30 : i64
+  %res10 = math.ipowsi %arg10_base, %arg10_power : i64
+  %i10 = arith.constant 10 : index
+  memref.store %res10, %result[%i10] : memref<?xi64>
+// CHECK: memref.store %[[cst1073741824]], %[[result]][%[[i10]]] : memref<?xi64>
+
+  %arg11_base = arith.constant 2 : i64
+  %arg11_power = arith.constant 48 : i64
+  %res11 = math.ipowsi %arg11_base, %arg11_power : i64
+  %i11 = arith.constant 11 : index
+  memref.store %res11, %result[%i11] : memref<?xi64>
+// CHECK: memref.store %[[cst281474976710656]], %[[result]][%[[i11]]] : memref<?xi64>
+
+  return
+}
+
+// CHECK-LABEL: @ipowsi16_fold(
+// CHECK-SAME: %[[result:.+]]: memref<?xi16>
+func.func @ipowsi16_fold(%result : memref<?xi16>) {
+// CHECK-DAG: %[[cst0:.+]] = arith.constant 0 : i16
+// CHECK-DAG: %[[cst1:.+]] = arith.constant 1 : i16
+// CHECK-DAG: %[[cst16384:.+]] = arith.constant 16384 : i16
+// CHECK-DAG: %[[cst_m1:.+]] = arith.constant -1 : i16
+// CHECK-DAG: %[[cst_m27:.+]] = arith.constant -27 : i16
+// CHECK-DAG: %[[i0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[i1:.+]] = arith.constant 1 : index
+// CHECK-DAG: %[[i2:.+]] = arith.constant 2 : index
+// CHECK-DAG: %[[i3:.+]] = arith.constant 3 : index
+// CHECK-DAG: %[[i4:.+]] = arith.constant 4 : index
+// CHECK-DAG: %[[i5:.+]] = arith.constant 5 : index
+// CHECK-DAG: %[[i6:.+]] = arith.constant 6 : index
+// CHECK-DAG: %[[i7:.+]] = arith.constant 7 : index
+// CHECK-DAG: %[[i8:.+]] = arith.constant 8 : index
+// CHECK-DAG: %[[i9:.+]] = arith.constant 9 : index
+// CHECK-DAG: %[[i10:.+]] = arith.constant 10 : index
+
+// --- Test power == 0 ---
+  %arg0_base = arith.constant 0 : i16
+  %arg0_power = arith.constant 0 : i16
+  %res0 = math.ipowsi %arg0_base, %arg0_power : i16
+  %i0 = arith.constant 0 : index
+  memref.store %res0, %result[%i0] : memref<?xi16>
+// CHECK: memref.store %[[cst1]], %[[result]][%[[i0]]] : memref<?xi16>
+
+  %arg1_base = arith.constant 10 : i16
+  %arg1_power = arith.constant 0 : i16
+  %res1 = math.ipowsi %arg1_base, %arg1_power : i16
+  %i1 = arith.constant 1 : index
+  memref.store %res1, %result[%i1] : memref<?xi16>
+// CHECK: memref.store %[[cst1]], %[[result]][%[[i1]]] : memref<?xi16>
+
+  %arg2_base = arith.constant -10 : i16
+  %arg2_power = arith.constant 0 : i16
+  %res2 = math.ipowsi %arg2_base, %arg2_power : i16
+  %i2 = arith.constant 2 : index
+  memref.store %res2, %result[%i2] : memref<?xi16>
+// CHECK: memref.store %[[cst1]], %[[result]][%[[i2]]] : memref<?xi16>
+
+// --- Test negative powers ---
+  %arg3_base = arith.constant 0 : i16
+  %arg3_power = arith.constant -1 : i16
+  %res3 = math.ipowsi %arg3_base, %arg3_power : i16
+  %i3 = arith.constant 3 : index
+  memref.store %res3, %result[%i3] : memref<?xi16>
+// No folding for ipowsi(0, x) for x < 0:
+// CHECK: %[[res3:.+]] = math.ipowsi %[[cst0]], %[[cst_m1]] : i16
+// CHECK: memref.store %[[res3]], %[[result]][%[[i3]]] : memref<?xi16>
+
+  %arg4_base = arith.constant 1 : i16
+  %arg4_power = arith.constant -10 : i16
+  %res4 = math.ipowsi %arg4_base, %arg4_power : i16
+  %i4 = arith.constant 4 : index
+  memref.store %res4, %result[%i4] : memref<?xi16>
+// CHECK: memref.store %[[cst1]], %[[result]][%[[i4]]] : memref<?xi16>
+
+  %arg5_base = arith.constant 2 : i16
+  %arg5_power = arith.constant -1 : i16
+  %res5 = math.ipowsi %arg5_base, %arg5_power : i16
+  %i5 = arith.constant 5 : index
+  memref.store %res5, %result[%i5] : memref<?xi16>
+// CHECK: memref.store %[[cst0]], %[[result]][%[[i5]]] : memref<?xi16>
+
+  %arg6_base = arith.constant -2 : i16
+  %arg6_power = arith.constant -1 : i16
+  %res6 = math.ipowsi %arg6_base, %arg6_power : i16
+  %i6 = arith.constant 6 : index
+  memref.store %res6, %result[%i6] : memref<?xi16>
+// CHECK: memref.store %[[cst0]], %[[result]][%[[i6]]] : memref<?xi16>
+
+  %arg7_base = arith.constant -1 : i16
+  %arg7_power = arith.constant -10 : i16
+  %res7 = math.ipowsi %arg7_base, %arg7_power : i16
+  %i7 = arith.constant 7 : index
+  memref.store %res7, %result[%i7] : memref<?xi16>
+// CHECK: memref.store %[[cst1]], %[[result]][%[[i7]]] : memref<?xi16>
+
+  %arg8_base = arith.constant -1 : i16
+  %arg8_power = arith.constant -11 : i16
+  %res8 = math.ipowsi %arg8_base, %arg8_power : i16
+  %i8 = arith.constant 8 : index
+  memref.store %res8, %result[%i8] : memref<?xi16>
+// CHECK: memref.store %[[cst_m1]], %[[result]][%[[i8]]] : memref<?xi16>
+
+// --- Test positive powers ---
+  %arg9_base = arith.constant -3 : i16
+  %arg9_power = arith.constant 3 : i16
+  %res9 = math.ipowsi %arg9_base, %arg9_power : i16
+  %i9 = arith.constant 9 : index
+  memref.store %res9, %result[%i9] : memref<?xi16>
+// CHECK: memref.store %[[cst_m27]], %[[result]][%[[i9]]] : memref<?xi16>
+
+  %arg10_base = arith.constant 2 : i16
+  %arg10_power = arith.constant 14 : i16
+  %res10 = math.ipowsi %arg10_base, %arg10_power : i16
+  %i10 = arith.constant 10 : index
+  memref.store %res10, %result[%i10] : memref<?xi16>
+// CHECK: memref.store %[[cst16384]], %[[result]][%[[i10]]] : memref<?xi16>
+
+  return
+}
+
+// CHECK-LABEL: @ipowsi8_fold(
+// CHECK-SAME: %[[result:.+]]: memref<?xi8>
+func.func @ipowsi8_fold(%result : memref<?xi8>) {
+// CHECK-DAG: %[[cst0:.+]] = arith.constant 0 : i8
+// CHECK-DAG: %[[cst1:.+]] = arith.constant 1 : i8
+// CHECK-DAG: %[[cst64:.+]] = arith.constant 64 : i8
+// CHECK-DAG: %[[cst_m1:.+]] = arith.constant -1 : i8
+// CHECK-DAG: %[[cst_m27:.+]] = arith.constant -27 : i8
+// CHECK-DAG: %[[i0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[i1:.+]] = arith.constant 1 : index
+// CHECK-DAG: %[[i2:.+]] = arith.constant 2 : index
+// CHECK-DAG: %[[i3:.+]] = arith.constant 3 : index
+// CHECK-DAG: %[[i4:.+]] = arith.constant 4 : index
+// CHECK-DAG: %[[i5:.+]] = arith.constant 5 : index
+// CHECK-DAG: %[[i6:.+]] = arith.constant 6 : index
+// CHECK-DAG: %[[i7:.+]] = arith.constant 7 : index
+// CHECK-DAG: %[[i8:.+]] = arith.constant 8 : index
+// CHECK-DAG: %[[i9:.+]] = arith.constant 9 : index
+// CHECK-DAG: %[[i10:.+]] = arith.constant 10 : index
+
+// --- Test power == 0 ---
+  %arg0_base = arith.constant 0 : i8
+  %arg0_power = arith.constant 0 : i8
+  %res0 = math.ipowsi %arg0_base, %arg0_power : i8
+  %i0 = arith.constant 0 : index
+  memref.store %res0, %result[%i0] : memref<?xi8>
+// CHECK: memref.store %[[cst1]], %[[result]][%[[i0]]] : memref<?xi8>
+
+  %arg1_base = arith.constant 10 : i8
+  %arg1_power = arith.constant 0 : i8
+  %res1 = math.ipowsi %arg1_base, %arg1_power : i8
+  %i1 = arith.constant 1 : index
+  memref.store %res1, %result[%i1] : memref<?xi8>
+// CHECK: memref.store %[[cst1]], %[[result]][%[[i1]]] : memref<?xi8>
+
+  %arg2_base = arith.constant -10 : i8
+  %arg2_power = arith.constant 0 : i8
+  %res2 = math.ipowsi %arg2_base, %arg2_power : i8
+  %i2 = arith.constant 2 : index
+  memref.store %res2, %result[%i2] : memref<?xi8>
+// CHECK: memref.store %[[cst1]], %[[result]][%[[i2]]] : memref<?xi8>
+
+// --- Test negative powers ---
+  %arg3_base = arith.constant 0 : i8
+  %arg3_power = arith.constant -1 : i8
+  %res3 = math.ipowsi %arg3_base, %arg3_power : i8
+  %i3 = arith.constant 3 : index
+  memref.store %res3, %result[%i3] : memref<?xi8>
+// No folding for ipowsi(0, x) for x < 0:
+// CHECK: %[[res3:.+]] = math.ipowsi %[[cst0]], %[[cst_m1]] : i8
+// CHECK: memref.store %[[res3]], %[[result]][%[[i3]]] : memref<?xi8>
+
+  %arg4_base = arith.constant 1 : i8
+  %arg4_power = arith.constant -10 : i8
+  %res4 = math.ipowsi %arg4_base, %arg4_power : i8
+  %i4 = arith.constant 4 : index
+  memref.store %res4, %result[%i4] : memref<?xi8>
+// CHECK: memref.store %[[cst1]], %[[result]][%[[i4]]] : memref<?xi8>
+
+  %arg5_base = arith.constant 2 : i8
+  %arg5_power = arith.constant -1 : i8
+  %res5 = math.ipowsi %arg5_base, %arg5_power : i8
+  %i5 = arith.constant 5 : index
+  memref.store %res5, %result[%i5] : memref<?xi8>
+// CHECK: memref.store %[[cst0]], %[[result]][%[[i5]]] : memref<?xi8>
+
+  %arg6_base = arith.constant -2 : i8
+  %arg6_power = arith.constant -1 : i8
+  %res6 = math.ipowsi %arg6_base, %arg6_power : i8
+  %i6 = arith.constant 6 : index
+  memref.store %res6, %result[%i6] : memref<?xi8>
+// CHECK: memref.store %[[cst0]], %[[result]][%[[i6]]] : memref<?xi8>
+
+  %arg7_base = arith.constant -1 : i8
+  %arg7_power = arith.constant -10 : i8
+  %res7 = math.ipowsi %arg7_base, %arg7_power : i8
+  %i7 = arith.constant 7 : index
+  memref.store %res7, %result[%i7] : memref<?xi8>
+// CHECK: memref.store %[[cst1]], %[[result]][%[[i7]]] : memref<?xi8>
+
+  %arg8_base = arith.constant -1 : i8
+  %arg8_power = arith.constant -11 : i8
+  %res8 = math.ipowsi %arg8_base, %arg8_power : i8
+  %i8 = arith.constant 8 : index
+  memref.store %res8, %result[%i8] : memref<?xi8>
+// CHECK: memref.store %[[cst_m1]], %[[result]][%[[i8]]] : memref<?xi8>
+
+// --- Test positive powers ---
+  %arg9_base = arith.constant -3 : i8
+  %arg9_power = arith.constant 3 : i8
+  %res9 = math.ipowsi %arg9_base, %arg9_power : i8
+  %i9 = arith.constant 9 : index
+  memref.store %res9, %result[%i9] : memref<?xi8>
+// CHECK: memref.store %[[cst_m27]], %[[result]][%[[i9]]] : memref<?xi8>
+
+  %arg10_base = arith.constant 2 : i8
+  %arg10_power = arith.constant 6 : i8
+  %res10 = math.ipowsi %arg10_base, %arg10_power : i8
+  %i10 = arith.constant 10 : index
+  memref.store %res10, %result[%i10] : memref<?xi8>
+// CHECK: memref.store %[[cst64]], %[[result]][%[[i10]]] : memref<?xi8>
+
+  return
+}

mlir/test/Dialect/Math/ops.mlir

func.func @round(%f: f32, %v: vector<4xf32>, %t: tensor<4x4x?xf32>) {
   // CHECK: %{{.*}} = math.round %[[F]] : f32
   %0 = math.round %f : f32
   // CHECK: %{{.*}} = math.round %[[V]] : vector<4xf32>
   %1 = math.round %v : vector<4xf32>
   // CHECK: %{{.*}} = math.round %[[T]] : tensor<4x4x?xf32>
   %2 = math.round %t : tensor<4x4x?xf32>
   return
 }
+
+// CHECK-LABEL: func @ipowsi(
+// CHECK-SAME:             %[[I:.*]]: i32, %[[V:.*]]: vector<4xi32>, %[[T:.*]]: tensor<4x4x?xi32>)
+func.func @ipowsi(%i: i32, %v: vector<4xi32>, %t: tensor<4x4x?xi32>) {
+  // CHECK: %{{.*}} = math.ipowsi %[[I]], %[[I]] : i32
+  %0 = math.ipowsi %i, %i : i32
+  // CHECK: %{{.*}} = math.ipowsi %[[V]], %[[V]] : vector<4xi32>
+  %1 = math.ipowsi %v, %v : vector<4xi32>
+  // CHECK: %{{.*}} = math.ipowsi %[[T]], %[[T]] : tensor<4x4x?xi32>
+  %2 = math.ipowsi %t, %t : tensor<4x4x?xi32>
+  return
+}