Update cpu_benchmarks.yml and gpu_benchmarks.yml to run the Gemma2-2B…

… HLO. PiperOrigin-RevId: 731434254
openxla · Feb 27, 2025 · dab6fe6 · dab6fe6
1 parent 0855570
commit dab6fe6
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 49 deletions.
diff --git a/.github/workflows/cpu_benchmarks.yml b/.github/workflows/cpu_benchmarks.yml
@@ -99,44 +99,24 @@ jobs:
       - name: Build run_hlo_module
         run: bazel build -c opt --dynamic_mode=off xla/tools:run_hlo_module
 
-      - name: Run HLO Module Benchmarks
-        continue-on-error: true
-        run: |
-          for file in xla/tests/fuzz/*.hlo; do
-            filename=$(basename "$file")
-            # Skip expected failed hlo files.
-            if [[ "$filename" == "rand_000060.hlo" || "$filename" == "rand_000067.hlo" || "$filename" == "rand_000072.hlo" ]]; then
-              echo "Skipping benchmark on $file"
-              continue
-            fi
-            echo "Running benchmark on $file"
-            ./bazel-bin/xla/tools/run_hlo_module --input_format=hlo --platform=CPU "$file"
-          done
-
-      - name: Build CPU Benchmarks
-        run: bazel build -c opt --dynamic_mode=off //xla/backends/cpu/benchmarks:*
+      - name: Run xla/tests/fuzz/rand_000001.hlo and collect runtime and compile time
+        run: ./bazel-bin/xla/tools/run_hlo_module --input_format=hlo --platform=CPU xla/tests/fuzz/rand_000001.hlo
 
-      - name: Run CPU benchmarks
-        continue-on-error: true
+      - name: Download and run Gemma2 2B HLO and collect runtime and compile time
         run: |
-          find ./bazel-bin/xla/backends/cpu/benchmarks/ -maxdepth 1 -type f -executable -name "*_test" -print0 | while IFS= read -r -d $'\0' benchmark; do
-            benchmark_name=$(basename "$benchmark" | sed 's/_test$//') 
-            echo "Running benchmark: $benchmark_name"
+          mkdir -p tmp_hlo
+          cd tmp_hlo
+          wget https://storage.googleapis.com/xla-benchmarking-temp/gemma2_2b_keras_jax.hlo gemma2_2b_keras_jax.hlo
+          cd ..
 
-            # Run the benchmark with default parameters.
-            $benchmark --benchmark_filter=".*"
-            $benchmark --benchmark_filter=".*" > "results/$benchmark_name.log" 2>&1
+      - name: Compute the cost of gemma2_2b_keras_jax.hlo
+        run: |
+          PWD=$(pwd)
+          bazel run //xla/tools:compute_cost -- --input=$PWD/tmp_hlo/gemma2_2b_keras_jax.hlo --format=hlo
 
-            # Check the exit code of the benchmark
-            if [ $? -ne 0 ]; then
-              echo "Error: Benchmark '$benchmark_name' failed. Check the log file: results/$benchmark_name.log"
-            else
-              echo "Benchmark '$benchmark_name' completed successfully."
-            fi
-          done
+      - name: Build hlo_runner_main
+        run: bazel build -c opt --dynamic_mode=off //xla/tools/multihost_hlo_runner:hlo_runner_main
 
-      - name: Upload Results ${{ matrix.build_target }}
-        uses: actions/upload-artifact@65c4c4a1ddee5b72f698fdd19549f0f0fb45cf08 # v4.6.0
-        with:
-          name: cpu-xla-benchmarks-${{ matrix.job_info.os }}
-          path: xla/results
+      - name: Run hlo_runner_main on gemma2_2b_keras_jax.hlo and collect profile stats
+        run: |
+          ./bazel-bin/xla/tools/multihost_hlo_runner/hlo_runner_main --device_type=host --log_output=True --use_spmd_partitioning --profile_execution=True tmp_hlo/gemma2_2b_keras_jax.hlo
diff --git a/.github/workflows/gpu_benchmarks.yml b/.github/workflows/gpu_benchmarks.yml
@@ -69,30 +69,46 @@ jobs:
         run: bazel build -c opt --config=cuda --dynamic_mode=off //xla/tools/multihost_hlo_runner:hlo_runner_main
 
       # TODO(juliagmt): Add more performance-critical HLOs to benchmark.
-      - name: Run hlo_opt and generate xspace.pb
+      - name: Run gpu_hlo_pass.hlo and dump XSpace with profile stats
         run: |
-          ./bazel-bin/xla/tools/multihost_hlo_runner/hlo_runner_main --device_type=gpu --log_output=True --use_spmd_partitioning --xla_gpu_dump_xspace_to=results/xspace.pb xla/tools/hlo_opt/tests/gpu_hlo_pass.hlo
+          ./bazel-bin/xla/tools/multihost_hlo_runner/hlo_runner_main --device_type=gpu --log_output=True --use_spmd_partitioning --profile_execution=True --xla_gpu_dump_xspace_to=results/gpu_hlo_pass_xspace.pb xla/tools/hlo_opt/tests/gpu_hlo_pass.hlo
 
       - name: Compute the cost of gpu_hlo_pass.hlo
         run: |
           PWD=$(pwd)
           bazel run //xla/tools:compute_cost -- --input=$PWD/xla/tools/hlo_opt/tests/gpu_hlo_pass.hlo --format=hlo --gpu
 
-      - name: Checkout juliagmt-google/xla
-        uses: actions/checkout@b4ffde65f46336ab88eb53be808477a3936bae11 # v4.1.1
-        with:
-          repository: juliagmt-google/xla
-          path: juliagmt-google-xla
+      - name: Download and run Gemma2 2B HLO and collect runtime and compile time
+        run: |
+          mkdir -p tmp_hlo
+          cd tmp_hlo
+          wget https://storage.googleapis.com/xla-benchmarking-temp/gemma2_2b_keras_jax.hlo gemma2_2b_keras_jax.hlo
+          cd ..
 
-      - name: Compute the device stats of gpu_hlo_pass.hlo
-        env:
-          XSPACE_PATH: ${{ github.workspace }}/results/xspace.pb
+      - name: Run gemma2_2b_keras_jax.hlo and dump XSpace with profile stats
         run: |
-          bazel run //xla/tools:get_device_stats_main -- --input="$XSPACE_PATH"
-        working-directory: juliagmt-google-xla
+          ./bazel-bin/xla/tools/multihost_hlo_runner/hlo_runner_main --device_type=gpu --log_output=True --use_spmd_partitioning --xla_gpu_dump_xspace_to=results/gemma2_2b_keras_jax_xspace.pb tmp_hlo/gemma2_2b_keras_jax.hlo
+
+      - name: Compute the cost of gemma2_2b_keras_jax.hlo
+        run: |
+          PWD=$(pwd)
+          bazel run //xla/tools:compute_cost -- --input=$PWD/tmp_hlo/gemma2_2b_keras_jax.hlo --format=hlo --gpu
+
+      - name: Build run_hlo_module
+        run: bazel build -c opt --config=cuda --dynamic_mode=off xla/tools:run_hlo_module
+
+      - name: Run gpu_hlo_pass.hlo and collect runtime and compile time
+        run: |
+          ./bazel-bin/xla/tools/run_hlo_module --input_format=hlo --platform=GPU xla/tools/hlo_opt/tests/gpu_hlo_pass.hlo
+
+      - name: Compute device stats
+        run: |
+          PWD=$(pwd)
+          bazel run //xla/tools:compute_gpu_device_stats_main -- --input=$PWD/results/gpu_hlo_pass_xspace.pb
+          bazel run //xla/tools:compute_gpu_device_stats_main -- --input=$PWD/results/gemma2_2b_keras_jax_xspace.pb
 
       - name: Upload XSpace
         uses: actions/upload-artifact@65c4c4a1ddee5b72f698fdd19549f0f0fb45cf08 # v4.6.0
         with:
           name: gpu-xla-benchmarks-xspace
-          path: results/xspace.pb
+          path: results/