Add SelectSimplify op to Enzyme Jax.

src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp

@@ -3217,6 +3217,38 @@ struct MinSimplify : public OpRewritePattern<mlir::stablehlo::MinOp> {
   }
 };
 
+struct SelectSimplify : public OpRewritePattern<mlir::stablehlo::SelectOp> {
+  using OpRewritePattern<mlir::stablehlo::SelectOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(mlir::stablehlo::SelectOp op,
+                                PatternRewriter &rewriter) const final {
+    SmallVector<DenseElementsAttr> constants;
+    constants.assign(op->getNumOperands(), DenseElementsAttr());
+    for (unsigned i = 0, e = op->getNumOperands(); i != e; ++i)
+      matchPattern(op->getOperand(i), m_Constant(&constants[i]));
+    if ((bool)constants[0] && (bool)constants[1] && (bool)constants[2]) {
+      if (isa<SplatElementsAttr>(constants[0])) {
+        if (constants[0].getSplatValue<bool>()) {
+          rewriter.replaceOp(op, op->getOperand(1));
+          return success();
+        } else {
+          rewriter.replaceOp(op, op->getOperand(2));
+          return success();
+        }
+      }
+
+      auto out = fromTensor(mlir::stablehlo::selectOp(
+          mlir::stablehlo::constantOp(constants[0]),
+          mlir::stablehlo::constantOp(constants[1]),
+          mlir::stablehlo::constantOp(constants[2]), op.getType()));
+
+      rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(op, op.getType(), out);
+      return success();
+    }
+    return failure();
+  }
+};
+
 struct CosSimplify : public OpRewritePattern<mlir::stablehlo::CosineOp> {
   using OpRewritePattern<mlir::stablehlo::CosineOp>::OpRewritePattern;
 
@@ -6188,7 +6220,7 @@ struct EnzymeHLOOptPass : public EnzymeHLOOptPassBase<EnzymeHLOOptPass> {
     auto context = getOperation()->getContext();
     RewritePatternSet patterns(context);
     patterns
-        .add<AddSimplify, SubSimplify, AndSimplify, MaxSimplify, MinSimplify,
+        .add<AddSimplify, SubSimplify, AndSimplify, MaxSimplify, MinSimplify, SelectSimplify,
              OrSimplify, NegateSimplify, MulSimplify, DivSimplify, RemSimplify,
              PowSimplify, SqrtSimplify, CosSimplify, SinSimplify, NoopSlice,
              SliceSlice, PadSimplify, ShiftRightLogicalSimplify,

src/enzyme_ad/jax/TransformOps/TransformOps.td

@@ -620,6 +620,10 @@ def ApplyConstPropThroughBarrierPatterns : EnzymeHLOPatternOp<
     "const_prop_through_barrier"> {
   let patterns = ["ConstPropThroughBarrier"];
 }
+def ApplySelectSimplifyPatterns : EnzymeHLOPatternOp<
+    "select_simplify"> {
+  let patterns = ["SelectSimplify"];
+}
 
 // TODO: better naming for parameters requires a static interface for
 // constructing them in search.

test/lit_tests/select.mlir

@@ -0,0 +1,24 @@
+// RUN: enzymexlamlir-opt --enzyme-hlo-opt %s | FileCheck %s
+
+func.func @f(%arg3:tensor<10xi32>, %arg4:tensor<10xi32>) -> tensor<10xi32> {
+  %c0 = stablehlo.constant dense<true> : tensor<10xi1>
+  %7 = stablehlo.select %c0, %arg3, %arg4 : tensor<10xi1>, tensor<10xi32>
+  return %7 : tensor<10xi32>
+}
+
+func.func @g() -> tensor<3xi32> {
+  %c0 = stablehlo.constant dense<[true, false, false]> : tensor<3xi1>
+    %c1 = stablehlo.constant dense<[42, 61, 76]> : tensor<3xi32>
+  %c2 = stablehlo.constant dense<[5, 36, 23]> : tensor<3xi32>
+
+  %7 = stablehlo.select %c0, %c1, %c2 : tensor<3xi1>, tensor<3xi32>
+  return %7 : tensor<3xi32>
+}
+
+CHECK:  func.func @f(%arg0: tensor<10xi32>, %arg1: tensor<10xi32>) -> tensor<10xi32>
+CHECK-NEXT   return %arg0 : tensor<10xi32>
+
+CHECK: func.func @g() -> tensor<3xi32> 
+CHECK-NEXT: %c = stablehlo.constant dense<[42, 36, 23]> : tensor<3xi32>
+CHECK-NEXT    return %c : tensor<3xi32>
+