Modularize Bazel build

WORKSPACE

@@ -1,62 +1,25 @@
 # add support for generating compile_commands
-load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
+load("//third_party/hedron_compile_commands:workspace.bzl", hedron_compile_commands_workspace = "repo")
+hedron_compile_commands_workspace()
 
-
-# Hedron's Compile Commands Extractor for Bazel
-# https://github.com/hedronvision/bazel-compile-commands-extractor
-http_archive(
-    name = "hedron_compile_commands",
-
-    # Replace the commit hash (0e990032f3c5a866e72615cf67e5ce22186dcb97) in both places (below) with the latest (https://github.com/hedronvision/bazel-compile-commands-extractor/commits/main), rather than using the stale one here.
-    # Even better, set up Renovate and let it do the work for you (see "Suggestion: Updates" in the README).
-    url = "https://github.com/hedronvision/bazel-compile-commands-extractor/archive/4f28899228fb3ad0126897876f147ca15026151e.tar.gz",
-    strip_prefix = "bazel-compile-commands-extractor-4f28899228fb3ad0126897876f147ca15026151e",
-    # When you first run this tool, it'll recommend a sha256 hash to put here with a message like: "DEBUG: Rule 'hedron_compile_commands' indicated that a canonical reproducible form can be obtained by modifying arguments sha256 = ..."
-)
 load("@hedron_compile_commands//:workspace_setup.bzl", "hedron_compile_commands_setup")
-hedron_compile_commands_setup()
 load("@hedron_compile_commands//:workspace_setup_transitive.bzl", "hedron_compile_commands_setup_transitive")
-hedron_compile_commands_setup_transitive()
 load("@hedron_compile_commands//:workspace_setup_transitive_transitive.bzl", "hedron_compile_commands_setup_transitive_transitive")
-hedron_compile_commands_setup_transitive_transitive()
 load("@hedron_compile_commands//:workspace_setup_transitive_transitive_transitive.bzl", "hedron_compile_commands_setup_transitive_transitive_transitive")
+
+hedron_compile_commands_setup()
+hedron_compile_commands_setup_transitive()
+hedron_compile_commands_setup_transitive_transitive()
 hedron_compile_commands_setup_transitive_transitive_transitive()
 
-load("//:workspace.bzl", "JAX_COMMIT", "JAX_SHA256", "ENZYME_COMMIT", "ENZYME_SHA256", "XLA_PATCHES")
-load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
-
-# http_archive(
-#     name = "rules_cc",
-#     sha256 = "a281b09e003a119bda2cea98c590129b783cc241fdd1a1de3baa656117e6cc3f"
-#     strip_prefix = "rules_cc-34f0e1f038cff9bf28d65d101c806b48f35bf92e",
-#     urls = [
-#         "https://github.com/bazelbuild/rules_cc/archive/34f0e1f038cff9bf28d65d101c806b48f35bf92e.tar.gz",
-#     ],
-# )
-# 
-# load("@rules_cc//cc:repositories.bzl", "rules_cc_dependencies")
-# 
-# rules_cc_dependencies()
-
-LLVM_TARGETS = ["X86", "AArch64", "AMDGPU", "NVPTX"]
-
-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 = ["//:patches/jax.patch"],
-)
+load("//third_party/jax:workspace.bzl", jax_workspace = "repo")
+jax_workspace()
 
-load("@jax//third_party/xla:workspace.bzl", "XLA_COMMIT", "XLA_SHA256")
-http_archive(
-    name = "xla",
-    sha256 = XLA_SHA256,
-    strip_prefix = "xla-" + XLA_COMMIT,
-    urls = ["https://github.com/wsmoses/xla/archive/{commit}.tar.gz".format(commit = XLA_COMMIT)],
-    patch_cmds = XLA_PATCHES,
-)
+load("//third_party/xla:workspace.bzl", xla_workspace = "repo")
+xla_workspace()
+
+load("//third_party/enzyme:workspace.bzl", enzyme_workspace = "repo")
+enzyme_workspace()
 
 load("@xla//third_party/py:python_init_rules.bzl", "python_init_rules")
 python_init_rules()
