Separate out optimization

Author: vimarsh6739

enzyme/Enzyme/MLIR/Passes/CMakeLists.txt

@@ -11,6 +11,7 @@ add_mlir_dialect_library(MLIREnzymeTransforms
     EnzymeBatchToTensorPass.cpp
     EnzymeWrapPass.cpp
     InlineEnzymeRegions.cpp
+    HoistEnzymeRegions.cpp
     LowerLLVMExtPass.cpp
     PrintActivityAnalysis.cpp
     PrintAliasAnalysis.cpp

enzyme/Enzyme/MLIR/Passes/HoistEnzymeRegions.cpp

@@ -0,0 +1,55 @@
+//===- HoistEnzymeRegions.cpp -LICM for  enzyme.autodiff_region ----------=== //
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements passes to hoist computations within autodiff_region ops
+// to the caller
+//
+//===----------------------------------------------------------------------===//
+
+#include "Dialect/Ops.h"
+#include "Interfaces/AutoDiffOpInterface.h"
+#include "Passes/Passes.h"
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Interfaces/FunctionInterfaces.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "mlir/Transforms/RegionUtils.h"
+#include "llvm/ADT/TypeSwitch.h"
+
+using namespace mlir;
+
+namespace mlir {
+namespace enzyme {
+#define GEN_PASS_DEF_HOISTENZYMEFROMREGIONPASS
+#include "Passes/Passes.h.inc"
+} // namespace enzyme
+} // namespace mlir
+
+namespace {
+
+struct HoistEnzymeAutoDiff : public OpRewritePattern<enzyme::AutoDiffRegionOp> {
+  using OpRewritePattern<enzyme::AutoDiffRegionOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(enzyme::AutoDiffRegionOp op,
+                                PatternRewriter &rewriter) const override {
+    
+    return success();
+  }
+};
+
+struct HoistEnzymeFromRegion
+    : public enzyme::impl::HoistEnzymeFromRegionPassBase<
+          HoistEnzymeFromRegion> {
+  void runOnOperation() override {
+    RewritePatternSet patterns(&getContext());
+    patterns.add<HoistEnzymeAutoDiff>(&getContext());
+    GreedyRewriteConfig config;
+    (void)applyPatternsGreedily(getOperation(), std::move(patterns), config);
+  }
+};
+} // namespace

enzyme/Enzyme/MLIR/Passes/InlineEnzymeRegions.cpp

@@ -27,7 +27,6 @@ namespace mlir {
 namespace enzyme {
 #define GEN_PASS_DEF_INLINEENZYMEINTOREGIONPASS
 #define GEN_PASS_DEF_OUTLINEENZYMEFROMREGIONPASS
-#define GEN_PASS_DEF_HOISTENZYMEFROMREGIONPASS
 #include "Passes/Passes.h.inc"
 } // namespace enzyme
 } // namespace mlir
@@ -262,16 +261,6 @@ struct OutlineEnzymeFromRegion
   }
 };
 
-struct HoistEnzymeFromRegion
-    : public enzyme::impl::HoistEnzymeFromRegionPassBase<
-          HoistEnzymeFromRegion> {
-  void runOnOperation() override {
-    RewritePatternSet patterns(&getContext());
-
-    GreedyRewriteConfig config;
-    (void)applyPatternsGreedily(getOperation(), std::move(patterns), config);
-  }
-};
 } // namespace
 
 bool mlir::enzyme::inlineAutodiffOp(enzyme::AutoDiffOp &op,