From 217aca2ac36254799872f9503b9c6254708ca6ff Mon Sep 17 00:00:00 2001
From: Jerry Wu <cheyuw@google.com>
Date: Mon, 7 Aug 2023 18:07:46 -0400
Subject: [PATCH] Reduce CPU benchmark run time (#116)

---
 .../jax_xla/benchmark_all.sh                  | 25 ++++++++++---------
 comparative_benchmark/tf_xla/benchmark_all.sh | 25 ++++++++++---------
 .../xla_hlo/benchmark_all.sh                  | 25 ++++++++++---------
 3 files changed, 39 insertions(+), 36 deletions(-)

diff --git a/comparative_benchmark/jax_xla/benchmark_all.sh b/comparative_benchmark/jax_xla/benchmark_all.sh
index 2fc6d5ae..b24ef026 100755
--- a/comparative_benchmark/jax_xla/benchmark_all.sh
+++ b/comparative_benchmark/jax_xla/benchmark_all.sh
@@ -64,7 +64,7 @@ if [ "${TARGET_DEVICE}" = "a2-highgpu-1g" ]; then
   ITERATIONS=50
 elif [ "${TARGET_DEVICE}" = "c2-standard-16" ]; then
   BENCHMARK_NAMES=("${CPU_BENCHMARK_NAMES[@]}")
-  ITERATIONS=20
+  ITERATIONS=5
 else
   echo "Unsupported target device ${TARGET_DEVICE}."
   exit 1
@@ -82,16 +82,17 @@ for benchmark_name in "${BENCHMARK_NAMES[@]}"; do
     --verbose
 done
 
+# Disable for now as it takes too long to run.
 # If running on CPU, also benchmark XLA CPU-Next.
 # Use a lower number of iterations since CPU-Next is slow.
-if [ "${TARGET_DEVICE}" = "c2-standard-16" ]; then
-  for benchmark_name in "${BENCHMARK_NAMES[@]}"; do
-    "${TD}/run_benchmarks.py" \
-      --benchmark_name="${benchmark_name}" \
-      --target_device="${TARGET_DEVICE}" \
-      --output="${OUTPUT_PATH}" \
-      --iterations=3 \
-      --compiler="xla_cpu_next" \
-      --verbose
-  done
-fi
+# if [ "${TARGET_DEVICE}" = "c2-standard-16" ]; then
+#   for benchmark_name in "${BENCHMARK_NAMES[@]}"; do
+#     "${TD}/run_benchmarks.py" \
+#       --benchmark_name="${benchmark_name}" \
+#       --target_device="${TARGET_DEVICE}" \
+#       --output="${OUTPUT_PATH}" \
+#       --iterations=3 \
+#       --compiler="xla_cpu_next" \
+#       --verbose
+#   done
+# fi
diff --git a/comparative_benchmark/tf_xla/benchmark_all.sh b/comparative_benchmark/tf_xla/benchmark_all.sh
index f13cbb45..7d9e1809 100755
--- a/comparative_benchmark/tf_xla/benchmark_all.sh
+++ b/comparative_benchmark/tf_xla/benchmark_all.sh
@@ -60,7 +60,7 @@ if [ "${TARGET_DEVICE}" = "a2-highgpu-1g" ]; then
   ITERATIONS=50
 elif [ "${TARGET_DEVICE}" = "c2-standard-16" ]; then
   BENCHMARK_NAMES=("${CPU_BENCHMARK_NAMES[@]}")
-  ITERATIONS=20
+  ITERATIONS=5
 else
   echo "Unsupported target device ${TARGET_DEVICE}."
   exit 1
@@ -78,16 +78,17 @@ for benchmark_name in "${BENCHMARK_NAMES[@]}"; do
     --verbose
 done
 
+# Disable for now as it takes too long to run.
 # If running on CPU, also benchmark XLA CPU-Next.
 # Use a lower number of iterations since CPU-Next is slow.
-if [ "${TARGET_DEVICE}" = "c2-standard-16" ]; then
-  for benchmark_name in "${BENCHMARK_NAMES[@]}"; do
-    "${TD}/run_benchmarks.py" \
-      --benchmark_name="${benchmark_name}" \
-      --target_device="${TARGET_DEVICE}" \
-      --output="${OUTPUT_PATH}" \
-      --iterations=3 \
-      --compiler="xla_cpu_next" \
-      --verbose
-  done
-fi
+# if [ "${TARGET_DEVICE}" = "c2-standard-16" ]; then
+#   for benchmark_name in "${BENCHMARK_NAMES[@]}"; do
+#     "${TD}/run_benchmarks.py" \
+#       --benchmark_name="${benchmark_name}" \
+#       --target_device="${TARGET_DEVICE}" \
+#       --output="${OUTPUT_PATH}" \
+#       --iterations=3 \
+#       --compiler="xla_cpu_next" \
+#       --verbose
+#   done
+# fi
diff --git a/comparative_benchmark/xla_hlo/benchmark_all.sh b/comparative_benchmark/xla_hlo/benchmark_all.sh
index 8a3d8029..1636eff9 100755
--- a/comparative_benchmark/xla_hlo/benchmark_all.sh
+++ b/comparative_benchmark/xla_hlo/benchmark_all.sh
@@ -95,17 +95,18 @@ for benchmark_name in "${BENCHMARK_NAMES[@]}"; do
     --verbose
 done
 
+# Disable for now as it takes too long to run.
 # If running on CPU, also benchmark XLA CPU-Next.
 # Use a lower number of iterations since CPU-Next is slow.
-if [ "${TARGET_DEVICE}" = "c2-standard-16" ]; then
-  for benchmark_name in "${BENCHMARK_NAMES[@]}"; do
-    "${TD}/run_benchmarks.py" \
-      --benchmark_name="${benchmark_name}" \
-      --target_device="${TARGET_DEVICE}" \
-      --hlo-tool="${XLA_TOOLS_DIR}/${HLO_TOOL}" \
-      --output="${OUTPUT_PATH}" \
-      --iterations=3 \
-      --compiler="xla_cpu_next" \
-      --verbose
-  done
-fi
+# if [ "${TARGET_DEVICE}" = "c2-standard-16" ]; then
+#   for benchmark_name in "${BENCHMARK_NAMES[@]}"; do
+#     "${TD}/run_benchmarks.py" \
+#       --benchmark_name="${benchmark_name}" \
+#       --target_device="${TARGET_DEVICE}" \
+#       --hlo-tool="${XLA_TOOLS_DIR}/${HLO_TOOL}" \
+#       --output="${OUTPUT_PATH}" \
+#       --iterations=3 \
+#       --compiler="xla_cpu_next" \
+#       --verbose
+#   done
+# fi