const prop for log and logp1

src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp

@@ -2209,6 +2209,57 @@ struct DynamicUpdateSliceConstProp final
   }
 };
 
+struct LogConstProp final : OpRewritePattern<mlir::stablehlo::LogOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(mlir::stablehlo::LogOp op,
+                                PatternRewriter &rewriter) const override {
+    // return if not constant
+    DenseElementsAttr inputAttr;
+    if (!matchPattern(op.getOperand(), m_Constant(&inputAttr)))
+      return failure();
+
+    // get the resultType
+    auto resultType = op.getType().cast<ShapedType>();
+
+    // Convert constant to tensor, compute log, then convert back to attribute
+    stablehlo::Tensor inputTen = mlir::stablehlo::constantOp(inputAttr);
+    auto out = mlir::stablehlo::logOp(inputTen, resultType);
+
+    // Replace with new constant op containing the computed result
+    rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(op, resultType,
+                                                       fromTensor(out));
+
+    return success();
+  }
+};
+
+struct LogPlusConstProp final : OpRewritePattern<mlir::stablehlo::Log1pOp> {
+
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(mlir::stablehlo::Log1pOp op,
+                                PatternRewriter &rewriter) const override {
+    // return if not constant
+    DenseElementsAttr inputAttr;
+    if (!matchPattern(op.getOperand(), m_Constant(&inputAttr)))
+      return failure();
+
+    // get the resultType
+    auto resultType = op.getType().cast<ShapedType>();
+
+    // Convert constant to tensor, compute log, then convert back to attribute
+    stablehlo::Tensor inputTen = mlir::stablehlo::constantOp(inputAttr);
+    auto out = mlir::stablehlo::log1pOp(inputTen, resultType);
+
+    // Replace with new constant op containing the computed result
+    rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(op, resultType,
+                                                       fromTensor(out));
+
+    return success();
+  }
+};
+
 struct ConcatConstProp final
     : OpRewritePattern<mlir::stablehlo::ConcatenateOp> {
   using OpRewritePattern::OpRewritePattern;
@@ -6963,8 +7014,9 @@ struct EnzymeHLOOptPass : public EnzymeHLOOptPassBase<EnzymeHLOOptPass> {
     patterns.add<ConvertConcat, DynamicUpdateToConcat, SliceOfDynamicUpdate,
                  SliceElementwise, SliceReshapeElementwise, SlicePad,
                  SliceReshapePad, DotReshapeDot, ConcatConstProp,
-                 DynamicUpdateSliceConstProp, ConcatFuse, ConcatToBroadcast,
-                 PadPad, PadReshapePad, ConcatPushBinop<stablehlo::AddOp>,
+                 DynamicUpdateSliceConstProp, LogConstProp, LogPlusConstProp,
+                 ConcatFuse, ConcatToBroadcast, PadPad, PadReshapePad,
+                 ConcatPushBinop<stablehlo::AddOp>,
                  ConcatPushBinop<stablehlo::MulOp>, ScatterToDynamicUpdateSlice,
                  ReduceConcat, ConcatSlice, SliceConcat, SliceReshapeConcat,
                  BinBroadcastSplat<stablehlo::AddOp>,

src/enzyme_ad/jax/TransformOps/TransformOps.td

@@ -153,6 +153,15 @@ def ApplyGatherSimplifyPatterns : EnzymeHLOPatternOp<
     "gather_simplify"> {
   let patterns = ["GatherSimplify"];
 }
+def ApplyLogConstProp : EnzymeHLOPatternOp<
+    "log_const_prop">{
+  let patterns = ["LogConstProp"];
+}
+
+def ApplyLogPlusConstProp : EnzymeHLOPatternOp<
+    "log_plus_one_const_prop">{
+  let patterns = ["LogPlusConstProp"];
+}
 
 // regular benefit
 def ApplyConvertConcatPatterns : EnzymeHLOPatternOp<

src/enzyme_ad/jax/primitives.py

@@ -172,6 +172,8 @@ def hlo_opts():
 dot_reshape_dot<1>;
 concat_const_prop<1>;
 dynamic_update_slice_const_prop<1>;
+log_const_prop<1>;
+log_plus_one_const_prop<1>;
 concat_fuse<1>;
 pad_reshape_pad<1>;
 pad_pad<1>;

test/lit_tests/log_const_prop.mlir

@@ -0,0 +1,30 @@
+// RUN: enzymexlamlir-opt %s --pass-pipeline="builtin.module(enzyme-hlo-opt)" | FileCheck %s
+
+module {
+  func.func @log_f32() -> tensor<2xf32> {  
+    %arg = stablehlo.constant dense<[1.000000e+00,2.000000e+00]> : tensor<2xf32>
+    %result = stablehlo.log %arg : tensor<2xf32> 
+    func.return %result : tensor<2xf32>
+  }
+
+
+  func.func @log_plus_one_op_test_f64() -> tensor<5xf64> {
+    %operand = stablehlo.constant dense<[0.0, -0.999, 7.0, 6.38905621, 15.0]> : tensor<5xf64>
+    %result = stablehlo.log_plus_one %operand : tensor<5xf64>
+    func.return %result : tensor<5xf64>
+  }
+
+
+}
+
+
+// CHECK: func.func @log_f32() -> tensor<2xf32> {
+// CHECK-NEXT: %cst = stablehlo.constant dense<[0.000000e+00, 0.693147182]> : tensor<2xf32>
+// CHECK-NEXT: return %cst : tensor<2xf32>
+// CHECK-NEXT: }
+
+
+// CHECK: func.func @log_plus_one_op_test_f64() -> tensor<5xf64> {
+// CHECK-NEXT: %cst = stablehlo.constant dense<[0.000000e+00, -6.9077552789821359, 2.0794415416798357, 2.0000000150316017, 2.7725887222397811]> : tensor<5xf64>
+// CHECK-NEXT: return %cst : tensor<5xf64>
+// CHECK-NEXT: }