Repository: EnzymeAD/Enzyme-JAX (language: python)

Issue:
PowSimplify leads to segfault on integer overflow
Pre-optimize IR

```mlir
module {
  func.func @main(%arg0: tensor<2x2xf64>) -> (tensor<i64>, tensor<2x2xf64>) {
    %0 = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<2x2xf64>) -> tensor<2x2xf64>
    %c = stablehlo.constant dense<2> : tensor<i64>
    %c_0 = stablehlo.constant dense<64> : tensor<i64>
    %1 = stablehlo.power %c, %c_0 : tensor<i64>
    %2 = stablehlo.transpose %1, dims = [] : (tensor<i64>) -> tensor<i64>
    %3 = stablehlo.transpose %0, dims = [1, 0] : (tensor<2x2xf64>) -> tensor<2x2xf64>
    return %2, %3 : tensor<i64>, tensor<2x2xf64>
  }
}
```

xref https://github.com/EnzymeAD/Reactant.jl/issues/368


Discussion / hints:
can you run this through a debug build of enzymexlamlir-opt and see what the backtrace is?
I have limited compute access (only my laptop) till the end of Dec to build enzyme-jax locally 😓.
It looks like this call is failing due to pow being applied to integers:

https://github.com/EnzymeAD/Enzyme-JAX/blob/f6587e37ff7298f2a1a273b08c24d69fca7ff30f/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp#L2932-L2933
From the debug build

