Fix XLA tools build

iree-org · Jan 5, 2024 · d8dfd14 · d8dfd14
1 parent ca9244f
commit d8dfd14
Show file tree

Hide file tree

Showing 16 changed files with 132 additions and 77 deletions.
diff --git a/.github/workflows/run_comparative_benchmark.yml b/.github/workflows/run_comparative_benchmark.yml
@@ -11,6 +11,7 @@ name: Comparative Benchmarks
 on:
   # Will only run when manually triggered.
   workflow_dispatch:
+  pull_request:
 
 concurrency:
   # A PR number if a pull request and otherwise the commit hash. This cancels
@@ -78,7 +79,7 @@ jobs:
         run: |
           mkdir -p "${XLA_TOOLS_DIR}"
           docker run --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
-            "gcr.io/iree-oss/openxla-benchmark/cuda11.8-cudnn8.9@sha256:02a54fbe596d6770a7e3bf7eac243c5bd71a3008cdc4da8cf8b8b79f95417710" \
+            "gcr.io/iree-oss/openxla-benchmark/cuda11.8-cudnn8.9@sha256:8ab3e5f30388c5e26a5cd88872e5bbe890a8ad49c4fe7aa3a471f122e9718e61" \
             ./comparative_benchmark/xla_hlo/build_xla_tools.sh \
               "${XLA_TOOLS_DIR}" \
               "${CUDA_VERSION}"
@@ -133,7 +134,7 @@ jobs:
           RESULTS_PATH="${RESULTS_DIR}/${XLA_HLO_RESULTS_JSON}"
           docker run --gpus all --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
             --env "OOBI_XLA_TOOLS_DIR=${XLA_TOOLS_DIR}" \
-            "gcr.io/iree-oss/openxla-benchmark/cuda11.8-cudnn8.9@sha256:02a54fbe596d6770a7e3bf7eac243c5bd71a3008cdc4da8cf8b8b79f95417710" \
+            "gcr.io/iree-oss/openxla-benchmark/cuda11.8-cudnn8.9@sha256:8ab3e5f30388c5e26a5cd88872e5bbe890a8ad49c4fe7aa3a471f122e9718e61" \
             ./comparative_benchmark/xla_hlo/benchmark_all.sh \
               "${TARGET_DEVICE}"\
               "${RESULTS_PATH}"
@@ -145,7 +146,7 @@ jobs:
         run: |
           RESULTS_PATH="${RESULTS_DIR}/${JAX_XLA_RESULTS_JSON}"
           docker run --gpus all --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
-            "gcr.io/iree-oss/openxla-benchmark/cuda11.8-cudnn8.9@sha256:02a54fbe596d6770a7e3bf7eac243c5bd71a3008cdc4da8cf8b8b79f95417710" \
+            "gcr.io/iree-oss/openxla-benchmark/cuda11.8-cudnn8.9@sha256:8ab3e5f30388c5e26a5cd88872e5bbe890a8ad49c4fe7aa3a471f122e9718e61" \
             ./comparative_benchmark/jax/benchmark_xla.sh \
               "${TARGET_DEVICE}"\
               "${RESULTS_PATH}"
@@ -158,7 +159,7 @@ jobs:
 #        run: |
 #          RESULTS_PATH="${RESULTS_DIR}/${JAX_IREE_RESULTS_JSON}"
 #          docker run --gpus all --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
-#            "gcr.io/iree-oss/openxla-benchmark/cuda11.8-cudnn8.9@sha256:02a54fbe596d6770a7e3bf7eac243c5bd71a3008cdc4da8cf8b8b79f95417710" \
+#            "gcr.io/iree-oss/openxla-benchmark/cuda11.8-cudnn8.9@sha256:8ab3e5f30388c5e26a5cd88872e5bbe890a8ad49c4fe7aa3a471f122e9718e61" \
 #            ./comparative_benchmark/jax/benchmark_iree.sh \
 #              "${TARGET_DEVICE}"\
 #              "${RESULTS_PATH}"
