Update XLA revision

.github/workflows/release.yml

@@ -4,14 +4,17 @@ on:
     tags:
       - "v*.*.*"
 
+env:
+  USE_BAZEL_VERSION: 7.4.1
+
 jobs:
   create_draft_release:
     if: github.ref_type == 'tag'
     permissions:
       contents: write
-    runs-on: ubuntu-20.04
+    runs-on: ubuntu-22.04
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
       - name: Create draft release
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
@@ -23,83 +26,84 @@ jobs:
   linux:
     name: "x86_64-linux-gnu-{cpu,tpu}"
     needs: [create_draft_release]
-    # We intentionally build on ubuntu 20 to compile against
+    # We intentionally build on ubuntu 22 to compile against
     # an older version of glibc
-    runs-on: ubuntu-20.04
+    runs-on: ubuntu-22.04
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
       - uses: erlef/setup-beam@v1
         with:
           otp-version: "24"
           elixir-version: "1.15.8"
       # Setup the compilation environment
-      - uses: abhinavsingh/setup-bazel@v3
-        with:
-          version: "6.5.0"
-      - uses: actions/setup-python@v2
+      - uses: bazel-contrib/[email protected]
+      - uses: actions/setup-python@v5
         with:
-          python-version: "3.9"
+          python-version: "3.11"
       - run: python -m pip install --upgrade pip numpy
+      - run: |
+          clang_version="18"
+          sudo apt-get update
+          sudo apt-get install -y wget gnupg software-properties-common lsb-release
+          wget -qO- https://apt.llvm.org/llvm.sh | sudo bash -s -- $clang_version
+          sudo update-alternatives --install /usr/bin/clang clang /usr/bin/clang-$clang_version 200
+          sudo update-alternatives --install /usr/bin/clang++ clang++ /usr/bin/clang++-$clang_version 200
       # Build and upload the archives
       - run: mix deps.get
       - run: .github/scripts/compile_and_upload.sh ${{ github.ref_name }}
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
           XLA_TARGET: cpu
-          CC: gcc-9
       - run: .github/scripts/compile_and_upload.sh ${{ github.ref_name }}
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
           XLA_TARGET: tpu
-          CC: gcc-9
 
   macos:
     name: "x86_64-darwin-cpu"
     needs: [create_draft_release]
-    runs-on: macos-12
+    runs-on: macos-13
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
       - run: brew install elixir
       - run: mix local.hex --force
       # Setup the compilation environment
-      - uses: abhinavsingh/setup-bazel@v3
+      - uses: bazel-contrib/[email protected]
+      - uses: actions/setup-python@v5
         with:
-          version: "6.5.0"
-      - uses: actions/setup-python@v2
-        with:
-          python-version: "3.9"
+          python-version: "3.11"
       - run: python -m pip install --upgrade pip numpy
       # Build and upload the archive
       - run: mix deps.get
       - run: .github/scripts/compile_and_upload.sh ${{ github.ref_name }}
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
           XLA_TARGET: cpu
-          CC: gcc-9
+          # This runner comes with Clang 14, which does not support the -mavxvnniint8
+          # CLI flag. We can install newer Clang, however at some point Bazel toolchains
+          # invoke xcrun clang, which always uses the system version from Xcode, ignoring
+          # whichever version we installed ourselves. With the flag below, we make sure
+          # this flag is not passed in the first place.
+          # See https://github.com/tensorflow/tensorflow/pull/87514
+          BUILD_FLAGS: "--define=xnn_enable_avxvnniint8=false"
 
   macos_arm:
-    name: "aarch64-darwin-cpu (cross-compiled)"
+    name: "aarch64-darwin-cpu"
     needs: [create_draft_release]
-    runs-on: macos-12
+    runs-on: macos-14
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
       - run: brew install elixir
       - run: mix local.hex --force
       # Setup the compilation environment
-      - uses: abhinavsingh/setup-bazel@v3
-        with:
-          version: "6.5.0"
-      - uses: actions/setup-python@v2
+      - uses: bazel-contrib/[email protected]
+      - uses: actions/setup-python@v5
         with:
-          python-version: "3.9"
+          python-version: "3.11"
       - run: python -m pip install --upgrade pip numpy
       # Build and upload the archive
       - run: mix deps.get
       - run: .github/scripts/compile_and_upload.sh ${{ github.ref_name }}
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
           XLA_TARGET: cpu
