bugfix else (#46)

EnzymeAD · Mar 10, 2024 · e652346 · e652346
1 parent 57ddcc2
commit e652346
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 3 deletions.
diff --git a/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp b/src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp
@@ -26,6 +26,7 @@
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 #include "mlir/Dialect/CommonFolders.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
 
 #define DEBUG_TYPE "enzyme"
 
@@ -156,8 +157,7 @@ struct SlicePad final : OpRewritePattern<mlir::stablehlo::SliceOp> {
           op.getLoc(), pad.getOperand(), start, end, step);
       rewriter.replaceOpWithNewOp<stablehlo::PadOp>(
           op, nslice, pad.getPaddingValue(), lpads, hpads, interiors);
-    }
-    {
+    } else {
       rewriter.replaceOpWithNewOp<stablehlo::SliceOp>(op, pad.getOperand(),
                                                       start, end, step);
     }

diff --git a/src/enzyme_ad/jax/Passes/Passes.h b/src/enzyme_ad/jax/Passes/Passes.h
@@ -32,6 +32,10 @@ namespace mhlo {
 class MhloDialect;
 } // end namespace mhlo
 
+namespace tensor {
+class TensorDialect;
+} // namespace tensor
+
 namespace stablehlo {
 class StablehloDialect;
 } // namespace stablehlo

diff --git a/src/enzyme_ad/jax/Passes/Passes.td b/src/enzyme_ad/jax/Passes/Passes.td
@@ -33,7 +33,8 @@ def ArithRaisingPass : Pass<"arith-raise"> {
 def EnzymeHLOOptPass : Pass<"enzyme-hlo-opt"> {
   let summary = "Optimize stablehlo";
   let dependentDialects = [
-    "stablehlo::StablehloDialect"
+    "stablehlo::StablehloDialect",
+    "tensor::TensorDialect"
   ];
   let constructor = "mlir::enzyme::createEnzymeHLOOptPass()";
 }