@@ -170,7 +171,7 @@ jobs:
         run: |
           RESULTS_PATH="${RESULTS_DIR}/${TF_XLA_RESULTS_JSON}"
           docker run --gpus all --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
-            "gcr.io/iree-oss/openxla-benchmark/cuda11.8-cudnn8.9@sha256:02a54fbe596d6770a7e3bf7eac243c5bd71a3008cdc4da8cf8b8b79f95417710" \
+            "gcr.io/iree-oss/openxla-benchmark/cuda11.8-cudnn8.9@sha256:8ab3e5f30388c5e26a5cd88872e5bbe890a8ad49c4fe7aa3a471f122e9718e61" \
             ./comparative_benchmark/tf_xla/benchmark_all.sh \
               "${TARGET_DEVICE}"\
               "${RESULTS_PATH}"
@@ -182,7 +183,7 @@ jobs:
         run: |
           RESULTS_PATH="${RESULTS_DIR}/${PT_INDUCTOR_RESULTS_JSON}"
           docker run --gpus all --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
-            "gcr.io/iree-oss/openxla-benchmark/cuda11.8-cudnn8.9@sha256:02a54fbe596d6770a7e3bf7eac243c5bd71a3008cdc4da8cf8b8b79f95417710" \
+            "gcr.io/iree-oss/openxla-benchmark/cuda11.8-cudnn8.9@sha256:8ab3e5f30388c5e26a5cd88872e5bbe890a8ad49c4fe7aa3a471f122e9718e61" \
             ./comparative_benchmark/pt_inductor/benchmark_all.sh \
               "${TARGET_DEVICE}"\
               "${RESULTS_PATH}"
@@ -223,7 +224,7 @@ jobs:
           RESULTS_PATH="${RESULTS_DIR}/${XLA_HLO_RESULTS_JSON}"
           docker run --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
             --env "OOBI_XLA_TOOLS_DIR=${XLA_TOOLS_DIR}" \
-            "gcr.io/iree-oss/openxla-benchmark/base@sha256:2dbee52eaa63e62137682f0eda701ac4cf59b8e16395daa757f6e1906b52dd82" \
+            "gcr.io/iree-oss/openxla-benchmark/base@sha256:5e13cda13bdbd6fcaeeae66f2d2ffe5a8c5ef88f763fcd489db7b87e04f40ccd" \
             ./comparative_benchmark/xla_hlo/benchmark_all.sh \
               "${TARGET_DEVICE}"\
               "${RESULTS_PATH}"
@@ -235,7 +236,7 @@ jobs:
         run: |
           RESULTS_PATH="${RESULTS_DIR}/${JAX_XLA_RESULTS_JSON}"
           docker run --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
-            "gcr.io/iree-oss/openxla-benchmark/base@sha256:2dbee52eaa63e62137682f0eda701ac4cf59b8e16395daa757f6e1906b52dd82" \
+            "gcr.io/iree-oss/openxla-benchmark/base@sha256:5e13cda13bdbd6fcaeeae66f2d2ffe5a8c5ef88f763fcd489db7b87e04f40ccd" \
             ./comparative_benchmark/jax/benchmark_xla.sh \
               "${TARGET_DEVICE}"\
               "${RESULTS_PATH}"
@@ -247,7 +248,7 @@ jobs:
         run: |
           RESULTS_PATH="${RESULTS_DIR}/${JAX_IREE_RESULTS_JSON}"
           docker run --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
-            "gcr.io/iree-oss/openxla-benchmark/base@sha256:2dbee52eaa63e62137682f0eda701ac4cf59b8e16395daa757f6e1906b52dd82" \
+            "gcr.io/iree-oss/openxla-benchmark/base@sha256:5e13cda13bdbd6fcaeeae66f2d2ffe5a8c5ef88f763fcd489db7b87e04f40ccd" \
             ./comparative_benchmark/jax/benchmark_iree.sh \
               "${TARGET_DEVICE}"\
               "${RESULTS_PATH}"