@@ -82,18 +45,15 @@ python_init_pip()
 load("@pypi//:requirements.bzl", "install_deps")
 install_deps()
 
-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)],
-)
-
 load("@xla//third_party/llvm:workspace.bzl", llvm = "repo")
+load("//:workspace.bzl", "LLVM_TARGETS")
 llvm("llvm-raw")
 load("@llvm-raw//utils/bazel:configure.bzl", "llvm_configure")
 llvm_configure(name = "llvm-project", targets = LLVM_TARGETS)
 
+load("@jax//third_party/flatbuffers:workspace.bzl", flatbuffers = "repo")
+flatbuffers()
+
 load("@xla//:workspace4.bzl", "xla_workspace4")
 xla_workspace4()
 
@@ -109,9 +69,6 @@ xla_workspace1()
 load("@xla//:workspace0.bzl", "xla_workspace0")
 xla_workspace0()
 
-load("@jax//third_party/flatbuffers:workspace.bzl", flatbuffers = "repo")
-flatbuffers()
-
 load(
     "@tsl//third_party/gpus/cuda/hermetic:cuda_json_init_repository.bzl",
     "cuda_json_init_repository",

third_party/enzyme/workspace.bzl

@@ -0,0 +1,12 @@
+"""Loads Enzyme."""
+
+load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
+load("//:workspace.bzl", "ENZYME_COMMIT", "ENZYME_SHA256")
+
+def repo():
+    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)],
+    )

third_party/hedron_compile_commands/workspace.bzl

@@ -0,0 +1,15 @@
+"""Loads Hedron's Compile Commands Extractor for Bazel."""
+# https://github.com/hedronvision/bazel-compile-commands-extractor
+
+load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
+
+def repo():
+    http_archive(
+        name = "hedron_compile_commands",
+
+        # Replace the commit hash (0e990032f3c5a866e72615cf67e5ce22186dcb97) in both places (below) with the latest (https://github.com/hedronvision/bazel-compile-commands-extractor/commits/main), rather than using the stale one here.
+        # Even better, set up Renovate and let it do the work for you (see "Suggestion: Updates" in the README).
+        url = "https://github.com/hedronvision/bazel-compile-commands-extractor/archive/4f28899228fb3ad0126897876f147ca15026151e.tar.gz",
+        strip_prefix = "bazel-compile-commands-extractor-4f28899228fb3ad0126897876f147ca15026151e",
+        # When you first run this tool, it'll recommend a sha256 hash to put here with a message like: "DEBUG: Rule 'hedron_compile_commands' indicated that a canonical reproducible form can be obtained by modifying arguments sha256 = ..."
+    )

third_party/jax/workspace.bzl

@@ -0,0 +1,14 @@
+"""Loads JAX."""
+
+load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
+load("//:workspace.bzl", "JAX_COMMIT", "JAX_SHA256")
+
+def repo():
+    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 = ["//:patches/jax.patch"],
+    )

third_party/xla/workspace.bzl

@@ -0,0 +1,14 @@
+"""Loads XLA."""
+
+load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
+load("@jax//third_party/xla:workspace.bzl", "XLA_COMMIT", "XLA_SHA256")
+load("//:workspace.bzl", "XLA_PATCHES")
+
+def repo():
+    http_archive(
+        name = "xla",
+        sha256 = XLA_SHA256,
+        strip_prefix = "xla-" + XLA_COMMIT,
+        urls = ["https://github.com/wsmoses/xla/archive/{commit}.tar.gz".format(commit = XLA_COMMIT)],
+        patch_cmds = XLA_PATCHES,
+    )

workspace.bzl

@@ -30,3 +30,5 @@ XLA_PATCHES = [
     "find . -type f -name BUILD -exec sed -i.bak2 's/\\/\\/xla\\/mlir\\/memref:friends/\\/\\/visibility:public/g' {} +",
     "find xla/mlir -type f -name BUILD -exec sed -i.bak3 's/\\/\\/xla:internal/\\/\\/\\/\\/visibility:public/g' {} +"
 ]
+
+LLVM_TARGETS = ["X86", "AArch64", "AMDGPU", "NVPTX"]