Pow to sqrt

EnzymeAD · Mar 10, 2024 · 91e6faa · 91e6faa
1 parent 650d115
commit 91e6faa
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 0 deletions.
diff --git a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp
@@ -1330,6 +1330,24 @@ struct PowSimplify : public OpRewritePattern<mlir::stablehlo::PowOp> {
       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();
+        }
+      }
+    }
+
     return failure();
   }
 };

diff --git a/test/lit_tests/powtosqrt.mlir b/test/lit_tests/powtosqrt.mlir
@@ -0,0 +1,15 @@
+// RUN: enzymexlamlir-opt --enzyme-hlo-opt %s | FileCheck %s
+
+module {
+
+  func.func @main(%a : tensor<2x2xf32>) -> tensor<2x2xf32> {
+    %half = stablehlo.constant dense<5.000000e-01> : tensor<2x2xf32>
+    %pd = stablehlo.power %a, %half : tensor<2x2xf32>
+    return %pd : tensor<2x2xf32>
+  }
+}
+
+// CHECK:  func.func @main(%arg0: tensor<2x2xf32>) -> tensor<2x2xf32> {
+// CHECK-NEXT:    %0 = stablehlo.sqrt %arg0 : tensor<2x2xf32>
+// CHECK-NEXT:    return %0 : tensor<2x2xf32>
+// CHECK-NEXT:  }