feat: register various stablehlo passes (#1391)

* feat: register various stablehlo passes

* feat: centralize dialect registration in EnzymeJAX

* fix: update passes

* chore: bump commit

* fix: build

* fix: linking errors

* chore: bump commit

Author: avik-pal

deps/ReactantExtra/API.cpp

@@ -39,11 +39,14 @@
 #include "src/enzyme_ad/jax/Dialect/Dialect.h"
 #include "src/enzyme_ad/jax/Implementations/XLADerivatives.h"
 #include "src/enzyme_ad/jax/Passes/Passes.h"
+#include "src/enzyme_ad/jax/RegistryUtils.h"
 #include "llvm/Support/TargetSelect.h"
 
 #include "mlir/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.h"
 #include "stablehlo/dialect/ChloOps.h"
 #include "stablehlo/dialect/StablehloOps.h"
+#include "stablehlo/transforms/Passes.h"
+#include "stablehlo/transforms/optimization/Passes.h"
 
 #include "absl/log/globals.h"
 #include "absl/log/initialize.h"
@@ -148,10 +151,6 @@
 
 #include "jaxlib/mosaic/dialect/tpu/tpu_dialect.h"
 
-// Triton did a dumb thing and their import is incompatible
-// We don't use so disabling until upstream fix
-// #include "triton/Dialect/Triton/IR/Dialect.h"
-
 #include "llvm/Support/ExtensibleRTTI.h"
 #include <llvm/Support/FileSystem.h>
 #include <llvm/Support/raw_ostream.h>
@@ -165,10 +164,6 @@ void registerRemoveTransformPass();
 void registerGenerateApplyPatternsPass();
 } // namespace enzyme
 
