Distributed passes boilerplate

Author: EganJ

src/enzyme_ad/jax/BUILD

@@ -584,6 +584,31 @@ gentbl_cc_library(
     ],
 )
 
+td_library(
+    name = "DistributedPassesTdFiles",
+    srcs = [
+    ],
+    deps = [
+        "@llvm-project//mlir:PassBaseTdFiles",
+    ],
+)
+
+gentbl_cc_library(
+    name = "DistributedPassesIncGen",
+    tbl_outs = [
+        (
+            [
+                "-gen-pass-decls",
+                "-name=distributed",
+            ],
+            "Passes/Distributed/Passes.h.inc",
+        ),
+    ],
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "Passes/Distributed/Passes.td",
+    deps = [":DistributedPassesTdFiles"],
+)
+
 td_library(
     name = "TesseraDialectTdFiles",
     srcs = [
@@ -715,6 +740,7 @@ cc_library(
     srcs = glob([
         "Implementations/*.cpp",
         "Passes/*.cpp",
+        "Passes/Distributed/*.cpp",
         "Dialect/*.cpp",
         "Dialect/Distributed/*.cpp",
         "Dialect/Tessera/*.cpp",
@@ -724,6 +750,7 @@ cc_library(
     hdrs = glob([
         "Implementations/*.h",
         "Passes/*.h",
+        "Passes/Distributed/*.h",
         "Dialect/*.h",
         "Dialect/Distributed/*.h",
         "Dialect/Tessera/*.h",
@@ -744,6 +771,7 @@ cc_library(
         ":DistributedDialectIncGen",
         ":DistributedInterfacesIncGen",
         ":DistributedOpsIncGen",
+        ":DistributedPassesIncGen",
         ":DistributedTypesIncGen",
         ":EnzymeHLOPatternsIncGen",
         ":EnzymeXLAAttrsIncGen",

src/enzyme_ad/jax/Passes/Distributed/Passes.h

@@ -0,0 +1,13 @@
+#ifndef ENZYMEXLA_DISTRIBUTED_PASSES_H
+#define ENZYMEXLA_DISTRIBUTED_PASSES_H
+
+namespace mlir::enzyme::distributed {
+
+#define GEN_PASS_DECL
+#include "src/enzyme_ad/jax/Passes/Distributed/Passes.h.inc"
+#define GEN_PASS_REGISTRATION
+#include "src/enzyme_ad/jax/Passes/Distributed/Passes.h.inc"
+
+} // namespace mlir::enzyme::distributed
+
+#endif // ENZYMEXLA_DISTRIBUTED_PASSES_H

src/enzyme_ad/jax/Passes/Distributed/Passes.td

@@ -0,0 +1,15 @@
+#ifndef ENZYMEXLA_DISTRIBUTED_PASSES
+#define ENZYMEXLA_DISTRIBUTED_PASSES
+
+include "mlir/Pass/PassBase.td"
+
+def EliminateConstantCommunicationPass : Pass<"eliminate-constant-communication"> {
+    let summary = "Replaces communicated constants with local constants";
+    let description = [{
+      This pass identifies send instructions with constant operands and replaces
+      the corresponding receive instructions with local constants.
+    }];
+    let dependentDialects = ["enzyme::distributed::DistributedDialect"];
+}
+
+#endif // ENZYMEXLA_DISTRIBUTED_PASSES

src/enzyme_ad/jax/RegistryUtils.cpp

@@ -86,6 +86,7 @@
 #include "mlir/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.h"
 
 #include "src/enzyme_ad/jax/Dialect/Ops.h"
+#include "src/enzyme_ad/jax/Passes/Distributed/Passes.h"
 #include "src/enzyme_ad/jax/Passes/Passes.h"
 
 #include "src/enzyme_ad/jax/Dialect/Distributed/Dialect.h"