@@ -259,7 +260,7 @@ jobs:
         run: |
           RESULTS_PATH="${RESULTS_DIR}/${TF_XLA_RESULTS_JSON}"
           docker run --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
-            "gcr.io/iree-oss/openxla-benchmark/base@sha256:2dbee52eaa63e62137682f0eda701ac4cf59b8e16395daa757f6e1906b52dd82" \
+            "gcr.io/iree-oss/openxla-benchmark/base@sha256:5e13cda13bdbd6fcaeeae66f2d2ffe5a8c5ef88f763fcd489db7b87e04f40ccd" \
             ./comparative_benchmark/tf_xla/benchmark_all.sh \
               "${TARGET_DEVICE}"\
               "${RESULTS_PATH}"
@@ -271,7 +272,7 @@ jobs:
         run: |
           RESULTS_PATH="${RESULTS_DIR}/${PT_INDUCTOR_RESULTS_JSON}"
           docker run --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
-            "gcr.io/iree-oss/openxla-benchmark/base@sha256:2dbee52eaa63e62137682f0eda701ac4cf59b8e16395daa757f6e1906b52dd82" \
+            "gcr.io/iree-oss/openxla-benchmark/base@sha256:5e13cda13bdbd6fcaeeae66f2d2ffe5a8c5ef88f763fcd489db7b87e04f40ccd" \
             ./comparative_benchmark/pt_inductor/benchmark_all.sh \
               "${TARGET_DEVICE}"\
               "${RESULTS_PATH}"

diff --git a/.github/workflows/run_convperf.yml b/.github/workflows/run_convperf.yml
@@ -11,6 +11,7 @@ name: convperf
 on:
   # Will only run when manually triggered.
   workflow_dispatch:
+  pull_request:
 
 concurrency:
   # A PR number if a pull request and otherwise the commit hash. This cancels
@@ -71,7 +72,7 @@ jobs:
           mkdir ${RESULTS_DIR}
           # Run under privileged mode because system-level scheduling is used.
           docker run --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" --privileged \
-            gcr.io/iree-oss/openxla-benchmark/convperf@sha256:fda453f1414085558a5c4a2250db0010841c73a082379997f507a58b19cb8aaf \
+            gcr.io/iree-oss/openxla-benchmark/convperf@sha256:964de007fd4b73467032e698945c03e04c46064a61766ee5c6c610d480b00ef4 \
           ./experimental/microbenchmarks/convperf/build_and_run_convperf.sh "${BUILD_DIR}" "${RESULTS_DIR}"
       - name: "Uploading results"
         run: |

diff --git a/.github/workflows/run_ggml_benchmark.yml b/.github/workflows/run_ggml_benchmark.yml
@@ -10,6 +10,7 @@ name: GGML Benchmarks
 
 on:
   workflow_dispatch:
+  pull_request:
 
 concurrency:
   # A PR number if a pull request and otherwise the commit hash. This cancels
@@ -77,7 +78,7 @@ jobs:
       - name: "Building GGML CPU"
         run: |
           docker run --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
-            "gcr.io/iree-oss/openxla-benchmark/base@sha256:2dbee52eaa63e62137682f0eda701ac4cf59b8e16395daa757f6e1906b52dd82" \
+            "gcr.io/iree-oss/openxla-benchmark/base@sha256:5e13cda13bdbd6fcaeeae66f2d2ffe5a8c5ef88f763fcd489db7b87e04f40ccd" \
             ./experimental/ggml/build_ggml.sh \
               "${TARGET_DEVICE}" \
               "${GGML_BUILD_DIR}"
@@ -88,7 +89,7 @@ jobs:
         run: |
           RESULTS_PATH="${RESULTS_DIR}/${GGML_RESULTS_JSON}"
           docker run --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
