gomlx · janpfeifer · Sep 28, 2024 · Sep 28, 2024 · Sep 28, 2024
diff --git a/c/BUILD b/c/BUILD
@@ -19,7 +19,7 @@ genrule(
 )
 
 genrule(
-    name = "gomlx_xlabuilder_lib",
+    name = "gomlx_xlabuilder_dynamic_lib",
     srcs = [
         "//gomlx/xlabuilder:libgomlx_xlabuilder.so",
     ],
@@ -30,13 +30,27 @@ genrule(
     """,
 )
 
+# Waiting for cc_static_library.
+#genrule(
+#    name = "gomlx_xlabuilder_static_lib",
+#    srcs = [
+#        "//gomlx/xlabuilder:libgomlx_xlabuilder.a",
+#    ],
+#    outs = ["lib/libgomlx_xlabuilder.a"],
+#    cmd = """
+#    mkdir -p "$(@D)"
+#    ln $(location //gomlx/xlabuilder:libgomlx_xlabuilder.a) "$@"
+#    """,
+#)
+
 # Tar packaging: notice we force date/time and ordering to be deterministic, to generate the
 # always same archive, if the contents are the same. See https://reproducible-builds.org/docs/archives/.
 genrule(
     name = "gomlx_xlabuilder",
     srcs = [
         ":gomlx_xlabuilder_headers_include",
-        ":gomlx_xlabuilder_lib",
+        ":gomlx_xlabuilder_dynamic_lib",
+ #       ":gomlx_xlabuilder_static_lib",
     ],
     outs = ["gomlx_xlabuilder.tar.gz"],
     cmd_bash = """

diff --git a/c/WORKSPACE b/c/WORKSPACE
@@ -24,29 +24,30 @@ http_archive(
 # = "2fc2bc05cc8aa04c0902555ab2f9f817c41886c3"  # From 2023-07-10
 # = "4f8381651977dff16b1d86bb4b198eb733c5f478"  # From 2023-10-22
 # = "5c3f21721d921751f3c64451f9be3b59f9cac8c7" # From 2024-04-16
-OPENXLA_XLA_COMMIT_HASH = "782b73d746ab5ad6698025f91a349c21b2eceba0"  # From 2024-06-15
-
+# OPENXLA_XLA_COMMIT_HASH = "782b73d746ab5ad6698025f91a349c21b2eceba0"  # From 2024-06-15
+OPENXLA_XLA_COMMIT_HASH = "0f30f33b7b7456cc6778fa47e57fc518d54fee3d"  # From 2024-09-28
 http_archive(
     name = "xla",
     # sha256 = "c86b68ce16b08ba1287c641b9790f5ba7944ebbfe7218710f11827c2a65569f4",  # From 2023-03-30
     # sha256 = "e05e41cd1d968d0bbacf32f9c2ccee6b5f783a7597b6e6668756953282cf90b9",  # From 2023-04-28
     # sha256 = "a2c905c02d2e9fa44109d163b7e8adedbb56b2f241a2bd95a962d549834f161e",  # From 2023-07-10
     # sha256 = "980e86702bd5a7722ec3ccd2412714ab62dbf0e2d58251af86ec60e1d832579c",  # From 2023-10-22
     # sha256 = "dd5f731b3de37c798bbd2f91ce7a5b99793f5ec07e75cfd680ecb3b9036c285b",  # From 2024-04-16
-    sha256 = "22747c91ac1154196d463f971b91cfe91f85c99a6e7e9a49c1c8c28a6c2d1729",  # From 2024-06-15
+    # sha256 = "22747c91ac1154196d463f971b91cfe91f85c99a6e7e9a49c1c8c28a6c2d1729",  # From 2024-06-15
+    sha256 = "14c80eac39f992d23575c1558acfc005799f2235e2286ee1629180389f490951",  # From 2024-09-28
     strip_prefix = "xla-" + OPENXLA_XLA_COMMIT_HASH,
     urls = [
         "https://github.com/openxla/xla/archive/{hash}.zip".format(hash = OPENXLA_XLA_COMMIT_HASH),
     ],
 )
 
 # Create @tsl repository.
-load("@xla//third_party:repo.bzl", "tf_vendored")
+#load("@xla//third_party:repo.bzl", "tf_vendored")
 
-tf_vendored(
-    name = "tsl",
-    relpath = "third_party/tsl",
-)
+#tf_vendored(
+#    name = "tsl",
+#    relpath = "third_party/tsl",
+#)
 
 # Initialize the XLA repository and all dependencies.
 
@@ -98,3 +99,50 @@ xla_workspace1()
 load("@xla//:workspace0.bzl", "xla_workspace0")
 
 xla_workspace0()
+
+load(
+    "@tsl//third_party/gpus/cuda/hermetic:cuda_json_init_repository.bzl",
+    "cuda_json_init_repository",
+)
+
+cuda_json_init_repository()
+
+load(
+    "@cuda_redist_json//:distributions.bzl",
+    "CUDA_REDISTRIBUTIONS",
+    "CUDNN_REDISTRIBUTIONS",
+)
+load(
+    "@tsl//third_party/gpus/cuda/hermetic:cuda_redist_init_repositories.bzl",
+    "cuda_redist_init_repositories",
+    "cudnn_redist_init_repository",
+)
+
+cuda_redist_init_repositories(
+    cuda_redistributions = CUDA_REDISTRIBUTIONS,
+)
+
+cudnn_redist_init_repository(
+    cudnn_redistributions = CUDNN_REDISTRIBUTIONS,
+)
+
+load(
+    "@tsl//third_party/gpus/cuda/hermetic:cuda_configure.bzl",
+    "cuda_configure",
+)
+
+cuda_configure(name = "local_config_cuda")
+
+load(
+    "@tsl//third_party/nccl/hermetic:nccl_redist_init_repository.bzl",
+    "nccl_redist_init_repository",
+)
+
+nccl_redist_init_repository()
+
+load(
+    "@tsl//third_party/nccl/hermetic:nccl_configure.bzl",
+    "nccl_configure",
+)
+
+nccl_configure(name = "local_config_nccl")
diff --git a/c/bazel.sh b/c/bazel.sh
@@ -61,14 +61,15 @@ PYTHON=${PYTHON:-python}  # Python, version should be > 3.7.
 # Check the OpenXLA version (commit hash) changed, and if so, download an
 # updated `openxla_xla_bazelrc` file from github.
 # TODO: include a sha256 verification of the file as well.
-if ! grep -q "OPENXLA_XLA_COMMIT_HASH" WORKSPACE ; then
+if ! egrep -q "^OPENXLA_XLA_COMMIT_HASH =" WORKSPACE ; then
   echo "Did not find OPENXLA_XLA_COMMIT_HASH in WORKSPACE file!?"
   exit 1
 fi
 OPENXLA_XLA_COMMIT_HASH="$(
-  grep -E "^OPENXLA_XLA_COMMIT_HASH[[:space:]]*=" WORKSPACE |\
+  grep -E "^OPENXLA_XLA_COMMIT_HASH[[:space:]]*=[[:space:]]*" WORKSPACE |\
     sed -n 's/^[^"]*"\([^"]*\)".*/\1/p'
 )"
+printf "OPENXLA_XLA_COMMIT_HASH=%s\n" "${OPENXLA_XLA_COMMIT_HASH}"
 OPENXLA_BAZELRC="openxla_xla_bazelrc"
 if [[ ! -e "${OPENXLA_BAZELRC}" || ! -e "${OPENXLA_BAZELRC}.version" \
   || "$(< "${OPENXLA_BAZELRC}.version")" != "${OPENXLA_XLA_COMMIT_HASH}" ]] ; then
@@ -79,6 +80,12 @@ else
     echo "File ${OPENXLA_BAZELRC} at version \"${OPENXLA_XLA_COMMIT_HASH}\" already exists, not fetching."
 fi
 
+# Attempts of enabling `cc_static_library`:
+# See https://github.com/bazelbuild/bazel/issues/1920
+# export USE_BAZEL_VERSION=last_green
+# STARTUP_FLAGS="${STARTUP_FLAGS} --experimental_cc_static_library"
+# Presumably, it will make to Bazel 7.4.0
+
 STARTUP_FLAGS="${STARTUP_FLAGS} ${OUTPUT_DIR}"
 STARTUP_FLAGS="${STARTUP_FLAGS} --bazelrc=${OPENXLA_BAZELRC}"
 STARTUP_FLAGS="${STARTUP_FLAGS} --bazelrc=xla_configure.bazelrc"

diff --git a/c/gomlx/xlabuilder/BUILD b/c/gomlx/xlabuilder/BUILD
@@ -16,9 +16,9 @@ gomlx_xlabuilder_deps = [
     "@xla//xla:util",
     "@xla//xla:xla_data_proto_cc",
     "@xla//xla:xla_proto_cc",
-    "@xla//xla/client:xla_builder",
-    "@xla//xla/client:xla_computation",
-    "@xla//xla/client/lib:math",
+    "@xla//xla/hlo/builder:xla_builder",
+    "@xla//xla/hlo/builder:xla_computation",
+    "@xla//xla/hlo/builder/lib:math",
 ]
 
 cc_library(
@@ -50,6 +50,14 @@ cc_binary(
     ],
 )
 
+# Waiting for Bazel 7.4.0:
+#cc_static_library(
+#    name = "libgomlx_xlabuilder.a",
+#    deps = gomlx_xlabuilder_deps + [
+#        ":xlabuilder",
+#    ],
+#)
+
 filegroup(
     name = "headers",
     srcs = glob(["*.h"]),

diff --git a/c/gomlx/xlabuilder/literal.h b/c/gomlx/xlabuilder/literal.h
@@ -28,7 +28,7 @@
 
 #ifdef __cplusplus
 // C++ only includes: these are not seen by the Go compiler.
-#include "xla/client/xla_builder.h"
+#include "xla/hlo/builder/xla_builder.h"
 #include "xla/shape.h"
 
 struct Literal;

diff --git a/c/gomlx/xlabuilder/op.h b/c/gomlx/xlabuilder/op.h
@@ -32,8 +32,8 @@
 
 #ifdef __cplusplus
 // C++ only includes: these are not seen by the Go compiler.
-#include "xla/client/xla_builder.h"
-#include "xla/client/xla_computation.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/hlo/builder/xla_computation.h"
 #include "xla/shape.h"
 
 typedef xla::XlaOp XlaOp;

diff --git a/c/gomlx/xlabuilder/xlabuilder.cpp b/c/gomlx/xlabuilder/xlabuilder.cpp
@@ -26,9 +26,9 @@
 #include "gomlx/xlabuilder/op.h"
 #include "gomlx/xlabuilder/gen_op_types.h"
 
-#include "xla/client/xla_builder.h"
-#include "xla/client/xla_computation.h"
-#include "xla/client/lib/arithmetic.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/hlo/builder/xla_computation.h"
+#include "xla/hlo/builder/lib/arithmetic.h"
 #include "xla/literal.h"
 #include "xla/statusor.h"
 #include "xla/types.h"

diff --git a/c/gomlx/xlabuilder/xlabuilder.h b/c/gomlx/xlabuilder/xlabuilder.h
@@ -29,8 +29,8 @@
 
 #ifdef __cplusplus
 // C++ only: Dependencies.
-#include "xla/client/xla_builder.h"
-#include "xla/client/xla_computation.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/hlo/builder/xla_computation.h"
 
 // Aliases.
 typedef xla::XlaBuilder XlaBuilder;