[NVIDIA GPU SPMD] Add runtime support to run windowed einsum in multiple streams

.github/workflows/buildifier.yml

@@ -0,0 +1,39 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+name: Buildifier
+on:
+  pull_request_target:
+
+env:
+  # Have `go install` place binaries in $PATH
+  GOBIN: "/usr/local/bin"
+
+jobs:
+  buildifier-lint:
+    runs-on: ubuntu-22.04
+    defaults:
+      run:
+        shell: bash
+    timeout-minutes: 1
+    if: |
+      github.event.sender.type == 'User' ||
+      contains(github.event.pull_request.body, 'FORCE_TEST_ACTIONS')
+    steps:
+      - name: "Checking out repository"
+        uses: actions/checkout@e2f20e631ae6d7dd3b768f56a5d2af784dd54791 # v2.5.0
+      - name: "Install buildifier"
+        run: go install github.com/bazelbuild/buildtools/buildifier@433ea85 # 6.4.0
+      - name: "Run buildifier"
+        run: buildifier --lint=warn --warnings=-out-of-order-load -r xla/

.github/workflows/trusted_partners.js

@@ -53,7 +53,7 @@ const get_email_domain = async ({github, username}) => {
 const filter_action = async ({github, context, domain}) => {
   const labels = ['kokoro:force-run'];
 
-  let assignees = ['radhakrishnaba', 'xla-rotation'];
+  let assignees = ['kamaljeeti', 'xla-rotation'];
   const title =
       context.payload.pull_request && context.payload.pull_request.title;
   const lowercased_title = (title || '').toLowerCase();

.gitignore

@@ -8,6 +8,9 @@ bazel-bin
 bazel-out
 bazel-testlogs
 
+# Ignore files produced by `configure`
+.tf_configure.bazelrc
+tools/python_bin_path.sh
 
 # Emacs autosaves
 *~

.kokoro/generate_index_html.sh

@@ -36,9 +36,7 @@ tee "$1" <<EOF
 </ul>
 <h2>Googlers-Only Links</h2>
 <ul>
-<li><a href="http://sponge/$KOKORO_BUILD_ID">Sponge</a></li>
 <li><a href="http://sponge2/$KOKORO_BUILD_ID">Sponge2</a></li>
-<li><a href="http://fusion/$KOKORO_BUILD_ID">Test Fusion</a></li>
 <li><a href="http://sponge/target:$KOKORO_JOB_NAME">Sponge - recent jobs</a></li>
 </ul>
 <h2>Non-Googler Links</h2>

.kokoro/jax/build.sh

@@ -0,0 +1,107 @@
+#!/bin/bash
+# Copyright 2022 Google LLC All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# -e: abort script if one command fails
+# -u: error if undefined variable used
+# -o pipefail: entire command fails if pipe fails. watch out for yes | ...
+# -o history: record shell history
+set -euox pipefail -o history
+
+# Builds + tests jaxlib against CL/PR version of XLA + JAX main.
+
+source "${KOKORO_GFILE_DIR}/utils.sh"
+
+function is_linux_gpu_job() {
+  [[ "$KOKORO_JOB_NAME" =~ tensorflow/xla/jax/.*gpu.* ]]
+}
+
+clone_main_jax() {
+  git clone https://github.com/google/jax.git
+}
+
+prelude() {
+  export JAX_ENABLE_X64=0
+
+  if is_linux_gpu_job ; then
+    export JAX_CUDA_VERSION=12
+    export JAX_CUDNN_VERSION=8.9
+    nvidia-smi
+    setup_env_vars_py39
+  else
+    setup_env_vars_py312
+  fi
+
+  cd "${KOKORO_ARTIFACTS_DIR}"
+
+  use_local_or_install_python
+  install_packages "$NUMPY_VERSION" "$SCIPY_VERSION"
+  clone_main_jax
+  # Install bazel
+  update_bazel_linux
+
+  cd jax
+
+}
+
+build_and_test_on_rbe_cpu() {
+  # Run the tests.
+  bazel \
+      test \
+      --verbose_failures=true \
+      --override_repository=xla="${KOKORO_ARTIFACTS_DIR}"/github/xla \
+      --config=avx_posix \
+      --config=mkl_open_source_only \
+      --config="rbe_cpu_linux_py3.12" \
+      --config=tensorflow_testing_rbe_linux \
+      --test_env=JAX_NUM_GENERATED_CASES=25 \
+      --test_output=errors \
+      -- //tests:cpu_tests //tests:backend_independent_tests
+}
+
+build_and_test_on_rbe_gpu() {
+  # Runs non-multiaccelerator tests with one GPU apiece.
+  # It appears --run_under needs an absolute path.
+
+  bazel \
+    test \
+    --verbose_failures=true \
+    --override_repository=xla="${KOKORO_ARTIFACTS_DIR}"/github/xla \
+    --config=avx_posix \
+    --config=mkl_open_source_only \
+    --config="rbe_linux_cuda12.3_nvcc_py3.9" \
+    --config=tensorflow_testing_rbe_linux \
+    --test_env=XLA_PYTHON_CLIENT_ALLOCATOR=platform \
+    --test_output=errors \
+    --test_env=JAX_SKIP_SLOW_TESTS=1 \
+    --test_env=TF_CPP_MIN_LOG_LEVEL=0 \
+    --test_env=JAX_EXCLUDE_TEST_TARGETS="PmapTest.testSizeOverflow" \
+    --test_tag_filters=-multiaccelerator \
+    -- //tests:gpu_tests //tests:backend_independent_tests
+}
+
+# Generate a templated results file to make output accessible to everyone
+"$KOKORO_ARTIFACTS_DIR"/github/xla/.kokoro/generate_index_html.sh "$KOKORO_ARTIFACTS_DIR"/index.html
+
+prelude
+
+if is_linux_gpu_job ; then
+  build_and_test_on_rbe_gpu
+else
+  build_and_test_on_rbe_cpu
+fi
+
+echo "bazel-testlogs (test results) location:"
+find "$KOKORO_ARTIFACTS_DIR" \
+  -type l,d -name bazel-testlogs || echo "bazel-testlogs not found"

.kokoro/linux/build.sh

@@ -44,16 +44,19 @@ RC_FILE="/usertools/cpu.bazelrc"
 TARGET_FILTER=""
 TAGS_FILTER="-no_oss,-oss_excluded,-oss_serial"
 ADDITIONAL_FLAGS=""
+RBE_CONFIG=""
 
 if is_linux_gpu_job ; then
     TAGS_FILTER="$TAGS_FILTER,gpu,requires-gpu-nvidia,-no_gpu"
     ADDITIONAL_FLAGS="$ADDITIONAL_FLAGS --run_under=//tools/ci_build/gpu_build:parallel_gpu_execute"
     RC_FILE="/usertools/gpu.bazelrc"
+    RBE_CONFIG="rbe_linux_cuda_nvcc"
     echo "***NOTE: nvidia-smi lists the highest CUDA version the driver supports, which may be different than the version of CUDA actually used!!***"
     nvidia-smi
 else
     TAGS_FILTER="$TAGS_FILTER,-gpu,-requires-gpu-nvidia"
     ADDITIONAL_FLAGS="$ADDITIONAL_FLAGS --config=nonccl"
+    RBE_CONFIG="rbe_linux_cpu"
 fi
 
 # Build & test XLA
@@ -65,7 +68,7 @@ docker exec xla bazel --bazelrc=$RC_FILE \
         --features=layering_check \
         --profile=/tf/pkg/profile.json.gz \
         --flaky_test_attempts=3 \
-        --config=rbe \
+        --config=$RBE_CONFIG \
         --jobs=150 \
         --nobuild_tests_only \
         $ADDITIONAL_FLAGS \
@@ -74,9 +77,7 @@ docker exec xla bazel --bazelrc=$RC_FILE \
 
 # Print build time statistics, including critical path.
 docker exec xla bazel analyze-profile "/tf/pkg/profile.json.gz"
-# TODO(ddunleavy): enable once container has clang-tidy
-# docker exec xla git config --global --add safe.directory /tf/xla
-# docker exec xla bash -c "bazel aquery --output=jsonproto \"mnemonic(CppCompile, //xla/...)\" | PYTHONPATH=.. python3 build_tools/lint/clang_tidy.py --changed_lines_only"
+
 # Stop container
 docker stop xla
 

.kokoro/macos/build.sh

@@ -81,7 +81,15 @@ bazel test \
     --output_filter="" \
     --macos_minimum_os=10.15 \
     --keep_going \
+    --test_output=errors \
     --config=nonccl \
     --build_tag_filters=$TAGS_FILTER  --test_tag_filters=$TAGS_FILTER \
     --test_size_filters=small,medium \
     -- //xla/... $TARGET_FILTER
+
+
+# We want to store the individual test logs in GCS after the build finishes.
+# Since the Bazel test output folder is in a different partition to the Kokoro
+# artifacts folder, we need to symlink the test output folder to be in the
+# artifacts folder.
+ln -s /Volumes/BuildData/bazel_output "$KOKORO_ARTIFACTS_DIR"

build_tools/docker/context/install_bazel.sh

@@ -1,4 +1,4 @@
-# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2023 The OpenXLA Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

build_tools/docker/context/install_python_deps.sh

@@ -1,4 +1,4 @@
-# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2023 The OpenXLA Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

build_tools/docker/dockerfiles/benchmarking.Dockerfile

@@ -1,4 +1,4 @@
-# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2023 The OpenXLA Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

build_tools/github_actions/build_xla.sh

@@ -1,6 +1,6 @@
 #!/bin/bash
 
-# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2023 The OpenXLA Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

build_tools/lint/BUILD

@@ -1,4 +1,4 @@
-# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2023 The OpenXLA Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -13,24 +13,14 @@
 # limitations under the License.
 # ============================================================================
 
-load(
-    "//xla:pytype.default.bzl",
-    "pytype_strict_binary",
-    "pytype_strict_library",
-)
+load("//xla:pytype.default.bzl", "pytype_strict_library")
 # Placeholder: load py_test
 
 package(
     # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
     licenses = ["notice"],
 )
 
-pytype_strict_binary(
-    name = "clang_tidy",
-    srcs = ["clang_tidy.py"],
-    deps = [":diff_parser"],
-)
-
 pytype_strict_library(
     name = "check_contents",
     srcs = ["check_contents.py"],
@@ -43,6 +33,11 @@ pytype_strict_library(
     visibility = ["//visibility:public"],
 )
 
+pytype_strict_library(
+    name = "generate_compile_commands",
+    srcs = ["generate_compile_commands.py"],
+)
+
 py_test(
     name = "check_contents_test",
     srcs = ["check_contents_test.py"],
@@ -71,3 +66,13 @@ py_test(
         "@absl_py//absl/testing:absltest",
     ],
 )
+
+py_test(
+    name = "generate_compile_commands_test",
+    srcs = ["generate_compile_commands_test.py"],
+    tags = ["no_oss"],
+    deps = [
+        ":generate_compile_commands",
+        "@absl_py//absl/testing:absltest",
+    ],
+)