-namespace triton {
-class TritonDialect;
-}
-
 } // namespace mlir
 
 namespace reactant {
@@ -1120,25 +1115,15 @@ extern "C" int PjRtLoadedExecutableNumPartitions(PjRtLoadedExecutable *exec) {
   return exec->num_partitions();
 }
 
-void prepareRegistry(mlir::DialectRegistry &registry);
-
 extern "C" void RegisterDialects(MlirContext cctx) {
   mlir::MLIRContext &context = *unwrap(cctx);
   DialectRegistry registry;
-  prepareRegistry(registry);
+  mlir::enzyme::prepareRegistry(registry);
+  mlir::enzyme::registerDialects(registry);
+  mlir::enzyme::registerInterfaces(registry);
+
   context.appendDialectRegistry(registry);
-  context.loadDialect<mlir::arith::ArithDialect>();
-  context.loadDialect<mlir::enzyme::EnzymeDialect>();
-  context.loadDialect<mlir::enzymexla::EnzymeXLADialect>();
-  // context.loadDialect<mlir::triton::TritonDialect>();
-  context.loadDialect<mlir::tpu::TPUDialect>();
-  context.loadDialect<mlir::tensor::TensorDialect>();
-  context.loadDialect<mlir::func::FuncDialect>();
-  context.loadDialect<mlir::mhlo::MhloDialect>();
-  context.loadDialect<mlir::stablehlo::StablehloDialect>();
-  context.loadDialect<mlir::chlo::ChloDialect>();
-  context.loadDialect<mlir::sdy::SdyDialect>();
-  context.loadDialect<mlir::LLVM::LLVMDialect>();
+  mlir::enzyme::loadAllRegisteredDialects(context);
 }
 
 #include "mlir/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.h"
@@ -1147,53 +1132,14 @@ extern "C" void RegisterDialects(MlirContext cctx) {
 #include "xla/service/spmd/shardy/sdy_round_trip/pipelines.h"
 
 extern "C" void InitializePasses(MlirDialectRegistry creg) {
-  mlir::registerenzymePasses();
-  enzyme::registerenzymexlaPasses();
-
-  // Register the standard passes we want.
-  mlir::registerTransformsPasses();
-  mlir::registerLowerAffinePass();
-  mlir::registerSCCPPass();
-  mlir::registerInlinerPass();
-  mlir::registerSymbolDCEPass();
-  mlir::registerLoopInvariantCodeMotionPass();
-  mlir::registerConvertSCFToOpenMPPass();
-  mlir::affine::registerAffinePasses();
-  mlir::registerReconcileUnrealizedCastsPass();
-
-  /*
-    registry.addExtension(+[](MLIRContext *ctx, LLVM::LLVMDialect *dialect) {
-      LLVM::LLVMFunctionType::attachInterface<MemRefInsider>(*ctx);
-      LLVM::LLVMArrayType::attachInterface<MemRefInsider>(*ctx);
-      LLVM::LLVMPointerType::attachInterface<MemRefInsider>(*ctx);
-      LLVM::LLVMStructType::attachInterface<MemRefInsider>(*ctx);
-      MemRefType::attachInterface<PtrElementModel<MemRefType>>(*ctx);
-      LLVM::LLVMStructType::attachInterface<
-          PtrElementModel<LLVM::LLVMStructType>>(*ctx);
-      LLVM::LLVMPointerType::attachInterface<
-          PtrElementModel<LLVM::LLVMPointerType>>(*ctx);
-      LLVM::LLVMArrayType::attachInterface<PtrElementModel<LLVM::LLVMArrayType>>(
-          *ctx);
-    });
-    */
-
-  // Transform dialect and extensions.
-  mlir::transform::registerInterpreterPass();
-  mlir::enzyme::registerGenerateApplyPatternsPass();
-  mlir::enzyme::registerRemoveTransformPass();
-
-  // xla + shardy specific passes
-  xla::sdy::registerSdyRoundTripExportPipeline();
-  xla::sdy::registerSdyRoundTripImportPipeline();
-  mlir::sdy::registerAllSdyPassesAndPipelines();
-  xla::sdy::registerStablehloExportPipeline();
-  xla::sdy::registerStablehloImportPipeline();
-  xla::sdy::registerStablehloImportShardingsPass();
+  mlir::enzyme::initializePasses();
 }
 
 extern "C" void InitializeRegistry(MlirDialectRegistry creg) {
   mlir::DialectRegistry &registry = *unwrap(creg);
-  prepareRegistry(registry);
+  mlir::enzyme::prepareRegistry(registry);
+  mlir::enzyme::registerDialects(registry);
+  mlir::enzyme::registerInterfaces(registry);
 
   mlir::registerLLVMDialectImport(registry);
   mlir::registerNVVMDialectImport(registry);
@@ -2326,7 +2272,8 @@ extern "C" mlir::sdy::TensorShardingAttr hloShardingToTensorShardingAttr(
 
   return mlir::sdy::TensorShardingAttr::get(
       context, meshName, tensorShardingAttr.getDimShardings(),
-      tensorShardingAttr.getReplicatedAxes(), tensorShardingAttr.getUnreducedAxes());
+      tensorShardingAttr.getReplicatedAxes(),
+      tensorShardingAttr.getUnreducedAxes());
 }
 
 #pragma endregion

deps/ReactantExtra/BUILD

@@ -761,25 +761,18 @@ cc_library(
         ],
 
     ) + [
-        "@enzyme_ad//src/enzyme_ad/jax:RegistryUtils.cpp",
         "@enzyme_ad//src/enzyme_ad/jax:gpu.cc",
         "@enzyme_ad//src/enzyme_ad/jax:cpu.cc",
-        # "@com_google_protobuf//:src/google/protobuf/io/coded_stream.cc",
-        # "@xla//xla:xla.pb.cc",
-        "@xla//xla:xla_data.pb.cc",
-        # "@xla//xla/stream_executor:device_description.pb.cc",
-        "@xla//xla/service:hlo.pb.cc",
-        #  # "@tsl//tsl/protobuf:dnn.pb.cc",
-        #"@tsl//tsl/protobuf:histogram.pb.cc",
-        #"@tsl//tsl/protobuf:bfc_memory_map_proto",bfc_memory_map.pb.cc",
         "@xla//xla/service/gpu:backend_configs.pb.cc",
         "@xla//xla:autotuning.pb.cc",
         "@xla//xla:autotune_results.pb.cc",
         "@xla//xla/service:buffer_assignment.pb.cc",
         ],
     hdrs = glob([
         "*.h",
-    ]),
+    ]) + [
+        "@enzyme_ad//src/enzyme_ad/jax:RegistryUtils.h",
+    ],
     copts = [
         "-Werror=unused-variable",
         "-Werror=unused-but-set-variable",
@@ -901,6 +894,7 @@ cc_library(
 "-Wl,-exported_symbol,_addSdyPropagationPipeline",
     ]}),
     deps = [
+        "@enzyme_ad//src/enzyme_ad/jax:RegistryUtils",
         "@enzyme//:EnzymeMLIR",
         "@llvm-project//mlir:AffineDialect",
         "@llvm-project//mlir:AllPassesAndDialects",
@@ -1031,6 +1025,21 @@ cc_library(
         "@xla//xla/pjrt/plugin/xla_cpu:xla_cpu_pjrt_client",
         "@xla//xla/pjrt/plugin/xla_gpu:xla_gpu_pjrt_client",
         "@xla//xla/pjrt/plugin/xla_tpu:xla_tpu_pjrt_client",
+        "@stablehlo//:linalg_passes",
+        "@stablehlo//:tosa_passes",
+        "@stablehlo//:stablehlo_passes",
+        "@stablehlo//:stablehlo_passes_optimization",
+        "@stablehlo//stablehlo/tests:check_ops",
+
+        "@tsl//tsl/platform:env",
+        "@xla//xla/tsl/protobuf:dnn_proto_cc_impl",
+        "@xla//xla/tsl/protobuf:histogram_proto_cc",
+        "@xla//xla/tsl/protobuf:histogram_proto_cc_impl",
+        "@xla//xla:xla_data_proto_cc_impl",
+        "@xla//xla/tsl/platform:env",
+        "@xla//xla/tsl/platform:errors",
+        "@xla//xla/service:hlo_proto_cc_impl",
+        "@com_google_absl//absl/status:statusor",
     ] + select({
         "@xla//xla/tsl:is_cuda_enabled_and_oss":[
             "@xla//xla/stream_executor:cuda_platform",

deps/ReactantExtra/WORKSPACE

@@ -9,7 +9,7 @@ http_archive(
     urls = ["https://github.com/wsmoses/nsync/archive/{commit}.tar.gz".format(commit = NSYNC_COMMIT)],
 )
 
-ENZYMEXLA_COMMIT = "bc7a3e5be1e08984329bd16753b07346eda07996"
+ENZYMEXLA_COMMIT = "4a19bc9910f34f4afcd2c2559c8776f7ca86ae4d"
 ENZYMEXLA_SHA256 = ""
 
 http_archive(