-          XLA_TARGET_PLATFORM: "aarch64-darwin"
-          # Explicitly cross-compile for arm64
-          BUILD_FLAGS: "--config=macos_arm64"
-          CC: gcc-9

.tool-versions

@@ -0,0 +1,2 @@
+bazel 7.4.1
+python 3.11.11

Makefile

@@ -9,7 +9,7 @@
 BUILD_MODE ?= opt # can also be dbg
 OPENXLA_GIT_REPO ?= https://github.com/openxla/xla.git
 
-OPENXLA_GIT_REV ?= fd58925adee147d38c25a085354e15427a12d00a
+OPENXLA_GIT_REV ?= cf7150f5755f7c5d985038410bd5c62c5a0debc2
 
 # Private configuration
 BAZEL_FLAGS = --define "framework_shared_object=false" -c $(BUILD_MODE)

README.md

@@ -51,8 +51,8 @@ In addition to building in a local environment, you can build the ROCm binary us
 the Docker-based scripts in [`builds/`](https://github.com/elixir-nx/xla/tree/main/builds). You may want to adjust the ROCm
 version in `rocm.Dockerfile` accordingly.
 
-When you encounter errors at runtime, you may want to set `ROCM_PATH=/opt/rocm-5.7.0`
-and `LD_LIBRARY_PATH="/opt/rocm-5.7.0/lib"` (with your respective version). For further
+When you encounter errors at runtime, you may want to set `ROCM_PATH=/opt/rocm-6.0.0`
+and `LD_LIBRARY_PATH="/opt/rocm-6.0.0/lib"` (with your respective version). For further
 issues, feel free to open an issue.
 
 #### `XLA_BUILD`
@@ -93,29 +93,29 @@ Keep in mind that the compilation usually takes a very long time.
 You will need the following installed in your system for the compilation:
 
   * [Git](https://git-scm.com/) for fetching XLA source
-  * [Bazel v6.5.0](https://bazel.build/) for compiling XLA
+  * [Bazel v7.4.1](https://bazel.build/) for compiling XLA
+  * [Clang 18](https://clang.llvm.org/) for compiling XLA
   * [Python3](https://python.org) with NumPy installed for compiling XLA
 
 ### Common issues
 
 #### Bazel version
 
-Use `bazel --version` to check your Bazel version, make sure you are using v6.5.0.
+Use `bazel --version` to check your Bazel version, make sure you are using v7.4.1.
 Most binaries are available on [Github](https://github.com/bazelbuild/bazel/releases),
 but it can also be installed with `asdf`:
 
 ```shell
 asdf plugin-add bazel
-asdf install bazel 6.5.0
-asdf global bazel 6.5.0
+asdf install bazel 7.4.1
+asdf global bazel 7.4.1
 ```
 
-#### GCC
+#### Clang
 
-You may have issues with newer and older versions of GCC. XLA builds are known to work
-with GCC versions between 7.5 and 9.3. If your system uses a newer GCC version, you can
-install an older version and tell Bazel to use it with `export CC=/path/to/gcc-{version}`
-where version is the GCC version you installed.
+XLA builds are known to work with Clang 18. On macOS clang comes as part of Xcode SDK
+and the version may be older, though for macOS we have precompiled archives, so you
+most likely don't need to worry about it.
 
 #### Python and asdf
 
@@ -136,7 +136,7 @@ There are two known workarounds:
    `direnv` along with the `asdf-direnv` plugin will explicitly set the paths for any binary specified
    in your project's `.tool-versions` files.
 
-If you still get the error, you can also try setting `PYTHON_BIN_PATH`, like `export PYTHON_BIN_PATH=/usr/bin/python3.9`.
+If you still get the error, you can also try setting `PYTHON_BIN_PATH`, like `export PYTHON_BIN_PATH=/usr/bin/python3.11`.
 
 After doing any of the steps above, it may be necessary to clear the build cache by removing ` ~/.cache/xla_build`
 (or the corresponding OS-specific cache location).

build.log

@@ -0,0 +1,2 @@
+ERROR: Cannot connect to the Docker daemon at unix:///Users/jonatanklosko/.docker/run/docker.sock. Is the docker daemon running?
+        0.34 real         0.08 user         0.08 sys

builds/Dockerfile

@@ -1,5 +1,5 @@
 ARG VARIANT
-ARG BASE_IMAGE="hexpm/elixir:1.15.8-erlang-24.3.4.17-ubuntu-focal-20240427"
+ARG BASE_IMAGE="hexpm/elixir:1.15.8-erlang-24.3.4.17-ubuntu-jammy-20250404"
 
 # Pre-stages for base image variants
 
@@ -42,25 +42,29 @@ ENV ROCM_PATH "/opt/rocm-${ROCM_VERSION}.0"
 FROM base-${VARIANT}
 
 # Set the missing UTF-8 locale, otherwise Elixir warns
-ENV LC_ALL C.UTF-8
+ENV LC_ALL=C.UTF-8
 
 # Make sure installing packages (like tzdata) doesn't prompt for configuration
 ARG DEBIAN_FRONTEND=noninteractive
 
 # We need to install "add-apt-repository" first
-RUN apt-get update && apt-get install -y software-properties-common && \
-  # Add repository with the latest git version
-  add-apt-repository ppa:git-core/ppa && \
+RUN apt-get update && \
   # Install basic system dependencies
   apt-get update && apt-get install -y ca-certificates curl git unzip wget && \
+  # Install Clang
+  clang_version="18" && \
+  apt-get install -y wget gnupg software-properties-common lsb-release && \
+  wget -qO- https://apt.llvm.org/llvm.sh | bash -s -- $clang_version && \
+  update-alternatives --install /usr/bin/clang clang /usr/bin/clang-$clang_version 100 && \
+  update-alternatives --install /usr/bin/clang++ clang++ /usr/bin/clang++-$clang_version 100 && \
   apt-get clean -y && rm -rf /var/lib/apt/lists/*
 
 # Install Bazel using Bazelisk (works for both amd and arm)
-RUN wget -O bazel "https://github.com/bazelbuild/bazelisk/releases/download/v1.18.0/bazelisk-linux-$(dpkg --print-architecture)" && \
+RUN wget -O bazel "https://github.com/bazelbuild/bazelisk/releases/download/v1.26.0/bazelisk-linux-$(dpkg --print-architecture)" && \
   chmod +x bazel && \
   mv bazel /usr/local/bin/bazel
 
-ENV USE_BAZEL_VERSION 6.5.0
+ENV USE_BAZEL_VERSION=7.4.1
 
 # Install Python and the necessary global dependencies
 RUN apt-get update && apt-get install -y python3 python3-pip && \
@@ -70,6 +74,8 @@ RUN apt-get update && apt-get install -y python3 python3-pip && \
 
 # Setup project files
 
+WORKDIR /xla
+
 ARG XLA_TARGET
 
 ENV XLA_TARGET=${XLA_TARGET}
@@ -80,7 +86,7 @@ COPY mix.exs mix.lock ./
 RUN mix deps.get
 
 COPY lib lib
-COPY Makefile ./
+COPY README.md Makefile ./
 COPY extension extension
 
 CMD [ "mix", "compile" ]

extension/BUILD

@@ -28,8 +28,12 @@ cc_binary(
     "//xla/stream_executor/host:host_platform",
     "//xla:literal",
     "//xla:shape_util",
-    "//xla:status",
-    "//xla:statusor",
+    "//xla/tsl/platform:status",
+    "//xla/tsl/platform:statusor",
+    "//xla/tsl/concurrency:async_value",
+    "@tsl//tsl/platform:platform_port",
+    "//xla/service:custom_call_target_registry",
+    "@llvm-project//llvm:config",
     "//xla:types",
     "//xla:util",
     "//xla/mlir/utils:error_util",
@@ -44,7 +48,6 @@ cc_binary(
     "//xla/client/lib:svd",
     "//xla/client/lib:self_adjoint_eig",
     "//xla/client/lib:sorting",
-    "//xla/pjrt:interpreter_device",
     "//xla/pjrt:pjrt_client",
     "//xla/pjrt:pjrt_compiler",
     "//xla/pjrt:tfrt_cpu_pjrt_client",
@@ -53,6 +56,8 @@ cc_binary(
     "//xla/pjrt/gpu:se_gpu_pjrt_client",
     "//xla/pjrt/distributed:client",
     "//xla/pjrt/distributed:service",
+    "//xla/service:metrics_proto_cc_impl",
+    "//xla/stream_executor/cuda:cuda_compute_capability_proto_cc_impl",
     "@com_google_absl//absl/types:span",
     "@com_google_absl//absl/types:optional",
     "@com_google_absl//absl/base:log_severity",
@@ -66,17 +71,19 @@ cc_binary(
     "@llvm-project//mlir:SparseTensorDialect",
     "@tsl//tsl/platform:errors",
     "@tsl//tsl/platform:fingerprint",
-    "@tsl//tsl/platform:ml_dtypes",
+    "@ml_dtypes_py//ml_dtypes:float8",
+    "@ml_dtypes_py//ml_dtypes:intn",
+    "@ml_dtypes_py//ml_dtypes:mxfloat",
     "@tsl//tsl/platform:statusor",
     "@tsl//tsl/platform:env_impl",
     "@tsl//tsl/platform:tensor_float_32_utils",
-    "@tsl//tsl/profiler/utils:time_utils_impl",
-    "@tsl//tsl/profiler/backends/cpu:annotation_stack_impl",
-    "@tsl//tsl/profiler/backends/cpu:traceme_recorder_impl",
-    "@tsl//tsl/protobuf:protos_all_cc_impl",
-    "@tsl//tsl/protobuf:dnn_proto_cc_impl",
-    "@tsl//tsl/framework:allocator",
-    "@tsl//tsl/framework:allocator_registry_impl",
+    "//xla/tsl/profiler/utils:time_utils_impl",
+    "//xla/tsl/profiler/backends/cpu:annotation_stack_impl",
+    "//xla/tsl/profiler/backends/cpu:traceme_recorder_impl",
+    "//xla/tsl/protobuf:protos_all_cc_impl",
+    "//xla/tsl/protobuf:dnn_proto_cc_impl",
+    "//xla/tsl/framework:allocator",
+    "//xla/tsl/framework:allocator_registry_impl",
     "//xla/tsl/util:determinism",
   ]
   # GRPC Dependencies (needed for PjRt distributed)
@@ -127,9 +134,9 @@ genrule(
   srcs = [
     ":xla_extension_dep_headers",
   ],
-  outs = ["include"],
+  outs = ["include.tar.gz"],
   cmd = """
-    mkdir $@
+    HEADERS_DIR=$$(mktemp -d)
     for f in $(SRCS); do
       d="$${f%/*}"
       d="$${d#bazel-out/*/genfiles/}"
@@ -162,12 +169,15 @@ genrule(
       # Remap ml_dtypes paths
       d="$${d/_virtual_includes\\/intn\\/ml_dtypes/ml_dtypes}"
       d="$${d/_virtual_includes\\/float8\\/ml_dtypes/ml_dtypes}"
-      mkdir -p "$@/$${d}"
-      cp "$${f}" "$@/$${d}/"
+      mkdir -p "$${HEADERS_DIR}/$${d}"
+      cp "$${f}" "$${HEADERS_DIR}/$${d}/"
     done
     # Files in xla/mlir_hlo include sibling headers from mhlo, so we
     # need to mirror them in includes
-    cp -r $@/xla/mlir_hlo/mhlo $@
+    cp -r $${HEADERS_DIR}/xla/mlir_hlo/mhlo $${HEADERS_DIR}
+    # Create the final tarball
+    tar czf "$@" -C $${HEADERS_DIR} .
+    rm -rf $${HEADERS_DIR}
     """,
   )
 
@@ -201,10 +211,12 @@ genrule(
   + if_with_tpu_support([
     ":libtpu_so"
   ]),
-  outs = ["lib"],
+  outs = ["lib.tar.gz"],
   cmd = """
-    mkdir $@
-    mv $(SRCS) $@
+    LIB_DIR=$$(mktemp -d)
+    mv $(SRCS) $${LIB_DIR}/
+    tar czf "$@" -C $${LIB_DIR} .
+    rm -rf $${LIB_DIR}
   """
 )
 
@@ -218,7 +230,13 @@ genrule(
   ],
   cmd = """
     mkdir xla_extension
-    cp -r $(SRCS) xla_extension
+    # Extract the lib tarball
+    mkdir xla_extension/lib
+    tar xzf $(location :xla_extension_lib) -C xla_extension/lib
+    # Extract the headers tarball
+    mkdir xla_extension/include
+    tar xzf $(location :xla_extension_headers) -C xla_extension/include
+    # Create the final package tarball
     tar czf "$@" xla_extension
   """
 )