Build updates (#48)

elixir-nx · Aug 10, 2023 · 2f2b4e6 · 2f2b4e6
1 parent c2b4a23
commit 2f2b4e6
Show file tree

Hide file tree

Showing 5 changed files with 72 additions and 69 deletions.
diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
@@ -39,54 +39,54 @@ jobs:
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
 
-  # linux x86_64 cpu musl
-  alpine:
-    runs-on: ubuntu-20.04
-    strategy:
-      fail-fast: false
-    container: hexpm/elixir:1.13.4-erlang-25.0.2-alpine-3.16.0
-    env:
-      GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-      XLA_TARGET: cpu
-      BAZEL_VERSION: "5.3.0"
-      BAZEL_SHA256SUM: "ee801491ff0ec3a562422322a033c9afe8809b64199e4a94c7433d4e14e6b921  bazel-5.3.0-dist.zip"
-      JAVA_HOME: "/usr/lib/jvm/default-jvm"
-    steps:
-      - name: Install system packages
-        run: |
-          apk update && apk upgrade && \
-          apk add --no-cache python3 py3-pip python3-dev py3-numpy && \
-          apk add --no-cache libstdc++ openjdk11 && \
-          apk add --no-cache bash curl git wget && \
-          apk add --no-cache musl-dev make libexecinfo libexecinfo-dev && \
-          apk add --no-cache coreutils gcc g++ linux-headers unzip zip && \
-          apk add --no-cache automake gcc subversion && \
-          apk add --no-cache github-cli && \
-          DIR=$(mktemp -d) && cd ${DIR} && \
-          curl -sLO https://github.com/bazelbuild/bazel/releases/download/${BAZEL_VERSION}/bazel-${BAZEL_VERSION}-dist.zip && \
-          echo ${BAZEL_SHA256SUM} | sha256sum --check && \
-          unzip bazel-${BAZEL_VERSION}-dist.zip && \
-          EXTRA_BAZEL_ARGS="--tool_java_runtime_version=local_jdk" bash ./compile.sh && \
-          cp ${DIR}/output/bazel /usr/local/bin/ && \
-          rm -rf ${DIR}
-      - name: Install hex
-        run: mix local.hex --force && mix local.rebar --force
-      # Prevent git from checking the repository owner and erroring with "dubious ownership"
-      - run: git config --global --add safe.directory '*'
-      - uses: actions/checkout@v3
-      # Build and upload the archive
-      - run: mix deps.get
-      - run: .github/scripts/compile_unless_exists.sh
-        env:
-          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-          XLA_TARGET: cpu
-      - run: .github/scripts/upload_archives.sh
-        env:
-          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+  # # linux x86_64 cpu musl
+  # alpine:
+  #   runs-on: ubuntu-20.04
+  #   strategy:
+  #     fail-fast: false
+  #   container: hexpm/elixir:1.13.4-erlang-25.0.2-alpine-3.16.0
+  #   env:
+  #     GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+  #     XLA_TARGET: cpu
+  #     BAZEL_VERSION: "5.3.0"
+  #     BAZEL_SHA256SUM: "ee801491ff0ec3a562422322a033c9afe8809b64199e4a94c7433d4e14e6b921  bazel-5.3.0-dist.zip"
+  #     JAVA_HOME: "/usr/lib/jvm/default-jvm"
+  #   steps:
+  #     - name: Install system packages
+  #       run: |
+  #         apk update && apk upgrade && \
+  #         apk add --no-cache python3 py3-pip python3-dev py3-numpy && \
+  #         apk add --no-cache libstdc++ openjdk11 && \
+  #         apk add --no-cache bash curl git wget && \
+  #         apk add --no-cache musl-dev make libexecinfo libexecinfo-dev && \
+  #         apk add --no-cache coreutils gcc g++ linux-headers unzip zip && \
+  #         apk add --no-cache automake gcc subversion && \
+  #         apk add --no-cache github-cli && \
+  #         DIR=$(mktemp -d) && cd ${DIR} && \
+  #         curl -sLO https://github.com/bazelbuild/bazel/releases/download/${BAZEL_VERSION}/bazel-${BAZEL_VERSION}-dist.zip && \
+  #         echo ${BAZEL_SHA256SUM} | sha256sum --check && \
+  #         unzip bazel-${BAZEL_VERSION}-dist.zip && \
+  #         EXTRA_BAZEL_ARGS="--tool_java_runtime_version=local_jdk" bash ./compile.sh && \
+  #         cp ${DIR}/output/bazel /usr/local/bin/ && \
+  #         rm -rf ${DIR}
+  #     - name: Install hex
+  #       run: mix local.hex --force && mix local.rebar --force
+  #     # Prevent git from checking the repository owner and erroring with "dubious ownership"
+  #     - run: git config --global --add safe.directory '*'
+  #     - uses: actions/checkout@v3
+  #     # Build and upload the archive
+  #     - run: mix deps.get
+  #     - run: .github/scripts/compile_unless_exists.sh
+  #       env:
+  #         GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+  #         XLA_TARGET: cpu
+  #     - run: .github/scripts/upload_archives.sh
+  #       env:
+  #         GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
 
   # darwin x86_64 cpu
   macos:
-    runs-on: macos-10.15
+    runs-on: macos-12
     steps:
       - uses: actions/checkout@v3
       - run: brew install elixir
@@ -112,7 +112,7 @@ jobs:
 
   # darwin aarch64 cpu (cross-compiled)
   macos_arm:
-    runs-on: macos-10.15
+    runs-on: macos-12
     steps:
       - uses: actions/checkout@v3
       - run: brew install elixir
@@ -158,10 +158,6 @@ jobs:
         include:
           # Note: we use exact references, because the images may be updated with cuDNN
           # version bumps unexpectedly
-          # - container: nvidia/cuda:11.1.1-cudnn8-devel-ubuntu20.04
-          - container: nvidia/cuda@sha256:bb8379ce06f114591a64031c31d4f0a42a711b7fe73a805711a1821835b36c41
-            xla_target: cuda111
-            python: "3.9"
           # - container: nvidia/cuda:11.4.3-cudnn8-devel-ubuntu20.04
           - container: nvidia/cuda@sha256:275d35f7985cfe12dc3e6e40a3df5fff2bb217b1c248f0f70dadaa298d1b938b
             xla_target: cuda114

diff --git a/Makefile b/Makefile
@@ -11,7 +11,7 @@ BUILD_MODE ?= opt # can also be dbg
 BUILD_CACHE ?= $(TEMP)/xla_extension
 OPENXLA_GIT_REPO ?= https://github.com/openxla/xla.git
 
-OPENXLA_GIT_REV ?= c1e4a16e77a7ba2000003ccade3ffba3749ada35
+OPENXLA_GIT_REV ?= b938cfdf2d4e9a5f69c494a316e92638c1a119ef
 
 # Private configuration
 BAZEL_FLAGS = --define "framework_shared_object=false" -c $(BUILD_MODE)

diff --git a/README.md b/README.md
@@ -24,7 +24,6 @@ only the host CPU.
 | cuda120 | CUDA 12.0+, cuDNN 8.8+ (recommended) |
 | cuda118 | CUDA 11.8+, cuDNN 8.7+ (recommended) |
 | cuda114 | CUDA 11.4+, cuDNN 8.2+ |
-| cuda111 | CUDA 11.1+, cuDNN 8.0.5+ |
 | cuda | CUDA x.y, cuDNN (building from source only) |
 | rocm | ROCm (building from source only) |
 

diff --git a/extension/BUILD b/extension/BUILD
@@ -1,10 +1,9 @@
 load("//xla/stream_executor:build_defs.bzl", "if_cuda_or_rocm",)
-load("@tsl//tsl:tsl.bzl", "if_with_tpu_support",)
 load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda",)
 load("@local_config_rocm//rocm:build_defs.bzl", "if_rocm",)
 load("@tsl//tsl:tsl.bzl", "tsl_grpc_cc_dependencies",)
 load("@tsl//tsl:tsl.bzl", "transitive_hdrs",)
-load("@bazel_tools//tools/build_defs/pkg:pkg.bzl", "pkg_tar",)
+load("@rules_pkg//pkg:tar.bzl", "pkg_tar")
 
 package(default_visibility=["//visibility:private"])
 
@@ -19,6 +18,7 @@ cc_binary(
     "//xla/service:memory_space_assignment_proto_cc_impl",
     "//xla/service/gpu:backend_configs_cc_impl",
     "//xla/stream_executor:dnn_proto_cc_impl",
+    "//xla/stream_executor/host:host_platform",
     "//xla:literal",
     "//xla:shape_util",
     "//xla:status",
@@ -44,12 +44,11 @@ cc_binary(
     "//xla/pjrt:tfrt_cpu_pjrt_client",
     "//xla/pjrt:pjrt_c_api_client",
     "//xla/pjrt:tpu_client",
-    "//xla/pjrt:pjrt_plugin_device_client",
-    "//xla/pjrt:pjrt_plugin_device_client_headers",
     "//xla/pjrt/distributed",
     "//xla/pjrt/gpu:se_gpu_pjrt_client",
     "//xla/pjrt/distributed:client",
     "//xla/pjrt/distributed:service",
+    "//xla:autotuning_proto_cc_impl",
     "@com_google_absl//absl/types:span",
     "@com_google_absl//absl/types:optional",
     "@com_google_absl//absl/base:log_severity",
@@ -63,7 +62,6 @@ cc_binary(
     "@llvm-project//mlir:SparseTensorDialect",
     "@tf_runtime//:core_runtime",
     "@tf_runtime//:hostcontext",
-    "@tf_runtime//third_party/llvm_derived:unique_any",
     "@tsl//tsl/platform:errors",
     "@tsl//tsl/platform:fingerprint",
     "@tsl//tsl/platform:float8",
@@ -73,7 +71,6 @@ cc_binary(
     "@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:autotuning_proto_cc_impl",
     "@tsl//tsl/protobuf:protos_all_cc_impl",
     "@tsl//tsl/protobuf:dnn_proto_cc_impl",
     "@tsl//tsl/framework:allocator",
@@ -82,9 +79,6 @@ cc_binary(
   ]
   # GRPC Dependencies (needed for PjRt distributed)
   + tsl_grpc_cc_dependencies()
-  + if_with_tpu_support([
-    "@tsl//tsl/core/tpu:tpu_api_dlsym_initializer",
-  ])
   + if_cuda_or_rocm([
     "//xla/service:gpu_plugin",
   ])
@@ -169,13 +163,17 @@ genrule(
   """
 )
 
-pkg_tar(
+# See https://github.com/bazelbuild/rules_pkg/issues/517#issuecomment-1492917994
+genrule(
   name = "xla_extension",
-  extension = "tar.gz",
-  mode = "0644",
-  package_dir = "xla_extension",
+  outs = ["xla_extension.tar.gz"],
   srcs = [
     ":xla_extension_lib",
     ":xla_extension_headers",
   ],
+  cmd = """
+    mkdir xla_extension
+    cp -r $(SRCS) xla_extension
+    tar czf "$@" xla_extension
+  """
 )
diff --git a/lib/xla.ex b/lib/xla.ex
@@ -78,10 +78,20 @@ defmodule XLA do
   def make_env() do
     bazel_build_flags_accelerator =
       case xla_target() do
-        "cuda" <> _ -> ["--config=cuda"]
-        "rocm" <> _ -> ["--config=rocm", "--action_env=HIP_PLATFORM=hcc"]
-        "tpu" <> _ -> ["--config=tpu"]
-        _ -> []
+        "cuda" <> _ ->
+          [
+            # See https://github.com/google/jax/blob/c9cf6b44239e373cba384936dcfeff60e39ad560/.bazelrc#L68
+            ~s/--config=cuda --action_env=TF_CUDA_COMPUTE_CAPABILITIES="sm_52,sm_60,sm_70,compute_80"/
+          ]
+
+        "rocm" <> _ ->
+          ["--config=rocm", "--action_env=HIP_PLATFORM=hcc"]
+
+        "tpu" <> _ ->
+          ["--config=tpu"]
+
+        _ ->
+          []
       end
 
     bazel_build_flags_cpu =