Migrate workspace and XLA module definitions to Bazel 8, updating MODULE.bazel files, BUILD rules, and related migration patches.

2024-02-12 12:43:23 +00:00 · 2024-02-12 12:43:23 +00:00 · 169a24307c
commit 169a24307c
parent 7e6103d876
12 changed files with 481 additions and 4690 deletions
--- a/MODULE.bazel
+++ b/MODULE.bazel
@ -5,15 +5,15 @@ module(
 new_git_repository = use_repo_rule("@bazel_tools//tools/build_defs/repo:git.bzl", "new_git_repository")

 bazel_dep(name = "bazel_skylib", version = "1.7.1")
-bazel_dep(name = "hermetic_cc_toolchain", version = "3.1.0")
+bazel_dep(name = "hermetic_cc_toolchain", version = "3.1.1")
 bazel_dep(name = "patchelf", version = "0.18.0")
 bazel_dep(name = "platforms", version = "0.0.10")
-bazel_dep(name = "rules_cc", version = "0.0.9")
+bazel_dep(name = "rules_cc", version = "0.0.17")
 bazel_dep(name = "rules_pkg", version = "1.0.1")
-bazel_dep(name = "rules_proto", version = "6.0.2")
-bazel_dep(name = "with_cfg.bzl", version = "0.5.0")
-bazel_dep(name = "buildifier_prebuilt", version = "6.4.0", dev_dependency = True)
-bazel_dep(name = "aspect_bazel_lib", version = "2.9.3")
+bazel_dep(name = "rules_proto", version = "7.1.0")
+bazel_dep(name = "with_cfg.bzl", version = "0.8.0")
+bazel_dep(name = "buildifier_prebuilt", version = "7.3.1", dev_dependency = True)
+bazel_dep(name = "aspect_bazel_lib", version = "2.11.0")
 bazel_dep(name = "rules_distroless", version = "0.3.8")

 bazel_lib_toolchains = use_extension("@aspect_bazel_lib//lib:extensions.bzl", "toolchains", dev_dependency = True)
@ -41,9 +41,9 @@ register_toolchains(
    "@zig_sdk//toolchain:linux_arm64_gnu.2.31",
 )

-bazel_dep(name = "rules_python", version = "0.37.0")
+bazel_dep(name = "rules_python", version = "0.40.0")
 bazel_dep(name = "rules_uv", version = "0.34.0")
-bazel_dep(name = "aspect_rules_py", version = "1.0.0-rc0")
+bazel_dep(name = "aspect_rules_py", version = "1.0.0")
 python = use_extension("@rules_python//python/extensions:python.bzl", "python")
 python.toolchain(python_version = "3.11")
 use_repo(python, "python_versions")
@ -98,7 +98,7 @@ llvm.configure(
 use_repo(llvm, "llvm-project")

 bazel_dep(name = "stablehlo", version = "20250117.0-c125b32")
-bazel_dep(name = "xla", version = "20250122.0-cc075be")
+bazel_dep(name = "xla", version = "20250123.0-cc075be")

 tsl = use_extension("@xla//:tsl.bzl", "tsl")
 use_repo(tsl, "tsl")
--- a/MODULE.bazel.lock
+++ b/MODULE.bazel.lock
--- a/runtimes/rocm/rocm.bzl
+++ b/runtimes/rocm/rocm.bzl
@ -106,10 +106,7 @@ _PACKAGES = {
        _CC_IMPORT_TPL.format(name = "MIOpen", shared_library = "libMIOpen.so.1"),
        _RUNFILES_TPL.format(name = "runfiles", glob = repr(["share/miopen/**"])),
    ]),
-    "rccl": "".join([
-        _CC_IMPORT_TPL.format(name = "rccl", shared_library = "librccl.so.1"),
-        _RUNFILES_TPL.format(name = "runfiles", glob = repr(["share/rccl/msccl-algorithms/**"])),
-    ]),
+    "rccl": _CC_IMPORT_TPL.format(name = "rccl", shared_library = "librccl.so.1"),
    "rocm-device-libs": _RUNFILES_TPL.format(name = "runfiles", glob = repr(["amdgcn/**"])),
    "hip-dev": _RUNFILES_TPL.format(name = "runfiles", glob = repr(["share/**"])),
    "rocblas": """\
--- a/third_party/modules/xla/20250123.0-cc075be/MODULE.bazel
+++ b/third_party/modules/xla/20250123.0-cc075be/MODULE.bazel
@ -0,0 +1,34 @@
+module(
+    name = "xla",
+    version = "20250123.0-cc075be",
+    compatibility_level = 1,
+)
+
+bazel_dep(name = "platforms", version = "0.0.8")
+bazel_dep(name = "bazel_skylib", version = "1.5.0")
+bazel_dep(name = "rules_cc", version = "0.0.17")
+bazel_dep(name = "rules_apple", version = "3.2.1", repo_name = "build_bazel_rules_apple")
+bazel_dep(name = "abseil-cpp", version = "20240116.0", repo_name = "com_google_absl")
+bazel_dep(name = "rules_python", version = "0.29.0")
+bazel_dep(name = "rules_proto", version = "6.0.0-rc1")
+bazel_dep(name = "rules_java", version = "7.3.2")
+bazel_dep(name = "rules_pkg", version = "0.9.1")
+bazel_dep(name = "zlib", version = "1.2.13")
+bazel_dep(name = "re2", version = "2024-02-01", repo_name = "com_googlesource_code_re2")
+bazel_dep(name = "rules_license", version = "0.0.8")
+
+bazel_dep(name = "stablehlo", version = "20250117.0-c125b32")
+
+tsl = use_extension("//:tsl.bzl", "tsl")
+use_repo(tsl, "tsl")
+
+xla_workspace = use_extension("//:workspace.bzl", "xla_workspace")
+use_repo(
+    xla_workspace,
+    "com_github_grpc_grpc",
+    "com_google_protobuf",
+    "local_config_cuda",
+    "local_config_remote_execution",
+    "local_config_rocm",
+    "local_config_tensorrt",
+)
--- a/third_party/modules/xla/20250123.0-cc075be/overlay/MODULE.bazel
+++ b/third_party/modules/xla/20250123.0-cc075be/overlay/MODULE.bazel
@ -0,0 +1,34 @@
+module(
+    name = "xla",
+    version = "20250123.0-cc075be",
+    compatibility_level = 1,
+)
+
+bazel_dep(name = "platforms", version = "0.0.8")
+bazel_dep(name = "bazel_skylib", version = "1.5.0")
+bazel_dep(name = "rules_cc", version = "0.0.17")
+bazel_dep(name = "rules_apple", version = "3.2.1", repo_name = "build_bazel_rules_apple")
+bazel_dep(name = "abseil-cpp", version = "20240116.0", repo_name = "com_google_absl")
+bazel_dep(name = "rules_python", version = "0.29.0")
+bazel_dep(name = "rules_proto", version = "6.0.0-rc1")
+bazel_dep(name = "rules_java", version = "7.3.2")
+bazel_dep(name = "rules_pkg", version = "0.9.1")
+bazel_dep(name = "zlib", version = "1.2.13")
+bazel_dep(name = "re2", version = "2024-02-01", repo_name = "com_googlesource_code_re2")
+bazel_dep(name = "rules_license", version = "0.0.8")
+
+bazel_dep(name = "stablehlo", version = "20250117.0-c125b32")
+
+tsl = use_extension("//:tsl.bzl", "tsl")
+use_repo(tsl, "tsl")
+
+xla_workspace = use_extension("//:workspace.bzl", "xla_workspace")
+use_repo(
+    xla_workspace,
+    "com_github_grpc_grpc",
+    "com_google_protobuf",
+    "local_config_cuda",
+    "local_config_remote_execution",
+    "local_config_rocm",
+    "local_config_tensorrt",
+)
--- a/third_party/modules/xla/20250123.0-cc075be/overlay/tsl.bzl
+++ b/third_party/modules/xla/20250123.0-cc075be/overlay/tsl.bzl
@ -0,0 +1,19 @@
+load("//third_party:repo.bzl", "tf_vendored")
+load("//third_party/py:python_init_repositories.bzl", "python_init_repositories")
+
+def _tsl_impl(mctx):
+    python_init_repositories(
+        requirements = {
+            "3.11": "//:requirements_lock_3_11.txt",
+        },
+    )
+    tf_vendored(name = "tsl", relpath = "third_party/tsl")
+    return mctx.extension_metadata(
+        reproducible = True,
+        root_module_direct_deps = ["tsl"],
+        root_module_direct_dev_deps = [],
+    )
+
+tsl = module_extension(
+    implementation = _tsl_impl,
+)
--- a/third_party/modules/xla/20250123.0-cc075be/overlay/workspace.bzl
+++ b/third_party/modules/xla/20250123.0-cc075be/overlay/workspace.bzl
@ -0,0 +1,52 @@
+load("@tsl//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls")
+load("@tsl//third_party/gpus:cuda_configure.bzl", "cuda_configure")
+load("@tsl//third_party/gpus:rocm_configure.bzl", "rocm_configure")
+load("@tsl//third_party/tensorrt:tensorrt_configure.bzl", "tensorrt_configure")
+load("@tsl//tools/toolchains/remote:configure.bzl", "remote_execution_configure")
+
+def _xla_workspace_impl(mctx):
+    cuda_configure(name = "local_config_cuda")
+    remote_execution_configure(name = "local_config_remote_execution")
+    rocm_configure(name = "local_config_rocm")
+    tensorrt_configure(name = "local_config_tensorrt")
+    tf_http_archive(
+        name = "com_github_grpc_grpc",
+        sha256 = "b956598d8cbe168b5ee717b5dafa56563eb5201a947856a6688bbeac9cac4e1f",
+        strip_prefix = "grpc-b54a5b338637f92bfcf4b0bc05e0f57a5fd8fadd",
+        system_build_file = "@tsl//third_party/systemlibs:grpc.BUILD",
+        patch_file = [
+            "@tsl//third_party/grpc:generate_cc_env_fix.patch",
+            "@tsl//third_party/grpc:register_go_toolchain.patch",
+        ],
+        system_link_files = {
+            "@tsl//third_party/systemlibs:BUILD": "bazel/BUILD",
+            "@tsl//third_party/systemlibs:grpc.BUILD": "src/compiler/BUILD",
+            "@tsl//third_party/systemlibs:grpc.bazel.grpc_deps.bzl": "bazel/grpc_deps.bzl",
+            "@tsl//third_party/systemlibs:grpc.bazel.grpc_extra_deps.bzl": "bazel/grpc_extra_deps.bzl",
+            "@tsl//third_party/systemlibs:grpc.bazel.cc_grpc_library.bzl": "bazel/cc_grpc_library.bzl",
+            "@tsl//third_party/systemlibs:grpc.bazel.generate_cc.bzl": "bazel/generate_cc.bzl",
+            "@tsl//third_party/systemlibs:grpc.bazel.protobuf.bzl": "bazel/protobuf.bzl",
+        },
+        urls = tf_mirror_urls("https://github.com/grpc/grpc/archive/b54a5b338637f92bfcf4b0bc05e0f57a5fd8fadd.tar.gz"),
+    )
+    tf_http_archive(
+        name = "com_google_protobuf",
+        patch_file = ["@tsl//third_party/protobuf:protobuf.patch"],
+        sha256 = "f66073dee0bc159157b0bd7f502d7d1ee0bc76b3c1eac9836927511bdc4b3fc1",
+        strip_prefix = "protobuf-3.21.9",
+        system_build_file = "@tsl//third_party/systemlibs:protobuf.BUILD",
+        system_link_files = {
+            "@tsl//third_party/systemlibs:protobuf.bzl": "protobuf.bzl",
+            "@tsl//third_party/systemlibs:protobuf_deps.bzl": "protobuf_deps.bzl",
+        },
+        urls = tf_mirror_urls("https://github.com/protocolbuffers/protobuf/archive/v3.21.9.zip"),
+    )
+    return mctx.extension_metadata(
+        reproducible = True,
+        root_module_direct_deps = "all",
+        root_module_direct_dev_deps = [],
+    )
+
+xla_workspace = module_extension(
+    implementation = _xla_workspace_impl,
+)
--- a/third_party/modules/xla/20250123.0-cc075be/patches/0001-bazel-migration-to-bazel-8.0.1.patch
+++ b/third_party/modules/xla/20250123.0-cc075be/patches/0001-bazel-migration-to-bazel-8.0.1.patch
@ -0,0 +1,42 @@
+From 6cf475b500521c1b8be06f590fdbc1818f0dc44b Mon Sep 17 00:00:00 2001
+From: Jean-Baptiste Dalido <jb@zml.ai>
+Date: Mon, 6 Jan 2025 13:33:13 +0100
+Subject: [PATCH] bazel: migration to bazel 8.0.1
+
+---
+ .bazelversion                                       | 2 +-
+ third_party/tsl/third_party/gpus/cuda_configure.bzl | 4 ++--
+ 2 files changed, 3 insertions(+), 3 deletions(-)
+
+diff --git a/.bazelversion b/.bazelversion
+index f22d756da3..fa5fce04b3 100644
+--- a/.bazelversion
+++ b/.bazelversion
+@@ -1 +1 @@
+-6.5.0
+8.0.1
+\ No newline at end of file
+diff --git a/third_party/tsl/third_party/gpus/cuda_configure.bzl b/third_party/tsl/third_party/gpus/cuda_configure.bzl
+index d62531152d..71d80a5a99 100644
+--- a/third_party/tsl/third_party/gpus/cuda_configure.bzl
+++ b/third_party/tsl/third_party/gpus/cuda_configure.bzl
+@@ -33,14 +33,14 @@ NB: DEPRECATED! Use `hermetic/cuda_configure` rule instead.
+ load(
+     "@bazel_tools//tools/cpp:lib_cc_configure.bzl",
+     "escape_string",
+-    "get_env_var",
+ )
+ load(
+     "@bazel_tools//tools/cpp:windows_cc_configure.bzl",
+-    "find_msvc_tool",
+     "find_vc_path",
+     "setup_vc_env_vars",
+ )
+load("@rules_cc//cc/private/toolchain:windows_cc_configure.bzl", "find_msvc_tool")
+load("@rules_cc//cc/private/toolchain:lib_cc_configure.bzl", "get_env_var")
+ load("//third_party/clang_toolchain:download_clang.bzl", "download_clang")
+ load(
+     "//third_party/remote_config:common.bzl",
+-- 
+2.39.3 (Apple Git-146)
+
--- a/third_party/modules/xla/20250123.0-cc075be/patches/0003-PJRT-C-API-Ensure-C-compliance-for-Profiler-Extensio.patch
+++ b/third_party/modules/xla/20250123.0-cc075be/patches/0003-PJRT-C-API-Ensure-C-compliance-for-Profiler-Extensio.patch
@ -0,0 +1,27 @@
+From 4db5de34f70d991fedbe28915c8239b97ba7a064 Mon Sep 17 00:00:00 2001
+From: Steeve Morin <steeve.morin@gmail.com>
+Date: Mon, 18 Mar 2024 17:17:34 +0100
+Subject: [PATCH 3/3] [PJRT C API] Ensure C compliance for Profiler Extension
+
+---
+ xla/pjrt/c/pjrt_c_api_profiler_extension.h | 2 ++
+ 1 file changed, 2 insertions(+)
+
+diff --git a/xla/pjrt/c/pjrt_c_api_profiler_extension.h b/xla/pjrt/c/pjrt_c_api_profiler_extension.h
+index c821916ad..89a596123 100644
+--- a/xla/pjrt/c/pjrt_c_api_profiler_extension.h
+++ b/xla/pjrt/c/pjrt_c_api_profiler_extension.h
+@@ -16,8 +16,10 @@ limitations under the License.
+ #ifndef XLA_PJRT_C_PJRT_C_API_PROFILER_EXTENSION_H_
+ #define XLA_PJRT_C_PJRT_C_API_PROFILER_EXTENSION_H_
+ 
+#ifdef __cplusplus
+ #include <cstddef>
+ #include <cstdint>
+#endif
+ 
+ #include "xla/backends/profiler/plugin/profiler_c_api.h"
+ #include "xla/pjrt/c/pjrt_c_api.h"
+-- 
+2.39.3 (Apple Git-146)
+
--- a/third_party/modules/xla/20250123.0-cc075be/source.json
+++ b/third_party/modules/xla/20250123.0-cc075be/source.json
@ -0,0 +1,15 @@
+{
+    "strip_prefix": "xla-cc075beb6148c2777da2b6749c63830856ee6c2a",
+    "url": "https://github.com/openxla/xla/archive/cc075beb6148c2777da2b6749c63830856ee6c2a.tar.gz",
+    "integrity": "sha256-oB8S38WZKEXBtZ6rARd0oL9SAtSmfsPM2xTuj3ylexc=",
+    "overlay": {
+        "tsl.bzl": "",
+        "workspace.bzl": "",
+        "MODULE.bazel": ""
+    },
+    "patch_strip": 1,
+    "patches": {
+        "0001-bazel-migration-to-bazel-8.0.1.patch": "",
+        "0003-PJRT-C-API-Ensure-C-compliance-for-Profiler-Extensio.patch": ""
+    }
+}
--- a/third_party/modules/xla/metadata.json
+++ b/third_party/modules/xla/metadata.json
@ -16,6 +16,7 @@
        "20241025.0-4663f04",
        "20250103.0-5f1fe6a",
        "20250122.0-cc075be",
+        "20250123.0-cc075be",
    ],
    "yanked_versions": {}
 }
--- a/zml/BUILD.bazel
+++ b/zml/BUILD.bazel
@ -14,7 +14,6 @@ zig_library(
        "*.zig",
        "aio/**/*.zig",
        "nn/**/*.zig",
-        "ops/**/*.zig",
        # TODO: test_runner.zig should not be here.
        # It's here for now because it seems that test_runner property in zig_test is misbehaving.
        # See https://github.com/zml/rules_zig/issues/2