-            "gcr.io/iree-oss/openxla-benchmark/base@sha256:2dbee52eaa63e62137682f0eda701ac4cf59b8e16395daa757f6e1906b52dd82" \
+            "gcr.io/iree-oss/openxla-benchmark/base@sha256:5e13cda13bdbd6fcaeeae66f2d2ffe5a8c5ef88f763fcd489db7b87e04f40ccd" \
             ./experimental/ggml/benchmark_ggml.sh \
               "${TARGET_DEVICE}" \
               "${GGML_BUILD_DIR}" \
@@ -118,7 +119,7 @@ jobs:
         run: |
           mkdir -p "${GGML_BUILD_DIR}"
           docker run --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
-            "gcr.io/iree-oss/openxla-benchmark/android@sha256:4c219b20213c7f26edf848b6924a95991c0d359463709d4ee55b499eeaa3349b" \
+            "gcr.io/iree-oss/openxla-benchmark/android@sha256:89921d48b01eade590bc84ba4ee7eda2fb436cbe041961d2b14ff2884fdbab7c" \
             ./experimental/ggml/build_ggml.sh \
               "${TARGET_DEVICE}" \
               "${GGML_BUILD_DIR}"

diff --git a/.github/workflows/run_mmperf.yml b/.github/workflows/run_mmperf.yml
@@ -70,7 +70,7 @@ jobs:
       - name: "Building mmperf for CPU"
         run: |
           docker run --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
-            gcr.io/iree-oss/openxla-benchmark/mmperf@sha256:c972ce5b2144de0786f103611fecbd88d93dd45ecd068f8c97d98c08677cee57 \
+            gcr.io/iree-oss/openxla-benchmark/mmperf@sha256:55720ca87862121c5dbb6027a09586ef7c3da59ed4bca74bfed8ba879c464cf3 \
           ./experimental/microbenchmarks/mmperf/build_mmperf.sh "cpu" "${BUILD_DIR}"
       - name: "Removing unused files"
         run: |
@@ -114,7 +114,7 @@ jobs:
       - name: "Building mmperf for CUDA"
         run: |
           docker run --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
-            gcr.io/iree-oss/openxla-benchmark/mmperf@sha256:c972ce5b2144de0786f103611fecbd88d93dd45ecd068f8c97d98c08677cee57 \
+            gcr.io/iree-oss/openxla-benchmark/mmperf@sha256:55720ca87862121c5dbb6027a09586ef7c3da59ed4bca74bfed8ba879c464cf3 \
           ./experimental/microbenchmarks/mmperf/build_mmperf.sh "cuda" "${BUILD_DIR}"
       - name: "Removing unused files"
         run: |
@@ -163,7 +163,7 @@ jobs:
         run: |
           mkdir ${RESULTS_DIR}
           docker run --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
-            gcr.io/iree-oss/openxla-benchmark/mmperf@sha256:c972ce5b2144de0786f103611fecbd88d93dd45ecd068f8c97d98c08677cee57 \
+            gcr.io/iree-oss/openxla-benchmark/mmperf@sha256:55720ca87862121c5dbb6027a09586ef7c3da59ed4bca74bfed8ba879c464cf3 \
           ./experimental/microbenchmarks/mmperf/run_mmperf.sh "${BUILD_DIR}" "${RESULTS_DIR}"
       - name: "Uploading results"
         run: |
@@ -195,7 +195,7 @@ jobs:
         run: |
           mkdir ${RESULTS_DIR}
           docker run --gpus all --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
-            gcr.io/iree-oss/openxla-benchmark/mmperf@sha256:c972ce5b2144de0786f103611fecbd88d93dd45ecd068f8c97d98c08677cee57 \
+            gcr.io/iree-oss/openxla-benchmark/mmperf@sha256:55720ca87862121c5dbb6027a09586ef7c3da59ed4bca74bfed8ba879c464cf3 \
           ./experimental/microbenchmarks/mmperf/run_mmperf.sh "${BUILD_DIR}" "${RESULTS_DIR}"
       - name: "Uploading results"
         run: |

