Create a workflow to run benchmarks

.github/workflows/benchmarks.yml

@@ -0,0 +1,231 @@
+name: Benchmarks
+
+on:
+  pull_request:
+    branches:
+      - main
+  workflow_dispatch:
+    inputs:
+      halt-for-connection:
+        description: 'Should this workflow run wait for a remote connection?'
+        type: choice
+        required: true
+        default: 'no'
+        options:
+        - 'yes'
+        - 'no'
+
+jobs:
+  build-xla-gpu-and-test:
+    runs-on: "linux-x86-g2-48-l4-4gpu" #linux-x86-n2-16 # Use a GPU-enabled runner 
+    container:
+      image: "gcr.io/tensorflow-testing/nosla-cuda12.3-cudnn9.1-ubuntu20.04-manylinux2014-multipython:latest"
+      options: --gpus all --privileged  # Might need privileged mode, use with caution
+
+    steps:
+      - name: Checkout XLA
+        uses: actions/checkout@b4ffde65f46336ab88eb53be808477a3936bae11 # v4.1.1
+        with:
+            repository: openxla/xla
+      # - name: Checkout repository
+      #   uses: actions/checkout@v3
+      #   with:
+      #     repository: juliagmt-google/xla  
+      #     path: xla
+
+      # - name: Print machine specs
+      #   run: |
+      #     nvidia-smi
+      #     free -h  # Memory information
+      #     df -h    # Disk space information
+      #     uname -a # Kernel information
+
+      # - name: Set WORKSPACE_DIR
+      #   env:
+      #     WORKSPACE_DIR: ${{ github.workspace }}
+      #   run: |
+      #     echo "WORKSPACE_DIR is: $WORKSPACE_DIR"
+
+      # - name: Create results directory
+      #   run: |
+      #     mkdir -p results
+      #     ls
+
+      # - name: Configure XLA for GPU backend
+      #   run: |
+      #     cd xla
+      #     ./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: Wait For Connection
+        uses: google-ml-infra/actions/ci_connection@main
+        with:
+          halt-dispatch-input: ${{ inputs.halt-for-connection }}
+
+      # - 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-critical HLOs to benchmark.
+      # - name: Run hlo_opt and generate xspace.pb
+      #   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=$WORKSPACE_DIR/results/xspace.pb $WORKSPACE_DIR/xla/tools/hlo_opt/tests/gpu_hlo_pass.hlo
+
+      # - name: Set XSPACE_PATH
+      #   env:
+      #     XSPACE_PATH: ${{ github.workspace }}/results/xspace.pb
+      #   run: |
+      #     echo "XSPACE_PATH is: $XSPACE_PATH"
+
+      # - name: Compute the cost of gpu_hlo_pass.hlo
+      #   run: |
+      #     bazel run //xla/tools:compute_cost -- --input=$WORKSPACE_DIR/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: Compute the device stats of gpu_hlo_pass.hlo
+      #   run: |
+      #     echo "XSPACE_PATH is: $XSPACE_PATH"
+      #     bazel run //xla/tools:get_device_stats_main -- --input=$XSPACE_PATH
+      #   working-directory: juliagmt-google-xla
+
+      # - name: Upload XSpace
+      #   uses: actions/upload-artifact@65c4c4a1ddee5b72f698fdd19549f0f0fb45cf08 # v4.6.0
+      #   with:
+      #     name: gpu-xla-benchmarks-xspace
+      #     path: $WORKSPACE_DIR/results
+
+      # - name: Print machine specs
+      #   run: |
+      #     lscpu
+      #     free -h  # Memory information
+      #     df -h    # Disk space information
+      #     uname -a # Kernel information
+
+      # - name: Create results directory
+      #   working-directory: xla
+      #   run: mkdir results
+
+      # - name: Wait For Connection
+      #   uses: google-ml-infra/actions/ci_connection@main
+      #   with:
+      #     halt-dispatch-input: ${{ inputs.halt-for-connection }}
+
+      # - name: Configure XLA
+      #   working-directory: xla
+      #   run: |
+      #     cd ..
+      #     ls
+      #     ./configure.py --backend CUDA --nccl
+
+  #     - name: Set TF_CPP_MAX_VLOG_LEVEL
+  #       working-directory: xla
+  #       run: echo "TF_CPP_MAX_VLOG_LEVEL=1" >> $GITHUB_ENV  # Use GITHUB_ENV to persist across steps
+
+      # - name: Build hlo_runner_main
+      #   working-directory: xla
+      #   run: bazel build -c opt --config=cuda --dynamic_mode=off //xla/tools/multihost_hlo_runner:hlo_runner_main
+
+      # - name: Wait For Connection
+      #   uses: google-ml-infra/actions/ci_connection@main
+      #   with:
+      #     halt-dispatch-input: ${{ inputs.halt-for-connection }}
+
+      # - name: Wait For Connection
+      #   uses: google-ml-infra/actions/ci_connection@main
+      #   with:
+      #     halt-dispatch-input: ${{ inputs.halt-for-connection }}
+
+      # - name: Run an HLO file
+      #   working-directory: xla
+      #   run: |
+      #     cd ..
+      #     ./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.pb xla/tools/hlo_opt/tests/gpu_hlo_pass.hlo 
+
+      # - name: Get Device Stats
+      #   working-directory: xla
+      #   run: |
+      #     cd ..
+      #     PWD=$(pwd)
+      #     bazel run //xla/tools:get_device_stats_main -- --input=$PWD/xla/results/xspace.pb
+
+      # - name: Wait For Connection
+      #   uses: google-ml-infra/actions/ci_connection@main
+      #   with:
+      #     halt-dispatch-input: ${{ inputs.halt-for-connection }}
+
+      # - name: Download parse_xla_logs.py
+      #   working-directory: xla
+      #   run: wget https://raw.githubusercontent.com/juliagmt-google/xla/main/.github/workflows/parse_xla_logs.py  
+
+      # - name: Parse XLA logs
+      #   working-directory: xla
+      #   run: python parse_xla_logs.py results/gpu_hlo_backend.log
+
+      # - name: Upload Results
+      #   uses: actions/upload-artifact@v4
+      #   with:
+      #     name: gpu-xla-benchmarks
+      #     path: xla/results
+
+
+# jobs:
+#   Tests:
+#     strategy:
+#       # Don't fail fast - want to see results for all builds even if one fails.
+#       fail-fast: false
+#       matrix:
+#         job_info:
+#           - os: "linux-x86-g2-48-l4-4gpu"
+#             container: "us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:latest",
+#             pretty_name: "Linux X86 runner with 4 NVIDIA L4 GPUs"
+#           # Expect more GPU types in the future.
+#     name: ${{ matrix.job_info.pretty_name }}
+#     runs-on: ${{ matrix.job_info.os }}
+#     container: ${{ matrix.job_info.container }}
+#     defaults:
+#       run:
+#         shell: bash
+#     timeout-minutes: 360
+#     steps:
+#       - name: Checkout XLA
+#         uses: actions/checkout@v4 # v4.1.1
+#         with:
+#            repository: openxla/xla  
+#       - name: Wait For Connection
+#         uses: google-ml-infra/actions/ci_connection@main
+#         with:
+#           halt-dispatch-input: ${{ inputs.halt-for-connection }}
+#       - 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: ./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: Upload XSpace
+#         uses: actions/upload-artifact@v4 # v4.1.1
+#         with:
+#           name: gpu-xla-benchmarks-xspace
+#           path: xla/results/xspace.pbtxt

