From 15ae277e711fb97340a17395e8b4c807130a9c8b Mon Sep 17 00:00:00 2001
From: "William S. Moses" <gh@wsmoses.com>
Date: Sun, 10 Mar 2024 04:44:09 -0400
Subject: [PATCH] Fix constprop on integers

---
 src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp | 151 +++++++++++++++-------
 1 file changed, 103 insertions(+), 48 deletions(-)

diff --git a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp
index c1c4c3748..9e11fc5f2 100644
--- a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp
+++ b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp
@@ -857,18 +857,32 @@ struct AddSimplify : public OpRewritePattern<mlir::stablehlo::AddOp> {
     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> {
-                  APFloat res2(a);
-                  res2.add(b, llvm::RoundingMode::NearestTiesToEven);
-                  return res2;
-                })) {
-      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> {
+                APFloat res2(a);
+                res2.add(b, llvm::RoundingMode::NearestTiesToEven);
+                return res2;
+              })) {
+        rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
+            op, op.getType(), res.cast<ElementsAttr>());
+        return success();
+      }
+    } else {
+      if (auto res = constFoldBinaryOpConditional<IntegerAttr,
+                                                  IntegerAttr::ValueType, void>(
+              constants,
+              [](const APInt &a, const APInt &b) -> std::optional<APInt> {
+                APInt res2(a);
+                res2 += b;
+                return res2;
+              })) {
+        rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
+            op, op.getType(), res.cast<ElementsAttr>());
+        return success();
+      }
     }
 
     return failure();
@@ -896,18 +910,32 @@ struct SubSimplify : public OpRewritePattern<mlir::stablehlo::SubtractOp> {
     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> {
-                  APFloat res2(a);
-                  res2.subtract(b, llvm::RoundingMode::NearestTiesToEven);
-                  return res2;
-                })) {
-      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> {
+                APFloat res2(a);
+                res2.subtract(b, llvm::RoundingMode::NearestTiesToEven);
+                return res2;
+              })) {
+        rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
+            op, op.getType(), res.cast<ElementsAttr>());
+        return success();
+      }
+    } else {
+      if (auto res = constFoldBinaryOpConditional<IntegerAttr,
+                                                  IntegerAttr::ValueType, void>(
+              constants,
+              [](const APInt &a, const APInt &b) -> std::optional<APInt> {
+                APInt res2(a);
+                res2 -= b;
+                return res2;
+              })) {
+        rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
+            op, op.getType(), res.cast<ElementsAttr>());
+        return success();
+      }
     }
 
     return failure();
@@ -1026,18 +1054,32 @@ struct MulSimplify : public OpRewritePattern<mlir::stablehlo::MulOp> {
     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> {
-                  APFloat res2(a);
-                  res2.multiply(b, llvm::RoundingMode::NearestTiesToEven);
-                  return res2;
-                })) {
-      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> {
+                APFloat res2(a);
+                res2.multiply(b, llvm::RoundingMode::NearestTiesToEven);
+                return res2;
+              })) {
+        rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
+            op, op.getType(), res.cast<ElementsAttr>());
+        return success();
+      }
+    } else {
+      if (auto res = constFoldBinaryOpConditional<IntegerAttr,
+                                                  IntegerAttr::ValueType, void>(
+              constants,
+              [](const APInt &a, const APInt &b) -> std::optional<APInt> {
+                APInt res2(a);
+                res2 *= b;
+                return res2;
+              })) {
+        rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
+            op, op.getType(), res.cast<ElementsAttr>());
+        return success();
+      }
     }
 
     return failure();
@@ -1067,18 +1109,31 @@ struct DivSimplify : public OpRewritePattern<mlir::stablehlo::DivOp> {
     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> {
-                  APFloat res2(a);
-                  res2.divide(b, llvm::RoundingMode::NearestTiesToEven);
-                  return res2;
-                })) {
-      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> {
+                APFloat res2(a);
+                res2.divide(b, llvm::RoundingMode::NearestTiesToEven);
+                return res2;
+              })) {
+        rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
+            op, op.getType(), res.cast<ElementsAttr>());
+        return success();
+      }
+    } else {
+      if (auto res = constFoldBinaryOpConditional<IntegerAttr,
+                                                  IntegerAttr::ValueType, void>(
+              constants,
+              [](const APInt &a, const APInt &b) -> std::optional<APInt> {
+                APInt res2(a);
+                return res2.sdiv(b);
+              })) {
+        rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
+            op, op.getType(), res.cast<ElementsAttr>());
+        return success();
+      }
     }
 
     return failure();