```
enzymexlamlir-opt: external/llvm-project/mlir/include/mlir/IR/BuiltinAttributes.h:423: auto mlir::DenseElementsAttr::getValues() const [with T = llvm::APFloat]: Assertion `succeeded(range) && "element type cannot be iterated"' failed.
PLEASE submit a bug report to https://github.com/llvm/llvm-project/issues/ and include the crash backtrace.
Stack dump:
0.	Program arguments: ./bazel-bin/enzymexlamlir-opt --enzyme-hlo-opt envs/pow.mlir
Stack dump without symbol names (ensure you have llvm-symbolizer in your PATH or set the environment var `LLVM_SYMBOLIZER_PATH` to point to it):
0  enzymexlamlir-opt 0x000055ba158fefce
1  enzymexlamlir-opt 0x000055ba158ff383
2  enzymexlamlir-opt 0x000055ba158fca84
3  enzymexlamlir-opt 0x000055ba158feaf3
4  libc.so.6         0x00007edd1444c1d0
5  libc.so.6         0x00007edd144a53f4
6  libc.so.6         0x00007edd1444c120 gsignal + 32
7  libc.so.6         0x00007edd144334c3 abort + 223
8  libc.so.6         0x00007edd144333df
9  libc.so.6         0x00007edd14444177
10 enzymexlamlir-opt 0x000055ba1207cc85
11 enzymexlamlir-opt 0x000055ba1207cd2f
12 enzymexlamlir-opt 0x000055ba1207cda5
13 enzymexlamlir-opt 0x000055ba11fe912c
14 enzymexlamlir-opt 0x000055ba11fda240
15 enzymexlamlir-opt 0x000055ba11fc4a20
16 enzymexlamlir-opt 0x000055ba120edc76
17 enzymexlamlir-opt 0x000055ba14ff1feb
18 enzymexlamlir-opt 0x000055ba14ff2a42
19 enzymexlamlir-opt 0x000055ba119404d8
20 enzymexlamlir-opt 0x000055ba14ff5d27
21 enzymexlamlir-opt 0x000055ba14ff2769
22 enzymexlamlir-opt 0x000055ba14fe3495
23 enzymexlamlir-opt 0x000055ba14fe4631
24 enzymexlamlir-opt 0x000055ba14fe594b
25 enzymexlamlir-opt 0x000055ba119404d8
26 enzymexlamlir-opt 0x000055ba14fe56cf
27 enzymexlamlir-opt 0x000055ba14fe48fc
28 enzymexlamlir-opt 0x000055ba14fe4ab6
29 enzymexlamlir-opt 0x000055ba11e0a548
30 enzymexlamlir-opt 0x000055ba11fd7f6a
31 enzymexlamlir-opt 0x000055ba154941f9
32 enzymexlamlir-opt 0x000055ba15497eda
33 enzymexlamlir-opt 0x000055ba119404d8
34 enzymexlamlir-opt 0x000055ba1549d95f
35 enzymexlamlir-opt 0x000055ba15494617
36 enzymexlamlir-opt 0x000055ba15494959
37 enzymexlamlir-opt 0x000055ba1549680a
38 enzymexlamlir-opt 0x000055ba154965f9
39 enzymexlamlir-opt 0x000055ba119239fb
40 enzymexlamlir-opt 0x000055ba119241dd
41 enzymexlamlir-opt 0x000055ba119248d0
42 enzymexlamlir-opt 0x000055ba11925ae7
43 enzymexlamlir-opt 0x000055ba157eaa42
44 enzymexlamlir-opt 0x000055ba157ea198
45 enzymexlamlir-opt 0x000055ba11924a5e
46 enzymexlamlir-opt 0x000055ba11924d6c
47 enzymexlamlir-opt 0x000055ba11924ffd
48 enzymexlamlir-opt 0x000055ba118d7642
49 libc.so.6         0x00007edd14434e08
50 libc.so.6         0x00007edd14434ecc __libc_start_main + 140
51 enzymexlamlir-opt 0x000055ba118d70b5
[1]    56060 IOT instruction (core dumped)  ./bazel-bin/enzymexlamlir-opt --enzyme-hlo-opt envs/pow.mlir
```

(I am not sure why there is the llvm-symbolizer warning. I definitely have it in my path)
> It looks like this call is failing due to pow being applied to integers:
> 
> [Enzyme-JAX/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp](https://github.com/EnzymeAD/Enzyme-JAX/blob/f6587e37ff7298f2a1a273b08c24d69fca7ff30f/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp#L2932-L2933)
> 
> Lines 2932 to 2933 in [f6587e3](/EnzymeAD/Enzyme-JAX/commit/f6587e37ff7298f2a1a273b08c24d69fca7ff30f)
> 
>  if (auto res = constFoldBinaryOpConditional<FloatAttr, FloatAttr::ValueType, 
>                                              void>(

so we just need to add a case for integers right?

Fix patch:
diff --git a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp
index 642c4a35..5db1d172 100644
--- a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp
+++ b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp
@@ -2929,40 +2929,68 @@ struct PowSimplify : public OpRewritePattern<mlir::stablehlo::PowOp> {
     for (unsigned i = 0, e = op->getNumOperands(); i != e; ++i)
       matchPattern(op->getOperand(i), m_Constant(&constants[i]));
 
-    if (auto res = constFoldBinaryOpConditional<FloatAttr, FloatAttr::ValueType,
-                                                void>(
-            constants,
-            [](const APFloat &a, const APFloat &b) -> std::optional<APFloat> {
-              if (a.getSizeInBits(a.getSemantics()) == 64 &&
-                  b.getSizeInBits(b.getSemantics()) == 64)
-                return APFloat(pow(a.convertToDouble(), b.convertToDouble()));
-
-              if (a.getSizeInBits(a.getSemantics()) == 32 &&
-                  b.getSizeInBits(b.getSemantics()) == 32)
-                return APFloat(powf(a.convertToFloat(), b.convertToFloat()));
-
-              return {};
-            })) {
-      rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
-          op, op.getType(), res.cast<ElementsAttr>());
-      return success();
-    }
+    if (op.getType().getElementType().isa<FloatType>()) {
+      if (auto res = constFoldBinaryOpConditional<FloatAttr,
+                                                  FloatAttr::ValueType, void>(
+              constants,
+              [](const APFloat &a, const APFloat &b) -> std::optional<APFloat> {
+                if (a.getSizeInBits(a.getSemantics()) == 64 &&
+                    b.getSizeInBits(b.getSemantics()) == 64)
+                  return APFloat(pow(a.convertToDouble(), b.convertToDouble()));
 
-    // pow(X, 0.5) -> sqrt(X)
-    {
-      DenseFPElementsAttr rhs;
-      if (matchPattern(op.getRhs(), m_Constant(&rhs))) {
-        bool allHalf = true;
-        for (auto v : rhs) {
-          if (!v.isExactlyValue(0.5)) {
-            allHalf = false;
-            break;
+                if (a.getSizeInBits(a.getSemantics()) == 32 &&
+                    b.getSizeInBits(b.getSemantics()) == 32)
+                  return APFloat(powf(a.convertToFloat(), b.convertToFloat()));
+
+                return {};
+              })) {
+        rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
+            op, op.getType(), res.cast<ElementsAttr>());
+        return success();
+      }
+
+      // pow(X, 0.5) -> sqrt(X)
+      {
+        DenseFPElementsAttr rhs;
+        if (matchPattern(op.getRhs(), m_Constant(&rhs))) {
+          bool allHalf = true;
+          for (auto v : rhs) {
+            if (!v.isExactlyValue(0.5)) {
+              allHalf = false;
+              break;
+            }
+          }
+          if (allHalf) {
+            rewriter.replaceOpWithNewOp<stablehlo::SqrtOp>(op, op.getLhs());
+            return success();
           }
         }
-        if (allHalf) {
-          rewriter.replaceOpWithNewOp<stablehlo::SqrtOp>(op, op.getLhs());
-          return success();
-        }
+      }
+    } else {
+      if (auto res = constFoldBinaryOpConditional<IntegerAttr,
+                                                  IntegerAttr::ValueType, void>(
+              constants,
+              [](const APInt &a, const APInt &b) -> std::optional<APInt> {
+                if (b.isNegative())
+                  return {}; // Ignore the negative case
+
+                APInt result = APInt(a.getBitWidth(), 1);
+                APInt base = a;
+                uint64_t exponent = b.getLimitedValue();
+
+                while (exponent > 0) {
+                  if (exponent % 2 == 1) {
+                    result *= base;
+                  }
+                  base *= base;
+                  exponent /= 2;
+                }
+
+                return result;
+              })) {
+        rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
+            op, op.getType(), res.cast<ElementsAttr>());
+        return success();
       }
     }
 


Test patch:
diff --git a/test/lit_tests/powint.mlir b/test/lit_tests/powint.mlir
new file mode 100644
index 00000000..c8066cb7
--- /dev/null
+++ b/test/lit_tests/powint.mlir
@@ -0,0 +1,45 @@
+// RUN: enzymexlamlir-opt --enzyme-hlo-opt %s | FileCheck %s
+
+module {
+  func.func @main() -> (tensor<i64>) {
+    %c = stablehlo.constant dense<2> : tensor<i64>
+    %c_0 = stablehlo.constant dense<4> : tensor<i64>
+    %1 = stablehlo.power %c, %c_0 : tensor<i64>
+    return %1 : tensor<i64>
+  }
+}
+
+// CHECK:  func.func @main() -> tensor<i64> {
+// CHECK-NEXT:    %c = stablehlo.constant dense<16> : tensor<i64>
+// CHECK-NEXT:    return %c : tensor<i64>
+// CHECK-NEXT:  }
+
+module {
+  func.func @main() -> (tensor<i64>) {
+    %c = stablehlo.constant dense<6> : tensor<i64>
+    %c_0 = stablehlo.constant dense<0> : tensor<i64>
+    %1 = stablehlo.power %c, %c_0 : tensor<i64>
+    return %1 : tensor<i64>
+  }
+}
+
+// CHECK:  func.func @main() -> tensor<i64> {
+// CHECK-NEXT:    %c = stablehlo.constant dense<1> : tensor<i64>
+// CHECK-NEXT:    return %c : tensor<i64>
+// CHECK-NEXT:  }
+
+module {
+  func.func @main() -> (tensor<i64>) {
+    %c = stablehlo.constant dense<6> : tensor<i64>
+    %c_0 = stablehlo.constant dense<-1> : tensor<i64>
+    %1 = stablehlo.power %c, %c_0 : tensor<i64>
+    return %1 : tensor<i64>
+  }
+}
+
+// CHECK:  func.func @main() -> tensor<i64> {
+// CHECK-NEXT:    %c = stablehlo.constant dense<6> : tensor<i64>
+// CHECK-NEXT:    %c_0 = stablehlo.constant dense<-1> : tensor<i64>
+// CHECK-NEXT:    %0 = stablehlo.power %c, %c_0 : tensor<i64>
+// CHECK-NEXT:    return %0 : tensor<i64>
+// CHECK-NEXT:  }