diff --git a/.github/workflows/run_mobile_comparative_benchmark.yml b/.github/workflows/run_mobile_comparative_benchmark.yml
@@ -13,6 +13,7 @@ on:
   schedule:
     # Scheduled to run at 09:00 UTC and 21:00 UTC.
     - cron: '0 09,21 * * *'
+  pull_request:
 
 concurrency:
   # A PR number if a pull request and otherwise the commit hash. This cancels
@@ -98,7 +99,7 @@ jobs:
               --env IREE_SOURCE_DIR="${IREE_SOURCE_DIR}" \
               --env IREE_INSTALL_DIR="${IREE_INSTALL_DIR}" \
               --env ANDROID_PLATFORM_VERSION="${ANDROID_PLATFORM_VERSION}" \
-              "gcr.io/iree-oss/openxla-benchmark/android@sha256:4c219b20213c7f26edf848b6924a95991c0d359463709d4ee55b499eeaa3349b" \
+              "gcr.io/iree-oss/openxla-benchmark/android@sha256:89921d48b01eade590bc84ba4ee7eda2fb436cbe041961d2b14ff2884fdbab7c" \
                 ./experimental/iree/build_iree.sh \
                     "${X86_BUILD_DIR}" \
                     "${ANDROID_BUILD_DIR}"
@@ -159,7 +160,7 @@ jobs:
 
           docker run --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
             --env OOBI_TEMP_DIR="${TF_SOURCE_DIR}" \
-            "gcr.io/iree-oss/openxla-benchmark/tflite-android@sha256:e02858dd5998dbae6798e11c8646bcff59135ffeb35e0f08a52977106979f1d2" \
+            "gcr.io/iree-oss/openxla-benchmark/tflite-android@sha256:90a46d61320cfeb43e344624837fd707ac1b045a400562582773c3eec591efbd" \
               ./experimental/tflite/build_tflite.sh \
                   "${X86_BUILD_DIR}" \
                   "${ANDROID_BUILD_DIR}"
@@ -226,7 +227,7 @@ jobs:
           docker run --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
             --env IREE_COMPILE_PATH="${IREE_COMPILE_PATH}" \
             --env OOBI_TEMP_DIR="${IREE_TEMP_DIR}" \
-            "gcr.io/iree-oss/openxla-benchmark/base@sha256:2dbee52eaa63e62137682f0eda701ac4cf59b8e16395daa757f6e1906b52dd82" \
+            "gcr.io/iree-oss/openxla-benchmark/base@sha256:5e13cda13bdbd6fcaeeae66f2d2ffe5a8c5ef88f763fcd489db7b87e04f40ccd" \
               ./experimental/iree/compile_workloads.sh \
                   "${TARGET_DEVICE}" \
                   "${BENCHMARK_ARTIFACTS_DIR}"
@@ -239,7 +240,7 @@ jobs:
           docker run --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
             --env OOBI_IREE_BENCHMARK_MODULE_PATH="${IREE_INSTALL_DIR}/bin/iree-benchmark-module" \
             --env OOBI_IREE_RUN_MODULE_PATH="${IREE_INSTALL_DIR}/bin/iree-run-module" \
-            "gcr.io/iree-oss/openxla-benchmark/base@sha256:2dbee52eaa63e62137682f0eda701ac4cf59b8e16395daa757f6e1906b52dd82" \
+            "gcr.io/iree-oss/openxla-benchmark/base@sha256:5e13cda13bdbd6fcaeeae66f2d2ffe5a8c5ef88f763fcd489db7b87e04f40ccd" \
             ./experimental/iree/benchmark_iree_x86.sh \
               "${TARGET_DEVICE}" \
               "${BENCHMARK_ARTIFACTS_DIR}" \
@@ -265,7 +266,7 @@ jobs:
           docker run --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
             --env OOBI_TEMP_DIR="${TFLITE_TEMP_DIR}" \
             --env TFLITE_BENCHMARK_BINARY="${TFLITE_BUILD_DIR}/benchmark_model" \
