Update benchmarks.yml

Author: juliagmt-google

.github/workflows/benchmarks.yml

@@ -28,12 +28,52 @@ jobs:
     steps:
       - name: Checkout XLA
         uses: actions/checkout@b4ffde65f46336ab88eb53be808477a3936bae11 # v4.1.1
+        with:
+          path: "openxla/xla"
+
+      - name: Print machine specs
+        run: |
+          nvidia-smi
+          free -h  # Memory information
+          df -h    # Disk space information
+          uname -a # Kernel information
 
       - name: Create results directory
+        run: mkdir results
+
+      - name: Configure XLA for GPU backend
+        run: |
+          ./configure.py --backend CUDA --nccl
+
+      - name: Set TF_CPP_MAX_VLOG_LEVEL
+        env:
+          TF_CPP_MAX_VLOG_LEVEL: 1
+        run: |
+          echo "TF_CPP_MAX_VLOG_LEVEL is: $TF_CPP_MAX_VLOG_LEVEL"
+
+      - name: Build hlo_runner_main
+        run: bazel build -c opt --config=cuda --dynamic_mode=off //xla/tools/multihost_hlo_runner:hlo_runner_main
+      # TODO(juliagmt): Add more performance-criticalHLOs to benchmark.
+      - name: Run xla/tools/hlo_opt/tests/gpu_hlo_pass.hlo
         run: |
-          mkdir results
-          ls
-          cd results
+          ./bazel-bin/xla/tools/multihost_hlo_runner/hlo_runner_main --device_type=gpu --log_output=True --use_spmd_partitioning --xla_gpu_dump_xspace_to=xla/results/xspace.pbtxt 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 # Replace with the current SHA of v4.1.1 or later
+        with:
+          repository: juliagmt-google/xla
+          path: juliagmt-google-xla
+
+      - name: Run get_device_stats_main
+        run: |
+          bazel run //xla/tools:get_device_stats_main -- --input=$PWD/xla/results/xspace.pb
+        working-directory: juliagmt-google-xla
+
       - name: Upload Artifacts
         uses: actions/upload-artifact@65c4c4a1ddee5b72f698fdd19549f0f0fb45cf08
         with: