Add raising skeleton

Author: vimarsh6739

src/enzyme_ad/jax/Passes/EnzymeBatchToStableHLOPass.cpp

@@ -0,0 +1,49 @@
+//===- EnzymeBatchToStableHLOPass.cpp  ------------------------------------===//
+//
+// 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 a pass to print the MLIR module
+//===----------------------------------------------------------------------===//
+
+#include "src/enzyme_ad/jax/Passes/Passes.h"
+#include "stablehlo/dialect/StablehloOps.h"
+#include "Enzyme/MLIR/Dialect/Dialect.h"
+#include "Enzyme/MLIR/Dialect/Ops.h"
+#include "src/enzyme_ad/jax/Utils.h"
+
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+namespace mlir {
+namespace enzyme {
+#define GEN_PASS_DEF_ENZYMEBATCHTOSTABLEHLOPASS
+#include "src/enzyme_ad/jax/Passes/Passes.h.inc"
+} // namespace enzyme
+} // namespace mlir
+
+using namespace mlir;
+using namespace mlir::enzyme;
+using namespace enzyme;
+namespace {
+struct EnzymeBatchToStableHLOPass
+    : public enzyme::impl::EnzymeBatchToStableHLOPassBase<
+          EnzymeBatchToStableHLOPass> {
+  void runOnOperation() override {
+    MLIRContext *context = &getContext();
+    RewritePatternSet patterns(context);
+    ConversionTarget target(*context);
+    target.addLegalDialect<stablehlo::StablehloDialect>();
+    target.addLegalDialect<enzyme::EnzymeDialect>();
+    target.addIllegalOp<enzyme::ConcatOp, enzyme::ExtractOp>();
+
+    if (failed(applyPartialConversion(getOperation(), target,
+                                      std::move(patterns)))) {
+      signalPassFailure();
+    }
+  };
+};
+} // namespace

src/enzyme_ad/jax/Passes/Passes.td

@@ -1008,4 +1008,13 @@ def SCFCPUify : Pass<"cpuify"> {
   Option<"method", "method", "std::string", /*default=*/"\"distribute\"", "Method of doing distribution">
   ];
 }
+
+def EnzymeBatchToStableHLOPass : Pass<"enzyme-batch-to-stablehlo"> {
+  let summary = "Legalize batching specific enzyme ops to stablehlo dialect";
+  let dependentDialects = [
+    "stablehlo::StablehloDialect",
+    "enzyme::EnzymeDialect"
+  ];
+}
+
 #endif