MLIR Reverse Mode (#37)

* WIP jax reverse more

* fix

* fix enzyme commit

* cleanup format

* fixup

* tmp

* continue

* fix

* fix

* fixup

* fix bug

* fixup

* more fixup

* sliceslice

* fix pad opt

* cleanup

* fix

Author: wsmoses

BUILD

@@ -18,6 +18,35 @@ py_package(
     packages = ["@//src/enzyme_ad/jax:enzyme_call.so", "@llvm-project//clang:builtin_headers_gen"],
 )
 
+cc_binary(
+    name = "enzymexlamlir-opt",
+    srcs = ["//src/enzyme_ad/jax:enzymexlamlir-opt.cpp"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//src/enzyme_ad/jax:XLADerivatives",
+        "@enzyme//:EnzymeMLIR",
+        "@llvm-project//mlir:AffineDialect",
+        "@llvm-project//mlir:AllPassesAndDialects",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:AsyncDialect",
+        "@llvm-project//mlir:ControlFlowDialect",
+        "@llvm-project//mlir:ConversionPasses",
+        "@llvm-project//mlir:DLTIDialect",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:GPUDialect",
+        "@llvm-project//mlir:LLVMDialect",
+        "@llvm-project//mlir:LinalgDialect",
+        "@llvm-project//mlir:MathDialect",
+        "@llvm-project//mlir:MemRefDialect",
+        "@llvm-project//mlir:MlirOptLib",
+        "@llvm-project//mlir:NVVMDialect",
+        "@llvm-project//mlir:OpenMPDialect",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:SCFDialect",
+        "@llvm-project//mlir:Transforms",
+    ],
+)
+
 py_wheel(
     name = "enzyme_ad",
     distribution = "enzyme_ad",

WORKSPACE

@@ -39,7 +39,7 @@ http_archive(
     strip_prefix = "xla-" + XLA_COMMIT,
     urls = ["https://github.com/wsmoses/xla/archive/{commit}.tar.gz".format(commit = XLA_COMMIT)],
     patch_args = ["-p1"],
-    patches = ["//:patches/xla.patch"],
+    patches = ["//:patches/xla.patch", "//:patches/xla2.patch", ],
 )
 
 PYRULES_COMMIT = "fe33a4582c37499f3caeb49a07a78fc7948a8949"
@@ -60,12 +60,8 @@ load("@rules_python//python/pip_install:repositories.bzl", "pip_install_dependen
 
 pip_install_dependencies()
 
-ENZYME_COMMIT = "0b621884bc531329095d202f042f6599a86614ec"
-ENZYME_SHA256 = "f9479530b08aeb3ecbf0c420d0e2f222fdf8bcf6c20a218271b365db3a3053ad"
-# local_repository(
-#     name = "enzyme",
-#     path = "../Enzyme/enzyme"
-# )
+ENZYME_COMMIT = "0a129ae7e45114a08f281e50632b9f967fae8396"
+ENZYME_SHA256 = "715982efd0a0ef8038e8ad35047e9c1941eb3f9cb038883342969b0bcc8915ad"
 
 http_archive(
     name = "enzyme",

patches/xla.patch

@@ -16,3 +16,38 @@
  )
 
  cc_library(
+
+--- a/xla/mlir/backends/cpu/transforms/BUILD
++++ b/xla/mlir/backends/cpu/transforms/BUILD
+@@ -4,7 +4,7 @@ load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library")
+ 
+ package(
+     # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+-    default_visibility = ["//xla:internal"],
++    default_visibility = ["//xla:friends"],
+     licenses = ["notice"],
+ )
+ 
+ gentbl_cc_library(
+
+--- a/xla/mlir/memref/BUILD
++++ b/xla/mlir/memref/BUILD
+@@ -1,6 +1,7 @@
+ package_group(
+     name = "friends",
+     packages = [
++        "public",
+         "//xla/mlir/...",
+         # copybara:uncomment_begin(google-only)
+         # # TODO(ezhulenev): Clean up dependencies that are leforvers from Autofusion project.
+
+--- a/xla/mlir/math/BUILD
++++ b/xla/mlir/math/BUILD
+@@ -1,6 +1,7 @@
+ package_group(
+     name = "friends",
+     packages = [
++        "public",
+         "//xla/mlir/...",
+         # copybara:uncomment_begin(google-only)
+         # # TODO(ezhulenev): Clean up dependencies that are leforvers from Autofusion project.

patches/xla2.patch

@@ -0,0 +1,10 @@
+--- a/xla/mlir/runtime/BUILD
++++ b/xla/mlir/runtime/BUILD
+@@ -19,6 +19,7 @@ package_group(
+         # TODO(ezhulenev): All targets depending on mlir must be under xla/mlir folder
+         "//xla/service/cpu/...",
+         "//xla/service/gpu/...",
++        "public",
+     ],
+ )
+ 

src/enzyme_ad/jax/BUILD

@@ -3,6 +3,7 @@ load("@pybind11_bazel//:build_defs.bzl", "pybind_extension", "pybind_library")
 load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library", "td_library")
 load("@llvm-project//llvm:tblgen.bzl", "gentbl")
 
+exports_files(["enzymexlamlir-opt.cpp"])
 licenses(["notice"])
 
 package(
@@ -29,10 +30,12 @@ pybind_library(
         "@llvm-project//llvm:AsmParser",
         "@llvm-project//llvm:CodeGen",
         "@llvm-project//llvm:Core",
+        "@llvm-project//llvm:MC",
         "@llvm-project//llvm:IRReader",
         "@llvm-project//llvm:Linker",
         "@llvm-project//llvm:OrcJIT",
         "@llvm-project//llvm:Support",
+        "@llvm-project//llvm:TargetParser",
     ],
 )
 
@@ -139,6 +142,16 @@ cc_library(
         ":stablehlo-derivatives",
         "@stablehlo//:stablehlo_ops",
         "@stablehlo//:stablehlo_passes",
+        "@stablehlo//:reference_ops",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:TensorDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:FunctionInterfaces",
+        "@llvm-project//mlir:ControlFlowInterfaces",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:CommonFolders",
+        "@llvm-project//mlir:Transforms",
         "@xla//xla/mlir_hlo",
         "@enzyme//:EnzymeMLIR",
     ]
@@ -174,6 +187,9 @@ pybind_library(
         "@xla//xla/client:client_library",
         "@xla//xla/client:executable_build_options",
         "@xla//xla/client:xla_computation",
+        "@xla//xla/service:service",
+        "@xla//xla/service:local_service",
+        "@xla//xla/service:local_service_utils",
         "@xla//xla/service:buffer_assignment_proto_cc",
         "@xla//xla/service:buffer_assignment_proto_cc_impl",
         "@xla//xla/service/cpu:cpu_executable",
@@ -191,23 +207,37 @@ pybind_library(
         "@xla//xla:xla_proto_cc",
         "@xla//xla:xla_proto_cc_impl",
 
+        "@stablehlo//:stablehlo_ops",
+
         # Make CPU target available to XLA.
         "@xla//xla/service:cpu_plugin",
 
         # MHLO stuff.
         "@xla//xla/mlir_hlo",
         "@xla//xla/translate/mhlo_to_hlo:mlir_hlo_to_hlo",
 
+        "@xla//xla/hlo/ir:hlo",
+
         # This is necessary for XLA protobufs to link
         "@com_google_protobuf//:protobuf",
 
         # MLIR dialects and parser.
+        "@llvm-project//llvm:Support",
         "@llvm-project//mlir:ArithDialect",
         "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:FuncExtensions",
+        "@llvm-project//mlir:TensorDialect",
+
         "@llvm-project//mlir:Parser",
+        "@llvm-project//mlir:Pass",
 
+        "@xla//xla/mlir_hlo:all_passes",
+        "@xla//xla:printer",
+
     # EnzymeMLIR
         "@enzyme//:EnzymeMLIR",
+
+        "@com_google_absl//absl/status:statusor",
 
         # Mosaic
 	"@jax//jaxlib/mosaic:tpu_dialect",
@@ -230,6 +260,19 @@ pybind_extension(
         "@com_google_absl//absl/status:statusor",
         "@stablehlo//:stablehlo_passes",
         "@xla//xla/stream_executor:stream_executor_impl",
+        "@xla//xla/mlir/backends/cpu/transforms:passes",
+        "@xla//xla/mlir/memref/transforms:passes",
+        "@xla//xla/mlir/math/transforms:passes",
+        "@xla//xla/mlir/runtime/transforms:passes",
+        "@xla//xla/mlir_hlo:deallocation_passes",
+        "@xla//xla/mlir_hlo:lhlo",
+        "@xla//xla/mlir_hlo:lhlo_gpu",
+        "@xla//xla/mlir_hlo:all_passes",
+        "@xla//xla/hlo/ir:hlo",
+        "@xla//xla/service/cpu:cpu_executable",
+        "@enzyme//:EnzymeStatic",
+        "@enzyme//:EnzymeMLIR",
+
     ],
     visibility = ["//visibility:public"],
 )