-            "gcr.io/iree-oss/openxla-benchmark/base@sha256:2dbee52eaa63e62137682f0eda701ac4cf59b8e16395daa757f6e1906b52dd82" \
+            "gcr.io/iree-oss/openxla-benchmark/base@sha256:5e13cda13bdbd6fcaeeae66f2d2ffe5a8c5ef88f763fcd489db7b87e04f40ccd" \
             ./experimental/tflite/benchmark_tflite_x86.sh \
               "${TARGET_DEVICE}" \
               "${TFLITE_RESULTS_FILE}"
@@ -288,7 +289,7 @@ jobs:
           gcloud storage cp "${TFLITE_RESULTS_GCS_PATH}" "${TFLITE_RESULTS_FILE}"
 
           docker run --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
-            "gcr.io/iree-oss/openxla-benchmark/base@sha256:2dbee52eaa63e62137682f0eda701ac4cf59b8e16395daa757f6e1906b52dd82" \
+            "gcr.io/iree-oss/openxla-benchmark/base@sha256:5e13cda13bdbd6fcaeeae66f2d2ffe5a8c5ef88f763fcd489db7b87e04f40ccd" \
             ./experimental/utils/generate_iree_vs_tflite_report.sh \
               "${IREE_RESULTS_FILE}" \
               "${TFLITE_RESULTS_FILE}" \
@@ -343,7 +344,7 @@ jobs:
           docker run --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
             --env IREE_COMPILE_PATH="${IREE_COMPILE_PATH}" \
             --env OOBI_TEMP_DIR="${OOBI_TEMP_DIR}" \
-            "gcr.io/iree-oss/openxla-benchmark/android@sha256:4c219b20213c7f26edf848b6924a95991c0d359463709d4ee55b499eeaa3349b" \
+            "gcr.io/iree-oss/openxla-benchmark/android@sha256:89921d48b01eade590bc84ba4ee7eda2fb436cbe041961d2b14ff2884fdbab7c" \
               ./experimental/iree/compile_workloads.sh \
                   "${TARGET_DEVICE}" \
                   "${BENCHMARK_ARTIFACTS_DIR}"

diff --git a/.github/workflows/update_model_artifacts.yml b/.github/workflows/update_model_artifacts.yml
@@ -10,6 +10,7 @@ name: Update Model Artifacts
 
 on:
   workflow_dispatch:
+  pull_request:
 
 concurrency:
   # A PR number if a pull request and otherwise the commit hash. This cancels
@@ -60,7 +61,7 @@ jobs:
             --env "AUTO_UPLOAD=0" \
             --env "JOBS=1" \
             --env "OUTPUT_DIR=jax" \
-            "gcr.io/iree-oss/openxla-benchmark/base@sha256:2dbee52eaa63e62137682f0eda701ac4cf59b8e16395daa757f6e1906b52dd82" \
+            "gcr.io/iree-oss/openxla-benchmark/base@sha256:5e13cda13bdbd6fcaeeae66f2d2ffe5a8c5ef88f763fcd489db7b87e04f40ccd" \
             "common_benchmark_suite/openxla/benchmark/comparative_suite/jax/scripts/generate_model_artifacts.sh" \
               "RESNET50_FP32_JAX_.+" \
               "BERT_LARGE_FP32_JAX_.+_BATCH(1|16|24|32|48|64|512)" \

diff --git a/comparative_benchmark/xla_hlo/build_xla_tools.sh b/comparative_benchmark/xla_hlo/build_xla_tools.sh
@@ -5,22 +5,33 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
+VENV_DIR="${OOBI_VENV_DIR:-xla-tools.venv}"
+PYTHON="${PYTHON:-/usr/bin/python3}"
 BUILD_DIR="${XLA_BUILD_DIR:-build-xla}"
 TOOLS_OUTPUT_DIR="${1:-"${XLA_TOOLS_OUTPUT_DIR}"}"
 CUDA_VERSION="${2:-"${XLA_CUDA_VERSION}"}"
 
