Added code to register passes

Author: jessicacotturone21

src/enzyme_ad/jax/BUILD

@@ -629,6 +629,31 @@ gentbl_cc_library(
     ],
 )
 
+td_library(
+    name = "TesseraPassesTdFiles",
+    srcs = [
+    ],
+    deps = [
+        "@llvm-project//mlir:PassBaseTdFiles",
+    ],
+)
+
+gentbl_cc_library(
+    name = "TesseraPassesIncGen",
+    tbl_outs = [
+        (
+            [
+                "-gen-pass-decls",
+                "-name=tessera",
+            ],
+            "Passes/Tessera/Passes.h.inc",
+        ),
+    ],
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "Passes/Tessera/Passes.td",
+    deps = [":TesseraPassesTdFiles"],
+)
+
 cc_library(
     name = "CheckedRewrite",
     hdrs = ["CheckedRewrite.h"],
@@ -643,6 +668,7 @@ cc_library(
     srcs = glob([
         "Implementations/*.cpp",
         "Passes/*.cpp",
+        "Passes/Tessera/*.cpp",
         "Dialect/*.cpp",
         "Dialect/Distributed/*.cpp",
         "Dialect/Tessera/*.cpp",
@@ -652,6 +678,7 @@ cc_library(
     hdrs = glob([
         "Implementations/*.h",
         "Passes/*.h",
+        "Passes/Tessera/*.h",
         "Dialect/*.h",
         "Dialect/Distributed/*.h",
         "Dialect/Tessera/*.h",
@@ -683,6 +710,7 @@ cc_library(
         ":StablehloOptPatternsIncGen",
         ":TesseraDialectIncGen",
         ":TesseraOpsIncGen",
+        ":TesseraPassesIncGen",
         ":chlo-derivatives",
         ":mhlo-derivatives",
         ":stablehlo-derivatives",

src/enzyme_ad/jax/Passes/Tessera/FuncToTessera.cpp

@@ -1,13 +1,14 @@
 //===----------------------------------------------------------------------===//
 //
-// This file implements patterns to convert the Func dialect to the Tessera
-// dialect.
+// This file implements patterns to convert operations in the Func dialect to 
+// operations in the Tessera dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "src/enzyme_ad/jax/Dialect/Tessera/Dialect.h"
+#include "src/enzyme_ad/jax/Passes/Tessera/Passes.h"
 #include "src/enzyme_ad/jax/Passes/Passes.h"
 
 using namespace mlir;

src/enzyme_ad/jax/Passes/Tessera/Passes.h

@@ -0,0 +1,18 @@
+#ifndef TESSERA_PASSES_H
+#define TESSERA_PASSES_H
+
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+namespace tessera {
+
+#define GEN_PASS_DECLS
+#include "Tessera/Passes/Tessera/Passes.h.inc"
+
+#define GEN_PASS_REGISTRATION
+#include "Tessera/Passes/Tessera/Passes.h.inc"
+
+} // namespace tessera
+} // namespace mlir
+
+#endif

src/enzyme_ad/jax/Passes/Tessera/TesseraToFunc.cpp

@@ -1,12 +1,14 @@
 //===----------------------------------------------------------------------===//
 //
-// This file implements patterns to convert the Tessera dialect to the Func
-// dialect.
+// This file implements patterns to convert operations in the Tessera dialect to 
+// operations in the Func dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/BuiltinOps.h"
 #include "src/enzyme_ad/jax/Dialect/Tessera/Dialect.h"
+#include "src/enzyme_ad/jax/Passes/Tessera/Passes.h"
 #include "src/enzyme_ad/jax/Passes/Passes.h"
 
 using namespace mlir;
@@ -103,7 +105,7 @@ class ReturnOpRewrite final : public OpRewritePattern<tessera::ReturnOp> {
 } // namespace
 
 //===----------------------------------------------------------------------===//
-// Pass to convert Func operations into Tessera operations
+// Pass to convert Tessera operations into Func operations
 //===----------------------------------------------------------------------===//
 
 struct TesseraToFuncPass

src/enzyme_ad/jax/RegistryUtils.cpp

@@ -87,6 +87,7 @@
 
 #include "src/enzyme_ad/jax/Dialect/Ops.h"
 #include "src/enzyme_ad/jax/Passes/Passes.h"
+#include "src/enzyme_ad/jax/Passes/Tessera/Passes.h"
 
 #include "src/enzyme_ad/jax/Dialect/Distributed/Dialect.h"
 #include "src/enzyme_ad/jax/Dialect/Tessera/Dialect.h"
@@ -294,6 +295,7 @@ void registerInterfaces(mlir::DialectRegistry &registry) {
 void initializePasses() {
   registerenzymePasses();
   enzyme::registerenzymexlaPasses();
+  enzyme::tessera::registerTesseraPasses();
 
   // Register the standard passes we want.
   mlir::registerCSEPass();