[mlir][math] Added constant folding for IPowI operation.

[vzakhari]

Summary

Split from the original D129809.

[vzakhari]

rebase

[vzakhari]

rebase

Please review.

[Mogball]

mlir/test/Dialect/Math/canonicalize_ipowi.mlir

440

why are these tests so complicated? why do you need memref in here?

[vzakhari]

mlir/test/Dialect/Math/canonicalize_ipowi.mlir

440

Thank you for looking into it, @Mogball!

Basically, I wanted to test multiple folding cases in one function, so I had to make use of the ipowi in some way (so that it does not get removed completely). I made the uses by storing the results into their own cells of the array passed as a pointer to the function. I guess I can just make the function return 11 results and get rid of the store operations. Or I can also split the function into 11 functions. Please let me know which of the options you prefer.

[Mogball]

mlir/test/Dialect/Math/canonicalize_ipowi.mlir

440

The second one.

The folder does look complicated. Do these tests cover each set of cases once?

[vzakhari]

mlir/test/Dialect/Math/canonicalize_ipowi.mlir

440

Yes, I tried to test each line of code at least once. For some of the cases I added more than one test, e.g. for positive powers I test negative and positive base values, whereas the folding code does not care about the base sign.

[Mogball]

mlir/test/Dialect/Math/canonicalize_ipowi.mlir

440

SG. Please split these up and document which case you are testing with each.

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

let description = [{
 
     Example:
 
     ```mlir
     // Scalar signed integer exponentiation.
     %a = math.ipowi %b, %c : i32
     ```
   }];
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // LogOp
 //===----------------------------------------------------------------------===//
 
 def Math_LogOp : Math_FloatUnaryOp<"log"> {
   let summary = "base-e logarithm of the specified value";

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());
   });
 }
 
 //===----------------------------------------------------------------------===//
+// IPowIOp folder
+//===----------------------------------------------------------------------===//
+
+OpFoldResult math::IPowIOp::fold(ArrayRef<Attribute> operands) {
+  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();
+}
+
+//===----------------------------------------------------------------------===//
 // LogOp folder
 //===----------------------------------------------------------------------===//
 
 OpFoldResult math::LogOp::fold(ArrayRef<Attribute> operands) {
   return constFoldUnaryOpConditional<FloatAttr>(
       operands, [](const APFloat &a) -> Optional<APFloat> {
         if (a.isNegative())
           return {};

mlir/test/Dialect/Math/canonicalize_ipowi.mlir

+// RUN: mlir-opt %s -canonicalize | FileCheck %s
+
+// CHECK-LABEL: @ipowi32_fold(
+// CHECK-SAME: %[[result:.+]]: memref<?xi32>
+func.func @ipowi32_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.ipowi %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.ipowi %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.ipowi %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.ipowi %arg3_base, %arg3_power : i32
+  %i3 = arith.constant 3 : index
+  memref.store %res3, %result[%i3] : memref<?xi32>
+// No folding for ipowi(0, x) for x < 0:
+// CHECK: %[[res3:.+]] = math.ipowi %[[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.ipowi %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.ipowi %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.ipowi %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.ipowi %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.ipowi %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.ipowi %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.ipowi %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.ipowi %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: @ipowi64_fold(
+// CHECK-SAME: %[[result:.+]]: memref<?xi64>
+func.func @ipowi64_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.ipowi %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.ipowi %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.ipowi %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.ipowi %arg3_base, %arg3_power : i64
+  %i3 = arith.constant 3 : index
+  memref.store %res3, %result[%i3] : memref<?xi64>
+// No folding for ipowi(0, x) for x < 0:
+// CHECK: %[[res3:.+]] = math.ipowi %[[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.ipowi %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.ipowi %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.ipowi %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.ipowi %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.ipowi %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.ipowi %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.ipowi %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.ipowi %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: @ipowi16_fold(
+// CHECK-SAME: %[[result:.+]]: memref<?xi16>
+func.func @ipowi16_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.ipowi %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.ipowi %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.ipowi %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.ipowi %arg3_base, %arg3_power : i16
+  %i3 = arith.constant 3 : index
+  memref.store %res3, %result[%i3] : memref<?xi16>
+// No folding for ipowi(0, x) for x < 0:
+// CHECK: %[[res3:.+]] = math.ipowi %[[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.ipowi %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.ipowi %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.ipowi %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.ipowi %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.ipowi %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.ipowi %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.ipowi %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: @ipowi8_fold(
+// CHECK-SAME: %[[result:.+]]: memref<?xi8>
+func.func @ipowi8_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.ipowi %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.ipowi %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.ipowi %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.ipowi %arg3_base, %arg3_power : i8
+  %i3 = arith.constant 3 : index
+  memref.store %res3, %result[%i3] : memref<?xi8>
+// No folding for ipowi(0, x) for x < 0:
+// CHECK: %[[res3:.+]] = math.ipowi %[[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.ipowi %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.ipowi %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.ipowi %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.ipowi %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.ipowi %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.ipowi %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.ipowi %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
+}