MLIR: post optimization pipeline

enzyme/.bazelversion

@@ -0,0 +1 @@
+6.5.0

enzyme/Enzyme/MLIR/Implementations/FuncAutoDiffOpInterfaceImpl.cpp

@@ -73,7 +73,8 @@ class AutoDiffCallFwd
         fn, RetActivity, ArgActivity, gutils->TA, returnPrimal, mode,
         freeMemory, width,
         /* addedType */ nullptr, type_args, volatile_args,
-        /* augmented */ nullptr);
+        /* augmented */ nullptr,
+        gutils->postpasses);
 
     SmallVector<Value> fwdArguments;
 
@@ -173,7 +174,7 @@ class AutoDiffCallRev
     auto revFn = gutils->Logic.CreateReverseDiff(
         fn, RetActivity, ArgActivity, gutils->TA, returnPrimal, returnShadow,
         mode, freeMemory, width, /*addedType*/ nullptr, type_args,
-        volatile_args, /*augmented*/ nullptr);
+        volatile_args, /*augmented*/ nullptr, gutils->postpasses);
 
     SmallVector<Value> revArguments;
 

enzyme/Enzyme/MLIR/Interfaces/EnzymeLogic.cpp

@@ -78,7 +78,7 @@ FunctionOpInterface mlir::enzyme::MEnzymeLogic::CreateForwardDiff(
     std::vector<DIFFE_TYPE> ArgActivity, MTypeAnalysis &TA,
     std::vector<bool> returnPrimals, DerivativeMode mode, bool freeMemory,
     size_t width, mlir::Type addedType, MFnTypeInfo type_args,
-    std::vector<bool> volatile_args, void *augmented) {
+    std::vector<bool> volatile_args, void *augmented, llvm::StringRef postpasses) {
   if (fn.getFunctionBody().empty()) {
     llvm::errs() << fn << "\n";
     llvm_unreachable("Differentiating empty function");
@@ -105,7 +105,7 @@ FunctionOpInterface mlir::enzyme::MEnzymeLogic::CreateForwardDiff(
   auto gutils = MDiffeGradientUtils::CreateFromClone(
       *this, mode, width, fn, TA, type_args, returnPrimalsP, returnShadowsP,
       RetActivity, ArgActivity, addedType,
-      /*omp*/ false);
+      /*omp*/ false, postpasses);
   ForwardCachedFunctions[tup] = gutils->newFunc;
 
   insert_or_assign2<MForwardCacheKey, FunctionOpInterface>(
@@ -195,10 +195,20 @@ FunctionOpInterface mlir::enzyme::MEnzymeLogic::CreateForwardDiff(
   if (!valid)
     return nullptr;
 
-  // if (PostOpt)
-  //  PPC.optimizeIntermediate(nf);
-  // if (EnzymePrint) {
-  //  llvm::errs() << nf << "\n";
-  //}
+  if (postpasses != "") {
+    mlir::PassManager pm(nf->getContext());
+    std::string error_message;
+    //llvm::raw_string_ostream error_stream(error_message);
+    mlir::LogicalResult result =
+        mlir::parsePassPipeline(postpasses, pm);
+    if (mlir::failed(result)) {
+      return nullptr;
+    }
+
+    if (!mlir::succeeded(pm.run(nf))) {
+      return nullptr;
+    }
+  }
+
   return nf;
 }

enzyme/Enzyme/MLIR/Interfaces/EnzymeLogic.h

@@ -196,14 +196,16 @@ class MEnzymeLogic {
                     std::vector<bool> returnPrimals, DerivativeMode mode,
                     bool freeMemory, size_t width, mlir::Type addedType,
                     MFnTypeInfo type_args, std::vector<bool> volatile_args,
-                    void *augmented);
+                    void *augmented, llvm::StringRef postpasses);
 
   FunctionOpInterface CreateReverseDiff(
       FunctionOpInterface fn, std::vector<DIFFE_TYPE> retType,
       std::vector<DIFFE_TYPE> constants, MTypeAnalysis &TA,
       std::vector<bool> returnPrimals, std::vector<bool> returnShadows,
       DerivativeMode mode, bool freeMemory, size_t width, mlir::Type addedType,
-      MFnTypeInfo type_args, std::vector<bool> volatile_args, void *augmented);
+      MFnTypeInfo type_args, std::vector<bool> volatile_args, void *augmented,
+      llvm::StringRef postpasses);
+
   void
   initializeShadowValues(SmallVector<mlir::Block *> &dominatorToposortBlocks,
                          MGradientUtilsReverse *gutils);

enzyme/Enzyme/MLIR/Interfaces/EnzymeLogicReverse.cpp

@@ -182,7 +182,8 @@ FunctionOpInterface MEnzymeLogic::CreateReverseDiff(
     std::vector<DIFFE_TYPE> constants, MTypeAnalysis &TA,
     std::vector<bool> returnPrimals, std::vector<bool> returnShadows,
     DerivativeMode mode, bool freeMemory, size_t width, mlir::Type addedType,
-    MFnTypeInfo type_args, std::vector<bool> volatile_args, void *augmented) {
+    MFnTypeInfo type_args, std::vector<bool> volatile_args, void *augmented,
+    llvm::StringRef postpasses) {
 
   if (fn.getFunctionBody().empty()) {
     llvm::errs() << fn << "\n";
@@ -214,7 +215,7 @@ FunctionOpInterface MEnzymeLogic::CreateReverseDiff(
 
   MGradientUtilsReverse *gutils = MGradientUtilsReverse::CreateFromClone(
       *this, mode, width, fn, TA, type_args, returnPrimalsP, returnShadowsP,
-      retType, constants, addedType);
+      retType, constants, addedType, postpasses);
 
   ReverseCachedFunctions[tup] = gutils->newFunc;
 
@@ -254,5 +255,20 @@ FunctionOpInterface MEnzymeLogic::CreateReverseDiff(
   if (!res.succeeded())
     return nullptr;
 
+  if (postpasses != "") {
+    mlir::PassManager pm(nf->getContext());
+    std::string error_message;
+    //llvm::raw_string_ostream error_stream(error_message);
+    mlir::LogicalResult result =
+        mlir::parsePassPipeline(postpasses, pm);
+    if (mlir::failed(result)) {
+      return nullptr;
+    }
+
+    if (!mlir::succeeded(pm.run(nf))) {
+      return nullptr;
+    }
+  }
+
   return nf;
 }

enzyme/Enzyme/MLIR/Interfaces/GradientUtils.cpp

@@ -37,13 +37,13 @@ mlir::enzyme::MGradientUtils::MGradientUtils(
     ArrayRef<DIFFE_TYPE> ReturnActivity, ArrayRef<DIFFE_TYPE> ArgDiffeTypes_,
     IRMapping &originalToNewFn_,
     std::map<Operation *, Operation *> &originalToNewFnOps_,
-    DerivativeMode mode, unsigned width, bool omp)
+    DerivativeMode mode, unsigned width, bool omp, llvm::StringRef postpasses)
     : newFunc(newFunc_), Logic(Logic), mode(mode), oldFunc(oldFunc_),
       invertedPointers(invertedPointers_), originalToNewFn(originalToNewFn_),
       originalToNewFnOps(originalToNewFnOps_), blocksNotForAnalysis(),
       activityAnalyzer(std::make_unique<enzyme::ActivityAnalyzer>(
           blocksNotForAnalysis, constantvalues_, activevals_, ReturnActivity)),
-      TA(TA_), TR(TR_), omp(omp), returnPrimals(returnPrimals),
+      TA(TA_), TR(TR_), omp(omp), postpasses(postpasses), returnPrimals(returnPrimals),
       returnShadows(returnShadows), width(width), ArgDiffeTypes(ArgDiffeTypes_),
       RetDiffeTypes(ReturnActivity) {}
 

enzyme/Enzyme/MLIR/Interfaces/GradientUtils.h

@@ -36,6 +36,7 @@ class MGradientUtils {
   MTypeAnalysis &TA;
   MTypeResults TR;
   bool omp;
+  llvm::StringRef postpasses;
   const llvm::ArrayRef<bool> returnPrimals;
   const llvm::ArrayRef<bool> returnShadows;
 
@@ -58,7 +59,7 @@ class MGradientUtils {
                  ArrayRef<DIFFE_TYPE> ArgDiffeTypes_,
                  IRMapping &originalToNewFn_,
                  std::map<Operation *, Operation *> &originalToNewFnOps_,
-                 DerivativeMode mode, unsigned width, bool omp);
+                 DerivativeMode mode, unsigned width, bool omp, llvm::StringRef postpasses);
   void erase(Operation *op) { op->erase(); }
   void replaceOrigOpWith(Operation *op, ValueRange vals) {
     for (auto &&[res, rep] : llvm::zip(op->getResults(), vals)) {
@@ -113,11 +114,11 @@ class MDiffeGradientUtils : public MGradientUtils {
                       ArrayRef<DIFFE_TYPE> RetActivity,
                       ArrayRef<DIFFE_TYPE> ArgActivity, IRMapping &origToNew_,
                       std::map<Operation *, Operation *> &origToNewOps_,
-                      DerivativeMode mode, unsigned width, bool omp)
+                      DerivativeMode mode, unsigned width, bool omp, llvm::StringRef postpasses)
       : MGradientUtils(Logic, newFunc_, oldFunc_, TA, TR, invertedPointers_,
                        returnPrimals, returnShadows, constantvalues_,
                        activevals_, RetActivity, ArgActivity, origToNew_,
-                       origToNewOps_, mode, width, omp),
+                       origToNewOps_, mode, width, omp, postpasses),
         initializationBlock(&*(newFunc.getFunctionBody().begin())) {}
 
   // Technically diffe constructor
@@ -127,7 +128,7 @@ class MDiffeGradientUtils : public MGradientUtils {
       const llvm::ArrayRef<bool> returnPrimals,
       const llvm::ArrayRef<bool> returnShadows,
       ArrayRef<DIFFE_TYPE> RetActivity, ArrayRef<DIFFE_TYPE> ArgActivity,
-      mlir::Type additionalArg, bool omp) {
+      mlir::Type additionalArg, bool omp, llvm::StringRef postpasses) {
     std::string prefix;
 
     switch (mode) {
@@ -163,7 +164,7 @@ class MDiffeGradientUtils : public MGradientUtils {
     return new MDiffeGradientUtils(
         Logic, newFunc, todiff, TA, TR, invertedPointers, returnPrimals,
         returnShadows, constant_values, nonconstant_values, RetActivity,
-        ArgActivity, originalToNew, originalToNewOps, mode, width, omp);
+        ArgActivity, originalToNew, originalToNewOps, mode, width, omp, postpasses);
   }
 };
 

enzyme/Enzyme/MLIR/Passes/EnzymeMLIRPass.cpp

@@ -31,6 +31,21 @@ struct DifferentiatePass : public DifferentiatePassBase<DifferentiatePass> {
 
   void runOnOperation() override;
 
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    mlir::PassManager pm(nf->getContext());
+    std::string error_message;
+    //llvm::raw_string_ostream error_stream(error_message);
+    mlir::LogicalResult result =
+        mlir::parsePassPipeline(postpasses, pm);
+    if (!mlir::failed(result)) {
+      pm.getDependentDialects(registry);
+    }
+
+    registry.insert<mlir::arith::ArithDialect, mlir::complex::ComplexDialect,
+                    mlir::cf::ControlFlowDialect, mlir::tensor::TensorDialect>();
+  }
+
   static std::vector<DIFFE_TYPE> mode_from_fn(FunctionOpInterface fn,
                                               DerivativeMode mode) {
     std::vector<DIFFE_TYPE> retTypes;
@@ -150,7 +165,7 @@ struct DifferentiatePass : public DifferentiatePassBase<DifferentiatePass> {
     FunctionOpInterface newFunc = Logic.CreateForwardDiff(
         fn, retType, constants, TA, returnPrimals, mode, freeMemory, width,
         /*addedType*/ nullptr, type_args, volatile_args,
-        /*augmented*/ nullptr);
+        /*augmented*/ nullptr, postpasses);
     if (!newFunc)
       return failure();
 
@@ -276,7 +291,7 @@ struct DifferentiatePass : public DifferentiatePassBase<DifferentiatePass> {
         Logic.CreateReverseDiff(fn, retType, arg_activities, TA, returnPrimals,
                                 returnShadows, mode, freeMemory, width,
                                 /*addedType*/ nullptr, type_args, volatile_args,
-                                /*augmented*/ nullptr);
+                                /*augmented*/ nullptr, postpasses);
     if (!newFunc)
       return failure();
 

enzyme/Enzyme/MLIR/Passes/Passes.td

@@ -19,6 +19,15 @@ def DifferentiatePass : Pass<"enzyme"> {
     "cf::ControlFlowDialect",
     "tensor::TensorDialect",
   ];
+  let options = [
+    Option<
+      /*C++ variable name=*/"postopt",
+      /*CLI argument=*/"postopt",
+      /*type=*/"std::string",
+      /*default=*/"",
+      /*description=*/"Optimization passes to apply to generated derivative functions"
+    >,
+  ],
   let constructor = "mlir::enzyme::createDifferentiatePass()";
 }