Add Docker build script for Linux ARM

.github/workflows/release.yml

@@ -180,36 +180,3 @@ jobs:
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
           XLA_TARGET: ${{ matrix.xla_target }}
-
-  linux_arm:
-    name: "aarch64-linux-gnu-cpu (cross-compiled)"
-    needs: [create_draft_release]
-    # We intentionally build on ubuntu 20 to compile against
-    # an older version of glibc
-    runs-on: ubuntu-20.04
-    steps:
-      - uses: actions/checkout@v3
-      - uses: erlef/setup-beam@v1
-        with:
-          otp-version: "24"
-          elixir-version: "1.12.3"
-      # Setup the compilation environment
-      - uses: abhinavsingh/setup-bazel@v3
-        with:
-          version: "6.1.2"
-      - uses: actions/setup-python@v2
-        with:
-          python-version: "3.9"
-      - run: python -m pip install --upgrade pip numpy
-      # Build and upload the archives
-      - run: mix deps.get
-      # Hide system OpenSSL as suggested in https://github.com/tensorflow/tensorflow/issues/48401#issuecomment-818377995
-      - run: sudo mv /usr/include/openssl /usr/include/openssl.original
-      - run: .github/scripts/compile_and_upload.sh ${{ github.ref_name }}
-        env:
-          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-          XLA_TARGET: cpu
-          XLA_TARGET_PLATFORM: "aarch64-linux-gnu"
-          # Explicitly cross-compile for arm64
-          BUILD_FLAGS: "--config=elinux_aarch64"
-          CC: gcc-9

builds/build.sh

@@ -7,7 +7,7 @@ cd "$(dirname "$0")/.."
 print_usage_and_exit() {
   echo "Usage: $0 <variant>"
   echo ""
-  echo "Compiles the project inside docker. Available variants: cuda118, cuda120."
+  echo "Compiles the project inside docker. Available variants: cpu, cuda118, cuda120."
   exit 1
 }
 
@@ -21,6 +21,14 @@ fi
 # [1]: https://docs.nvidia.com/deeplearning/cudnn/archives/index.html
 
 case "$1" in
+  "cpu")
+    docker build -t xla-cpu -f builds/cpu.Dockerfile \
+      --build-arg XLA_TARGET=cpu \
+      .
+
+    docker run --rm -v $(pwd)/builds/output/cpu/build:/build -v $(pwd)/builds/output/cpu/.cache:/root/.cache xla-cpu
+  ;;
+
   "cuda118")
     docker build -t xla-cuda118 -f builds/cuda.Dockerfile \
       --build-arg CUDA_VERSION=11.8.0 \

builds/cpu.Dockerfile

@@ -0,0 +1,43 @@
+FROM hexpm/elixir:1.15.4-erlang-26.0.2-ubuntu-focal-20230126 AS elixir
+
+# Set the missing UTF-8 locale, otherwise Elixir warns
+ENV LC_ALL C.UTF-8
+
+# Make sure installing packages (like tzdata) doesn't prompt for configuration
+ENV 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 && \
+  # Install basic system dependencies
+  apt-get update && apt-get install -y ca-certificates curl git unzip wget
+
+# 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)" && \
+  chmod +x bazel && \
+  mv bazel /usr/local/bin/bazel
+
+ENV USE_BAZEL_VERSION 6.1.0
+
+# Install Python and the necessary global dependencies
+RUN apt-get install -y python3 python3-pip && \
+  ln -s /usr/bin/python3 /usr/bin/python && \
+  python -m pip install --upgrade pip numpy
+
+# ---
+
+ARG XLA_TARGET
+
+ENV XLA_TARGET=${XLA_TARGET}
+ENV XLA_CACHE_DIR=/build
+ENV XLA_BUILD=true
+
+COPY mix.exs mix.lock ./
+RUN mix deps.get
+
+COPY lib lib
+COPY Makefile Makefile.win ./
+COPY extension extension
+
+CMD [ "mix", "compile" ]

builds/cuda.Dockerfile

@@ -25,13 +25,12 @@ RUN cuda_version="${CUDA_VERSION}" && \
   cudnn_package_version="$(apt-cache madison libcudnn8 | grep -o "${cudnn_version}.*-1+cuda${cuda_version%.*}")" && \
   apt-get install -y --allow-downgrades --allow-change-held-packages libcudnn8=$cudnn_package_version libcudnn8-dev=$cudnn_package_version
 
-# Install Bazel
-RUN apt-get install -y apt-transport-https curl gnupg && \
-  curl -fsSL https://bazel.build/bazel-release.pub.gpg | gpg --dearmor > bazel.gpg && \
-  mv bazel.gpg /etc/apt/trusted.gpg.d/ && \
-  echo "deb [arch=amd64] https://storage.googleapis.com/bazel-apt stable jdk1.8" | tee /etc/apt/sources.list.d/bazel.list && \
-  apt-get update && apt-get install -y bazel-6.1.2 && \
-  ln -s /usr/bin/bazel-6.1.2 /usr/bin/bazel
+# 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)" && \
+  chmod +x bazel && \
+  mv bazel /usr/local/bin/bazel
+
+ENV USE_BAZEL_VERSION 6.1.2
 
 # Install Python and the necessary global dependencies
 RUN apt-get install -y python3 python3-pip && \