automatic differentiation with Enzyme

XLA_VERSION

@@ -1 +1 @@
-2fb20601f1cc6cab7f29f8bc73d90cd31e74bba0
+4ec7e2a7721ace136bae967e6881fa7035d0b35a

dev.sh

@@ -1,11 +1,6 @@
 #!/bin/sh -e
 
-install_xla () {
-  if [ -z "$2" ]; then
-    echo "Usage: install_xla <xla-revision> <install-path>."
-    exit 1;
-  fi
-
+install_git_repository () {
   if [ "$(ls -A "$2")" ]; then
     echo "Directory at path $2 is not empty, refusing to install XLA to this directory."
     exit 1;
@@ -14,8 +9,16 @@ install_xla () {
   (
     cd "$2"
     git init
-    git remote add origin https://github.com/openxla/xla
+    git remote add origin "$3"
     git fetch --depth 1 origin "$1"
     git checkout FETCH_HEAD
   )
 }
+
+install_xla () {
+  install_git_repository "$1" "$2" https://github.com/openxla/xla
+}
+
+install_enzyme () {
+  install_git_repository "$1" "$2" https://github.com/EnzymeAD/Enzyme-JAX.git
+}

pjrt-plugins/xla-cpu/.gitignore

@@ -0,0 +1 @@
+xla/

pjrt-plugins/xla-cpu/build.sh

@@ -23,7 +23,8 @@ case $osu in
     ;;
 esac
 
-xla_dir=$(mktemp -d)
+xla_dir=pjrt-plugins/xla-cpu/xla
+mkdir "$xla_dir"
 install_xla "$xla_version" "$xla_dir"
 (
   cd "$xla_dir"

pjrt-plugins/xla-cuda/build.sh

@@ -18,7 +18,8 @@ case $osu in
     ;;
 esac
 
-xla_dir=$(mktemp -d)
+xla_dir=pjrt-plugins/xla-cuda/xla
+mkdir "$xla_dir"
 install_xla "$xla_version" "$xla_dir"
 (
   cd "$xla_dir"

spidr/backend/.gitignore

@@ -1 +1,2 @@
+/Enzyme-JAX
 /xla

spidr/backend/BUILD

@@ -12,21 +12,37 @@ cc_binary(
     linkshared = True,
     linkstatic = True,
     srcs = [
+        "//src/Enzyme-JAX",
+        "//src/Enzyme",
+        "//src/llvm",
+        "//src/mlir",
+        "//src/stablehlo",
         "//src/xla",
         "//src/xla/client",
         "//src/xla/hlo/builder",
         "//src/xla/hlo/builder/lib",
+        "//src/xla/hlo/translate",
+        "//src/xla/mlir_hlo/mhlo/IR",
         "//src/xla/pjrt",
         "//src/xla/pjrt/c",
+        "//src/xla/service",
         "//src",
     ],
     deps = [
+        "//src/Enzyme-JAX",
+        "//src/Enzyme",
+        "//src/llvm",
+        "//src/mlir",
+        "//src/stablehlo",
         "//src/xla",
         "//src/xla/client",
         "//src/xla/hlo/builder",
         "//src/xla/hlo/builder/lib",
+        "//src/xla/hlo/translate",
+        "//src/xla/mlir_hlo/mhlo/IR",
         "//src/xla/pjrt",
         "//src/xla/pjrt/c",
+        "//src/xla/service",
         "//src",
     ],
 )

spidr/backend/BUILD.patch

@@ -0,0 +1,15 @@
+--- BUILD	2025-02-04 00:03:30.224943988 +0000
++++ BUILD.diff	2025-02-04 00:02:33.364805000 +0000
+@@ -385,12 +385,6 @@
+         "Passes/*.h",
+         "Dialect/*.h",
+     ]),
+-    copts = [
+-        "-Werror=unused-variable",
+-        "-Werror=unused-but-set-variable",
+-        "-Werror=return-type",
+-        "-Werror=unused-result",
+-    ],
+     deps = [
+         ":EnzymeXLAOpsIncGen",
+         ":EnzymeXLAPassesIncGen",

spidr/backend/ENZYME_JAX_VERSION

@@ -0,0 +1 @@
+ae84d009fe9385e8752a8ee8406115d662d24967

spidr/backend/VERSION

@@ -1 +1 @@
-0.0.16
+0.0.17

spidr/backend/WORKSPACE

@@ -1,3 +1,5 @@
+### xla
+
 # this must be a local repository not http archive
 # so we can run ./configure.py before invoking bazel
 local_repository(name = "xla", path = "xla")
@@ -28,3 +30,34 @@ xla_workspace0()
 
 load("@tsl//third_party/gpus/cuda/hermetic:cuda_configure.bzl", "cuda_configure")
 cuda_configure(name = "local_config_cuda")
+
+# Enzyme-JAX
+
+# should we use http_archive for enzyme-jax?
+# how best to use enzyme's XLA version?
+local_repository(name = "enzyme-jax", path = "Enzyme-JAX")
+load("@enzyme-jax//:workspace.bzl", "JAX_COMMIT", "JAX_SHA256", "ENZYME_COMMIT", "ENZYME_SHA256")
+
+load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
+
+http_archive(
+    name = "hedron_compile_commands",
+    url = "https://github.com/hedronvision/bazel-compile-commands-extractor/archive/4f28899228fb3ad0126897876f147ca15026151e.tar.gz",
+    strip_prefix = "bazel-compile-commands-extractor-4f28899228fb3ad0126897876f147ca15026151e",
+)
+
+http_archive(
+    name = "jax",
+    sha256 = JAX_SHA256,
+    strip_prefix = "jax-" + JAX_COMMIT,
+    urls = ["https://github.com/google/jax/archive/{commit}.tar.gz".format(commit = JAX_COMMIT)],
+    patch_args = ["-p1"],
+    patches = ["@enzyme-jax//:patches/jax.patch"],
+)
+
+http_archive(
+    name = "enzyme",
+    sha256 = ENZYME_SHA256,
+    strip_prefix = "Enzyme-" + ENZYME_COMMIT + "/enzyme",
+    urls = ["https://github.com/EnzymeAD/Enzyme/archive/{commit}.tar.gz".format(commit = ENZYME_COMMIT)],
+)

spidr/backend/build.sh

@@ -3,7 +3,8 @@
 script_dir=$(CDPATH="" cd -- "$(dirname -- "$0")" && pwd)
 cd "$script_dir/../.."
 . ./dev.sh
-rev="$(cat XLA_VERSION)"
+xla_rev="$(cat XLA_VERSION)"
+enzyme_rev="$(cat spidr/backend/ENZYME_JAX_VERSION)"
 
 osu="$(uname)"
 case $osu in
@@ -26,9 +27,14 @@ esac
 (
   cd spidr/backend
   mkdir xla
-  install_xla "$rev" xla
+  install_xla "$xla_rev" xla
   (cd xla; ./configure.py --backend=cpu --os=$os)
+  # depending on Enzyme-JAX is problematic as it fixes the XLA version. Can we only depend on enzyme?
+  # seems unlikely that they could decouple XLA entirely. They almost certainly can't decouple stablehlo
+  mkdir Enzyme-JAX
+  install_enzyme "$enzyme_rev" Enzyme-JAX
+  patch Enzyme-JAX/src/enzyme_ad/jax/BUILD < BUILD.patch
   bazel build //:c_xla
-  rm -rf xla
+  rm -rf xla Enzyme-JAX
 )
 mv "spidr/backend/bazel-bin/libc_xla.$ext" "libc_xla-$os-$arch.$ext"

spidr/backend/src/Enzyme-JAX/BUILD

@@ -0,0 +1,14 @@
+cc_library(
+    name = "Enzyme-JAX",
+    linkstatic = True,
+    alwayslink = True,
+    visibility = ["//visibility:public"],
+    srcs = glob(["**/*.cpp"]),
+    hdrs = glob(["**/*.h"]),
+    deps = [
+        "@enzyme//:EnzymeMLIR",
+        "@enzyme-jax//src/enzyme_ad/jax:XLADerivatives",
+        "@llvm-project//mlir:AllPassesAndDialects",
+        "//src/mlir",
+    ],
+)

spidr/backend/src/Enzyme-JAX/src/enzyme_ad/jax/Implementations/XLADerivatives.cpp

@@ -0,0 +1,33 @@
+/*
+Copyright 2025 Joel Berkeley
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+#include "Enzyme/MLIR/Dialect/Dialect.h"
+#include "Enzyme/MLIR/Dialect/Ops.h"
+#include "Enzyme/MLIR/Passes/Passes.h"
+#include "src/enzyme_ad/jax/Implementations/XLADerivatives.h"
+
+#include "../../../../../mlir/IR/DialectRegistry.h"
+
+extern "C" {
+    void registerStableHLODialectAutoDiffInterface(DialectRegistry* registry) {
+        auto registry_ = reinterpret_cast<mlir::DialectRegistry*>(registry);
+        mlir::enzyme::registerStableHLODialectAutoDiffInterface(*registry_);
+    }
+
+    void registerCHLODialectAutoDiffInterface(DialectRegistry* registry) {
+        auto registry_ = reinterpret_cast<mlir::DialectRegistry*>(registry);
+        mlir::enzyme::registerCHLODialectAutoDiffInterface(*registry_);
+    }
+}

spidr/backend/src/Enzyme-JAX/src/enzyme_ad/jax/Passes/Passes.cpp

@@ -0,0 +1,26 @@
+/*
+Copyright 2025 Joel Berkeley
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+#include "mlir/Pass/PassManager.h"
+#include "src/enzyme_ad/jax/Passes/Passes.h"
+
+#include "../../../../../mlir/Pass/PassManager.h"
+
+extern "C" {
+    void PassManager_addPass_ArithRaisingPass(PassManager& s) {
+        auto& s_ = reinterpret_cast<mlir::PassManager&>(s);
+        s_.addPass(mlir::enzyme::createArithRaisingPass());
+    }
+}

spidr/backend/src/Enzyme/BUILD

@@ -0,0 +1,13 @@
+cc_library(
+    name = "Enzyme",
+    linkstatic = True,
+    alwayslink = True,
+    srcs = glob(["**/*.cpp"]),
+    hdrs = glob(["**/*.h"]),
+    deps = [
+        "@enzyme//:EnzymeMLIR",
+        "@llvm-project//mlir:IR",
+        "//src/mlir",
+    ],
+    visibility = ["//visibility:public"],
+)

spidr/backend/src/Enzyme/MLIR/Dialect/Dialect.cpp

@@ -0,0 +1,30 @@
+/*
+Copyright 2025 Joel Berkeley
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+#include "mlir/IR/DialectRegistry.h"
+#include "Enzyme/MLIR/Dialect/Dialect.h"
+
+#include "../../../mlir/IR/DialectRegistry.h"
+#include "../../../mlir/IR/MLIRContext.h"
+
+extern "C" {
+    void DialectRegistry_insert_EnzymeDialect(DialectRegistry& s) {
+        reinterpret_cast<mlir::DialectRegistry&>(s).insert<mlir::enzyme::EnzymeDialect>();
+    }
+
+    void MLIRContext_loadDialect_EnzymeDialect(MLIRContext& s) {
+        reinterpret_cast<mlir::MLIRContext&>(s).loadDialect<mlir::enzyme::EnzymeDialect>();
+    }
+}

spidr/backend/src/Enzyme/MLIR/Implementations/CoreDialectsAutoDiffImplementations.cpp

@@ -0,0 +1,25 @@
+/*
+Copyright 2025 Joel Berkeley
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+#include "Enzyme/MLIR/Implementations/CoreDialectsAutoDiffImplementations.h"
+
+#include "../../../mlir/IR/DialectRegistry.h"
+
+extern "C" {
+    void registerCoreDialectAutodiffInterfaces(DialectRegistry* registry) {
+        auto registry_ = reinterpret_cast<mlir::DialectRegistry*>(registry);
+        mlir::enzyme::registerCoreDialectAutodiffInterfaces(*registry_);
+    }
+}

spidr/backend/src/Enzyme/MLIR/Passes/Passes.cpp

@@ -0,0 +1,31 @@
+/*
+Copyright 2025 Joel Berkeley
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+#include "Enzyme/MLIR/Passes/Passes.h"
+#include "mlir/Pass/PassManager.h"
+
+#include "../../../mlir/Pass/PassManager.h"
+
+extern "C" {
+    void registerenzymePasses() {
+        // where on earth does this come from?
+        mlir::registerenzymePasses();
+    }
+
+    void PassManager_addPass_RemoveUnusedEnzymeOpsPass(PassManager& s) {
+        auto& s_ = reinterpret_cast<mlir::PassManager&>(s);
+        s_.addPass(mlir::enzyme::createRemoveUnusedEnzymeOpsPass());
+    }
+}