+TD="$(cd $(dirname $0) && pwd)"
+VENV_DIR="${VENV_DIR}" PYTHON="${PYTHON}" source "${TD}/setup_venv.sh"
+
 mkdir -p "${BUILD_DIR}"
 pushd "${BUILD_DIR}"
 
 git clone https://github.com/openxla/xla.git
 cd xla
 
+pip list
+
 # Log the git version of XLA repo.
 git log --oneline --graph --max-count=1
 
 bazel build -c opt --config=cuda \
-  --action_env TF_CUDA_COMPUTE_CAPABILITIES="8.0" \
   --action_env GCC_HOST_COMPILER_PATH="/usr/bin/x86_64-linux-gnu-gcc-11" \
+  --action_env GCC="/usr/bin/x86_64-linux-gnu-gcc-11" \
+  --action_env TF_NEED_ROCM=0 \
+  --action_env TF_NEED_CUDA=1 \
+  --action_env TF_CUDA_CLANG=0 \
+  --action_env TF_CUDA_COMPUTE_CAPABILITIES="8.0" \
   --action_env LD_LIBRARY_PATH="/usr/local/cuda-${CUDA_VERSION}/lib64:" \
   --action_env CUDA_TOOLKIT_PATH="/usr/local/cuda-${CUDA_VERSION}" \
   --copt=-Wno-switch \
@@ -29,6 +40,7 @@ RUN_HLO_RUNNER_MAIN_PATH="$(realpath bazel-bin/xla/tools/multihost_hlo_runner/hl
 
 bazel build -c opt --copt=-Wno-switch \
   --action_env GCC_HOST_COMPILER_PATH="/usr/bin/x86_64-linux-gnu-gcc-11" \
+  --action_env GCC="/usr/bin/x86_64-linux-gnu-gcc-11" \
   xla/tools:run_hlo_module
 RUN_HLO_MODULE_PATH="$(realpath bazel-bin/xla/tools/run_hlo_module)"
 

diff --git a/comparative_benchmark/xla_hlo/requirements.txt b/comparative_benchmark/xla_hlo/requirements.txt
@@ -1,2 +1,3 @@
 numpy
 requests
+scipy
diff --git a/devtools/docker/dockerfiles/android.Dockerfile b/devtools/docker/dockerfiles/android.Dockerfile
@@ -6,12 +6,29 @@
 
 # An image for cross-compiling towards Android.
 
-FROM gcr.io/iree-oss/openxla-benchmark/base@sha256:2dbee52eaa63e62137682f0eda701ac4cf59b8e16395daa757f6e1906b52dd82
+FROM gcr.io/iree-oss/openxla-benchmark/base@sha256:5e13cda13bdbd6fcaeeae66f2d2ffe5a8c5ef88f763fcd489db7b87e04f40ccd
 
-ARG NDK_VERSION=r26b
+######## Python ########
+WORKDIR /install-python
+
+ARG PYTHON_VERSION=3.11
+
+COPY devtools/docker/context/python_build_requirements.txt devtools/docker/context/install_python_deps.sh ./
+RUN ./install_python_deps.sh "${PYTHON_VERSION}" \
+  && apt-get -y install python-is-python3 \
+  && rm -rf /install-python
+
+ENV PYTHON_BIN /usr/bin/python3
+ENV TF_PYTHON_VERSION "${PYTHON_VERSION}"
+
+WORKDIR /
+
+######## Android NDK ########
+ARG NDK_VERSION=r25c
 WORKDIR /install-ndk
 
 ENV ANDROID_NDK "/usr/src/android-ndk-${NDK_VERSION}"
+ENV ANDROID_NDK_API_LEVEL "25"
 
 RUN wget -q "https://dl.google.com/android/repository/android-ndk-${NDK_VERSION}-linux.zip" \
     && unzip -q "android-ndk-${NDK_VERSION}-linux.zip" -d /usr/src/  \