.github/workflows/parse_xla_logs.py

@@ -0,0 +1,120 @@
+import re
+import sys
+
+class BenchmarkResult:  # Define the class here
+    def __init__(self):
+        self.name = ""
+        self.flavor = ""
+        self.description = ""
+        self.xla_flags = ""
+        self.build_flags = []
+        self.source_uri = ""
+        self.device_name = ""
+        self.time_us = 0.0
+        self.last_average_loss = 0.0
+        self.memory_read_written = 0
+        self.buffer_allocation = []
+        self.nvptx_compilation_time_us = 0
+        self.hlo_passes_time_us = 0
+        self.run_backend_time_us = 0
+        self.device_time_us = 0.0
+        self.device_memcpy_time_us = 0.0
+        self.xprof_session_id = ""
+        self.error = ""
+        self.custom_metrics = []
+
+class LogParser:
+    def __init__(self):
+        self.compilation_time = {}  # You'll need to initialize this appropriately
+    # Units defined in google3/third_party/tensorflow/tsl/platform/numbers.cc
+    _TIME_UNITS = {"us": 1e-6, "ms": 1e-3, "s": 1, "min": 60, "h": 3600}
+    _TIME_REGEXP = re.compile(r"time: (\d+\.?\d*) (%s)" % "|".join(_TIME_UNITS))
+    _SIZE_UNITS = "KMGTPE"
+    _SIZE_REGEXP = re.compile(r"(\d+\.?\d*)([%s]i)?B" % _SIZE_UNITS)
+
+    def _ParseTimeFromLog(self, time_str: bytes, position: int = -1) -> int:
+      """Returns the time in microseconds parsed from XLA logs."""
+      match = self._TIME_REGEXP.search(time_str.decode())
+      assert match, "Unable to parse the time on log line %d" % position
+      exp_ = self._TIME_UNITS[match.group(2)] * 1e6
+      return int(float(match.group(1)) * exp_)
+
+    def _ParseBytesFromLog(self, size_str: bytes, position: int = -1) -> int:
+      """Returns the size in bytes parsed from XLA logs."""
+      match = self._SIZE_REGEXP.search(size_str.decode())
+      assert match, "Unable to parse the size on log line %d" % position
+      prefix = match.group(2) or "-"
+      exp_ = pow(2, 10 * (self._SIZE_UNITS.find(prefix[0]) + 1))
+      return int(float(match.group(1)) * exp_)
+
+    def _ParseLogLine(self, result, line, position) -> None:
+    """Parses a single XLA log line and updates the BenchmarkResult proto.
+
+    It looks for certain text patterns and updates the result proto with
+    compilation stats and Xprof session id.
+
+    Args:
+      result: a benchmark result proto to be updated
+      line: a single line of benchmark log output
+      position: line number (for debugging)
+    """
+      # Log output generated by --vmodule=nvptx_compiler=1
+      if b"NVPTXCompiler::CompileTargetBinary - CompileToPtx" in line:
+        thread_id = re.split(rb"\s+", line)[2]
+        self.compilation_time[thread_id] += self._ParseTimeFromLog(line, position)
+        return
+
+      # Log output generated by --vmodule=gpu_compiler=1
+      if b"HLO memory read+written:" in line:
+        result.memory_read_written += self._ParseBytesFromLog(line, position)
+        return
+      if b"GpuCompiler::RunHloPasses for" in line:
+        result.hlo_passes_time_us += self._ParseTimeFromLog(line, position)
+        return
+      if b"GpuCompiler::RunBackend for" in line:
+        result.run_backend_time_us += self._ParseTimeFromLog(line, position)
+        return
+
+      # Log output generated by --vmodule=bfc_allocator=2
+      if b"New Peak memory usage" in line and b"for GPU" in line:
+        match = re.search(rb"(\d+) bytes", line)
+        assert match, "Unable to parse the size on log line %d" % position
+        alloc_size = int(match.group(1))
+        if alloc_size > max(result.buffer_allocation, default=0):
+          result.buffer_allocation[:] = [alloc_size]
+        return
+
+      # Log output generated by --xprof_end_2_end_upload
+      if b"XprofResponse uploaded to http://xprof/" in line:
+        match = re.search(rb"session_id=([\w\-]+)", line)
+        assert match, "Unable to parse the XProf link"
+        result.xprof_session_id = match.group(1)
+      return
+
+    def parse_log_file(self, log_file_path):
+        result = BenchmarkResult()
+        with open(log_file_path, "rb") as f:
+            for i, line in enumerate(f):
+                self._ParseLogLine(result, line, i + 1)
+
+        # Update compilation time (total across threads) in the result object
+        for time_taken in self.compilation_time.values():
+            result.nvptx_compilation_time_us += time_taken
+        return result  # Return the result object
+
+if __name__ == "__main__":
+    if len(sys.argv) != 2:
+        print("Usage: python parse_log.py <log_file_path>")
+        sys.exit(1)
+
+    log_file_path = sys.argv[1]
+
+    parser = LogParser()
+    result = parser.parse_log_file(log_file_path)
+
+    # Access results
+    print(f"Memory Read+Written: {result.memory_read_written} bytes")
+    print(f"NVPTX Compilation Time: {result.nvptx_compilation_time_us} us")
+    print(f"HLO Passes Time: {result.hlo_passes_time_us} us")
+    print(f"Run Backend Time: {result.run_backend_time_us} us")
+    print(f"Buffer Allocation: {result.buffer_allocation} bytes")