chore: import upstream snapshot with attribution

third_party/xla/.github/workflows/benchmarks/README.md

@@ -0,0 +1,237 @@
+# Onboard New Microbenchmarks to OpenXLA
+
+This guide provides step-by-step instructions for contributing new
+microbenchmarks to OpenXLA microbenchmarking infrastructure.
+
+## Overview
+
+The OpenXLA microbenchmarking system is designed to automatically detect
+performance regressions and track performance trends across different hardware
+backends (CPU, GPU) in presubmit, postsubmit and nightly workflows. By adding
+your microbenchmark, you help ensure XLA's performance remains robust for your
+specific use cases.
+
+The process involves:
+
+1.  **Preparing your Benchmark Artifact:** Ensuring your HLO file is OSS-friendly.
+2.  **Defining the Benchmark Configuration:** Adding an entry to a benchmark registry file.
+3.  **Establishing a Baseline:** Adding initial performance thresholds if your benchmark will run in presubmit/postsubmit/nightly jobs.
+
+## Prerequisites
+
+*   **Benchmark Artifact (HLO):** You should have your benchmark ready as either
+    an HLO text file (`.hlo`) (Note: StableHLO MLIR text file (`.mlir`) will be
+    supported later).
+    *   For small artifacts, you can place them in `xla/tools/benchmarks/hlo/`.
+    *   **[Not yet supported]** For larger artifacts, upload them to a GCS
+        bucket (e.g., `gs://xla-benchmarking-temp/your-benchmark.hlo`) and
+        ensure it's publicly readable.
+*   **GitHub Access:** You'll need to create a Pull Request (PR) to the OpenXLA repository.
+
+## Step-by-Step Guide
+
+### Step 1: Prepare Your Benchmark Artifact
+
+Ensure your HLO file is ready and accessible.
+
+*   **Store in the XLA Repository**
+    1.  Place your `.hlo` file in the `xla/xla/tools/benchmarks/hlo/` directory.
+    2.  Make sure your hlo benchmarks run < 15min/20min/30min for
+        presubmit/postsubmit/nightly workflows.
+
+### Step 2: Define the Benchmark Configuration
+
+You'll need to add a new entry to a benchmark registry YAML file. For most
+community contributions, this will be
+`xla/xla/tools/benchmarks/registries/default_registry.yml`.
+
+Each benchmark configuration is a YAML object with the following key fields:
+
+*   `name`: A unique, descriptive name for your benchmark (e.g., `"my_model_attention_layer"`).
+*   `description`: A brief explanation of what the benchmark measures.
+*   `owner`: Your GitHub handle or relevant team alias (e.g., `"your-github-username@"`).
+*   `input_artifact`:
+    *   `input_format`: Currently we support `HLO_TEXT`, and `STABLEHLO_MLIR` will be supported in the future.
+    *   `artifact_path`: (If stored in repo) Relative path from `xla`, e.g., `xla/tools/benchmarks/hlo/my_new_benchmark.hlo`.
+    *   `artifact_gcs_bucket_path`: (If stored in GCS) Full GCS URL.
+*   `model_source_info`: A list of strings describing the origin of the benchmark (e.g., `["Gemma2 2B"]`).
+*   `hardware_targets`: A list defining on which hardware configurations this
+    benchmark should run. Each target has:
+    *   `hardware_category`: e.g., `GPU_L4`, `CPU_X86`, `GPU_B200`.
+    *   `topology`:
+        *   `num_hosts`: Number of hosts (default: 1).
+        *   `num_devices_per_host`: Number of devices per host (default: 1).
+        *   `multi_host`: `true` or `false`.
+    *   `multi_device`: `true` or `false`.
+    *   `target_metrics`: A list of metrics to collect, e.g., `[GPU_DEVICE_TIME, PEAK_GPU_MEMORY]`.
+    *   `run_frequencies`: When to run this benchmark, e.g., `[PRESUBMIT, POSTSUBMIT]`, `[SCHEDULED]`.
+*   `update_frequency_policy`: How often this benchmark definition should
+    be reviewed, e.g., `QUARTERLY`.
+*   `xla_compilation_flags` (Optional): List of XLA flags, e.g.,
+    `["--xla_gpu_enable_cudnn_fusion=false"]`.
+*   `runtime_flags` (Optional): List of flags for the
+    `multihost_hlo_runner`, e.g., `["--num_repeats=5"]`.
+*   `github_labels` (Optional): GitHub labels to manually trigger this
+    specific benchmark.
+
+**Example: Adding "gemma3\_1b\_flax\_sample\_loop" to `default_registry.yml`**
+
+```yaml
+# xla/xla/tools/benchmarks/registries/default_registry.yml
+benchmarks: [
+  # ... existing benchmarks ...
+  {
+    name: "gemma3_1b_flax_sample_loop"
+    description: "Gemma3 1B in Flax Sample Loop."
+    owner: "company-A@" # Replace with your GitHub handle or team
+    input_artifact: {
+      input_format: HLO_TEXT, # Or STABLEHLO_MLIR
+      artifact_path: "xla/tools/benchmarks/hlo/gemma3_1b_flax_sample_loop.hlo"
+      # Option 2 (for large hlo):
+      #`artifact_gcs_bucket_path`: (If stored in GCS) Full GCS URL (not supported yet).
+    }
+    model_source_info: ["Gemma3 1B"] # Describe the source of your HLO
+    hardware_targets: [{
+      hardware_category: GPU_L4
+      topology: { num_hosts: 1, num_devices_per_host: 1, multi_host: false, multi_device: false }
+      target_metrics: [GPU_DEVICE_TIME, GPU_DEVICE_MEMCPY_TIME]
+      run_frequencies: [PRESUBMIT, POSTSUBMIT] # Run on PRs in presubmit and postsubmit
+      runtime_flags: ["--num_repeats=5"] # Example: run 5 times to reduce noise
+    },
+    {
+      hardware_category: CPU_X86
+      topology: { num_hosts: 1, num_devices_per_host: 1, multi_host: false, multi_device: false }
+      target_metrics: [CPU_TIME, WALL_TIME]
+      run_frequencies: [PRESUBMIT] # Only run on PRs for presubmit
+      runtime_flags: ["--num_repeats=5"]
+    }]
+    update_frequency_policy: QUARTERLY # Review this benchmark definition quarterly
+  }
+]
+```
+
+### Step 3: Establish a Baseline
+
+1.  **Determine Baseline Values:**
+    *   The best way to get initial baseline values is to run your benchmark
+    manually on the target hardware or let it run once in postsubmit after your
+    initial PR (without presubmit blocking) is merged.
+    *   Promote your benchmarks from postsubmit to presubmit once you get
+        stable results for baseline values.
+    *   Run the benchmark multiple times (e.g., using `--num_repeats=5` or
+        more) and take the median or a stable average.
+    *   Benchmarks must run < 15min for presubmit, < 20min for postsubmit and < 30min for nightly.
+
+2.  **Add to `presubmit_baseline.yml`:**
+    Edit the file `xla/xla/tools/benchmarks/baseline/presubmit_baseline.yml`. The
+    key for each entry is the `config_id`.
+
+    **Note on `config_id` generation:**
+    `config_id` follows the below pattern:
+    `"{benchmark_name}_{hardware_category_simplified}_{topology_simplified}_{workflow_type}"`.
+    *   `hardware_category_simplified`: e.g., `l4` (for `GPU_L4`), `b200` (for `GPU_B200`), `x86` (for `CPU_X86`).
+    *   `topology_simplified`: e.g., `1h1d` for 1 host, 1 device. <!-- disableFinding(LINE_OVER_80) -->
+    *   `workflow_type`: e.g., `presubmit`, `postsubmit`, `scheduled`.
+
+    If unsure, you can check the GitHub Actions workflow logs for the
+    `generate_benchmark_matrices.py` script output, which will show the generated
+    `config_id`s.
+
+    For each metric you want to track in presubmit (must be in `target_metrics`
+    in the registry):
+
+    *   `baseline_ms`: The baseline performance in milliseconds.
+    *   `threshold`: The maximum allowed regression percentage (e.g.,
+        `0.30` for 30%).
+    *   **Note on metrics:** Currently, we support `GPU_DEVICE_TIME`,
+        `GPU_DEVICE_MEMCPY_TIME` for GPU, and `CPU_TIME`, `WALL_TIME` for CPU.
+
+**Example: Adding baseline for "gemma3\_1b\_flax\_sample\_loop"**
+
+Assuming the `name` is `"gemma3_1b_flax_sample_loop"`:
+
+*   For `GPU_L4`, `1` host, `1` device: `config_id` becomes
+    `gemma3_1b_flax_sample_loop_l4_1h1d_presubmit`
+*   For `CPU_X86`, `1` host, `1` device: `config_id` becomes
+    `gemma3_1b_flax_sample_loop_x86_1h1d_presubmit`
+
+```yaml
+# xla/xla/tools/benchmarks/baseline/presubmit_baseline.yml
+{
+  # ... existing baselines ...
+
+  "gemma3_1b_flax_sample_loop_l4_1h1d_presubmit": {
+    "GPU_DEVICE_TIME": {
+      "baseline_ms": 4,  # Your measured baseline
+      "threshold": 0.30
+    },
+    "GPU_DEVICE_MEMCPY_TIME": {
+      "baseline_ms": 10, # Your measured baseline
+      "threshold": 0.30
+    }
+  },
+  "gemma3_1b_flax_sample_loop_x86_1h1d_presubmit": {
+    "CPU_TIME": {
+      "baseline_ms": 8000, # Your measured baseline
+      "threshold": 0.30
+    },
+    "WALL_TIME": {
+      "baseline_ms": 1300, # Your measured baseline
+      "threshold": 0.30
+    }
+  }
+}
+```
+
+### Step 4: Create a Pull Request
+
+1.  Commit your changes:
+    *   The HLO file (if added to the repo).
+    *   The updated benchmark registry file (e.g., `default_registry.yml`).
+    *   The updated `presubmit_baseline.yml` (if applicable).
+2.  Push your branch and open a Pull Request against the `openxla/xla` main branch.
+3.  A member of the OpenXLA repository or organization will need to review your
+    PR for safety before the CI system is invoked.
+     * **Note**: This step happens
+    automatically for organization members and most Googlers, but require
+    manual review for external contributors.
+4.  Once approved, the CI system will pick up your new benchmark configuration.
+    *   If it's a `PRESUBMIT` benchmark, it will run against your PR and check for regressions based on the baseline you provided.
+    *   If it's `POSTSUBMIT` or `SCHEDULED`, it will run after your PR is merged.
+5.  Monitor the CI checks. If the presubmit check fails due to your new
+    benchmark (e.g., performance is significantly different from your initial
+    baseline), you might need to adjust the baseline values in
+    `presubmit_baseline.yml` and update your PR.
+
+## Best Practices
+
+*   **Establish Baselines First:** Since, a baseline value per metric is required, always add the benchmark with only
+    `POSTSUBMIT` or `SCHEDULED` frequency first to establish stable baseline
+    values. Once it runs a few times and you have stable performance data, you
+    can add `PRESUBMIT` and the corresponding baseline entry in a follow-up PR.
+*   **Meaningful Names and Descriptions:** Make it easy for others to
+    understand what your benchmark does.
+*   **Targeted Metrics:** Only include relevant metrics in `target_metrics`.
+*   **Noise Reduction:** Use `runtime_flags: ["--num_repeats=X"]` (e.g., X=5 or
+    10) to run the benchmark multiple times within a single execution, which
+    helps in getting more stable measurements. The runner typically reports the
+    median or average.
+*   **Keep Baselines Updated:** If your benchmark's performance characteristics
+    change significantly (due to XLA improvements or changes in the benchmark
+    itself), the baseline values in `presubmit_baseline.yml` will need to be
+    updated. This is usually done by the benchmark owner or XLA maintainers.
+
+## Troubleshooting
+
+*   **Workflow Failures:** Check the GitHub Actions logs for detailed error
+    messages. The logs for the "Compare Benchmarks" step are particularly useful
+    for presubmit issues.
+*   **Incorrect `config_id`:** If your presubmit benchmark isn't being picked up
+    or matched to a baseline, double-check the `config_id` format in
+    `presubmit_baseline.yml`.
+*   **Performance Fluctuations:** Microbenchmarks can be sensitive to noise.
+    Ensure you're using `--num_repeats` and that your baseline reflects typical
+    performance.
+
+If you encounter issues, feel free to ask for help on the OpenXLA communication
+channels or tag the juliagmt-google@ on your PR.

third_party/xla/.github/workflows/benchmarks/build_binaries.sh

@@ -0,0 +1,135 @@
+# Copyright 2025 The OpenXLA Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+# .github/workflows/benchmarks/build_binaries.sh
+# TODO(juliagmt): convert this to a python script.
+# Builds HLO runner and stats computation binaries using build.py.
+# Expects HARDWARE_CATEGORY (matching enum values like CPU_X86, GPU_L4, etc.)
+# and OUTPUT_DIR environment variables to be set.
+# Outputs: runner_binary, stats_binary, device_type_flag to GITHUB_OUTPUT
+#!/bin/bash
+set -e # Exit immediately if a command exits with a non-zero status.
+set -u # Treat unset variables as an error when substituting.
+# set -o pipefail # Causes pipelines to fail if any command fails (see Run script)
+
+echo "--- Configuring and Building Binaries ---"
+echo "Workspace: $(pwd)"
+echo "Hardware Category: $HARDWARE_CATEGORY"
+echo "Output Directory (for profile.json.gz): $OUTPUT_DIR"
+
+# Sanitize HARDWARE_CATEGORY for use in filenames.
+# Replaces non-alphanumeric with underscore, converts to lowercase.
+HW_CATEGORY_SLUG=$(echo "${HARDWARE_CATEGORY:-UNSPECIFIED}" | tr '[:upper:]' '[:lower:]' | sed 's/[^A-Z0-9_]/_/g')
+
+# --- 1. Configure Backend (using ./configure.py if available) ---
+# This part can remain, as build.py doesn't currently handle `configure.py`
+configure_backend() {
+  echo "Configuring backend using ./configure.py if present..."
+  if [ ! -f "./configure.py" ]; then
+    echo "INFO: ./configure.py not found. Skipping configuration step."
+    return
+  fi
+
+  local hw_category_upper_for_configure
+  hw_category_upper_for_configure=$(echo "${HARDWARE_CATEGORY:-UNSPECIFIED}" | tr '[:lower:]' '[:upper:]')
+
+  case "$hw_category_upper_for_configure" in
+    CPU_X86 | CPU_ARM64)
+      echo "Running: ./configure.py --backend=CPU"
+      ./configure.py --backend=CPU || echo "INFO: CPU Configure script failed or is not applicable."
+      ;;
+    GPU_L4 | GPU_B200)
+      echo "Running: ./configure.py --backend=CUDA --cuda_compiler=nvcc"
+      ./configure.py --backend=CUDA --cuda_compiler=nvcc || echo "INFO: GPU Configure script failed or is not applicable."
+      ;;
+    *)
+      echo "INFO: Unknown hardware category '$hw_category_upper_for_configure'"
+      ;;
+  esac
+  echo "Configuration step finished."
+}
+
+# --- 2. Main Build Logic using build.py ---
+declare BAZEL_BIN_DIR="bazel-bin"
+declare runner_binary_path=""
+declare stats_binary_path=""
+declare device_type_flag_value=""
+
+configure_backend
+
+echo "Building with build.py for HARDWARE_CATEGORY: $HARDWARE_CATEGORY"
+
+BUILD_TYPE=""
+case "$HARDWARE_CATEGORY" in
+  CPU_X86)
+    BUILD_TYPE="XLA_LINUX_X86_CPU_128_VCPU_PRESUBMIT_GITHUB_ACTIONS"
+    runner_binary_path="./$BAZEL_BIN_DIR/xla/tools/multihost_hlo_runner/hlo_runner_main"
+    stats_binary_path="./$BAZEL_BIN_DIR/xla/tools/compute_xspace_stats_main"
+    device_type_flag_value="host"
+    ;;
+  CPU_ARM64)
+    BUILD_TYPE="XLA_LINUX_ARM64_CPU_48_VCPU_PRESUBMIT_GITHUB_ACTIONS"
+    runner_binary_path="./$BAZEL_BIN_DIR/xla/tools/multihost_hlo_runner/hlo_runner_main"
+    stats_binary_path="./$BAZEL_BIN_DIR/xla/tools/compute_xspace_stats_main"
+    device_type_flag_value="host"
+    ;;
+  GPU_L4)
+    BUILD_TYPE="XLA_LINUX_X86_GPU_L4_16_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS" # Or _48_VCPU if that's the more common
+    runner_binary_path="./$BAZEL_BIN_DIR/xla/tools/multihost_hlo_runner/hlo_runner_main_gpu"
+    stats_binary_path="./$BAZEL_BIN_DIR/xla/tools/compute_xspace_stats_main_gpu"
+    device_type_flag_value="gpu"
+    ;;
+  GPU_B200)
+    BUILD_TYPE="XLA_LINUX_X86_GPU_A4_224_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS"
+    runner_binary_path="./$BAZEL_BIN_DIR/xla/tools/multihost_hlo_runner/hlo_runner_main_gpu"
+    stats_binary_path="./$BAZEL_BIN_DIR/xla/tools/compute_xspace_stats_main_gpu"
+    device_type_flag_value="gpu"
+    ;;
+  *)
+    echo "::error::Unsupported HARDWARE_CATEGORY: '$HARDWARE_CATEGORY'. This script is configured to handle specific values from the HardwareCategory enum (CPU_X86, CPU_ARM64, GPU_L4, GPU_B200)."
+    exit 1
+    ;;
+esac
+
+echo "Executing build with build.py for build type: $BUILD_TYPE"
+# Run build.py with the determined build type
+python3 build_tools/ci/build.py --build="$BUILD_TYPE" || {
+  echo "::error::build.py failed for $BUILD_TYPE!"
+  exit 1
+}
+
+
+# --- 3. Verify Binaries and Set Outputs ---
+echo "Verifying binary existence..."
+if [ -z "$runner_binary_path" ] || [ ! -f "$runner_binary_path" ]; then
+  echo "::error::Runner binary path not set or binary '$runner_binary_path' not found after build for $HARDWARE_CATEGORY!"
+  exit 1
+fi
+if [ -z "$stats_binary_path" ] || [ ! -f "$stats_binary_path" ]; then
+  echo "::error::Stats binary path not set or binary '$stats_binary_path' not found after build for $HARDWARE_CATEGORY!"
+  exit 1
+fi
+echo "Binaries verified: $runner_binary_path, $stats_binary_path"
+
+echo "Setting step outputs..."
+{
+  echo "runner_binary=$runner_binary_path"
+  echo "stats_binary=$stats_binary_path"
+  echo "device_type_flag=$device_type_flag_value"
+} >> "$GITHUB_OUTPUT"
+
+echo "  Runner binary: $runner_binary_path"
+echo "  Stats binary: $stats_binary_path"
+echo "  Device type flag: $device_type_flag_value"
+echo "--- Build Script Finished ---"

third_party/xla/.github/workflows/benchmarks/compare_with_baseline.py

@@ -0,0 +1,340 @@
+# Copyright 2025 The OpenXLA Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""Compares benchmark results to predefined baselines and checks for regressions."""
+import argparse
+import json
+import os
+import sys
+import yaml
+
+
+def load_results_data(results_json_file):
+  """Loads and parses the results JSON file."""
+  try:
+    with open(results_json_file, "r") as f:
+      results_data = json.load(f)
+  except json.JSONDecodeError as e:
+    print(
+        f"::error file={results_json_file}::Failed to parse results JSON: {e}"
+    )
+    sys.exit(1)
+  except OSError as e:  # Catch other potential file reading errors
+    print(f"::error file={results_json_file}::Error reading results JSON: {e}")
+    sys.exit(1)
+  return results_data
+
+
+def load_baseline_data(baseline_yaml_file):
+  """Loads and parses the baseline YAML file."""
+  try:
+    with open(baseline_yaml_file, "r") as f:
+      baseline_data_full = yaml.safe_load(f)
+  except yaml.YAMLError as e:
+    print(
+        f"::error file={baseline_yaml_file}::Failed to parse baseline YAML: {e}"
+    )
+    sys.exit(1)
+  except OSError as e:  # Catch other potential file reading errors
+    print(
+        f"::error file={baseline_yaml_file}::Error reading baseline YAML: {e}"
+    )
+    sys.exit(1)
+  return baseline_data_full
+
+
+def validate_baseline_data(baseline_data_full, config_id, baseline_yaml_file):
+  """Validates the loaded baseline data and extracts config-specific baselines."""
+  if not baseline_data_full or not isinstance(baseline_data_full, dict):
+    print(
+        f"::warning file={baseline_yaml_file}::Baseline YAML is empty or"
+        " not a dictionary. Skipping comparison."
+    )
+    sys.exit(0)
+
+  if config_id not in baseline_data_full:
+    print(
+        f"::notice::No baseline found for config_id '{config_id}' in"
+        f" {baseline_yaml_file}. Skipping comparison."
+    )
+    sys.exit(0)
+
+  config_baselines = baseline_data_full[config_id]
+  if not isinstance(config_baselines, dict):
+    print(
+        f"::warning file={baseline_yaml_file},title=Invalid Baseline"
+        f" Structure::Baseline entry for '{config_id}' is not a"
+        " dictionary. Skipping."
+    )
+    sys.exit(0)
+  return config_baselines
+
+
+def compare_metrics(
+    results_data, config_baselines, results_json_file, config_id
+):
+  """Compares metrics from results to baselines and reports findings."""
+  regressions_found = False
+  summary_messages = ["\n--- Comparison Summary ---"]
+
+  # --- IMPORTANT: Metric Extraction Logic ---
+  # This assumes your results.json (produced by compute_xspace_stats_main or
+  # fallback)
+  # has a structure like:
+  # {
+  #   ...,
+  #   "metrics": {
+  #     "GPU_DEVICE_TIME": { "value": 150.0, "unit": "ms", ... },
+  #     "GPU_DEVICE_MEMCPY_TIME": { "value": 1.2, "unit": "ms", ... },
+  #     "PEAK_GPU_MEMORY": { "value": 12.45, "unit": "GB", ... },
+  #     ...
+  #   }
+  # }
+  # If your actual results.json structure is different, you MUST adapt this
+  # section.
+  actual_metrics_container = results_data.get("metrics")
+  metric_to_baseline_key = {
+      "WALL_TIME": "baseline_ms",
+      "GPU_DEVICE_TIME": "baseline_ms",
+      "GPU_DEVICE_MEMCPY_TIME": "baseline_ms",
+      "CPU_TIME": "baseline_ms",
+      "PEAK_CPU_MEMORY": "baseline_gb",
+      "PEAK_GPU_MEMORY": "baseline_gb",
+  }
+
+  if not actual_metrics_container or not isinstance(
+      actual_metrics_container, dict
+  ):
+    summary_messages.append(
+        "::warning title=Missing Metrics in Results::'metrics' key not found"
+        f" or not a dictionary in '{results_json_file}'. Cannot perform"
+        " comparison."
+    )
+    # Depending on your policy, if metrics are always expected,
+    # this could be sys.exit(1)
+    # For now, it will skip comparisons and pass if no metrics are found.
+    print("\n".join(summary_messages))
+    sys.exit(0)  # Exit cleanly if no metrics to compare
+
+  for metric_name, baseline_info in config_baselines.items():
+    if (
+        not isinstance(baseline_info, dict)
+        or "threshold" not in baseline_info
+        or not {"baseline_ms", "baseline_gb"}.intersection(baseline_info.keys())
+    ):
+      summary_messages.append(
+          f"::warning title=Malformed Baseline::Metric '{metric_name}' in"
+          f" baseline for '{config_id}' is missing 'threshold', or is missing"
+          " both 'baseline_ms' and 'baseline_gb', or is not structured as a"
+          " dictionary. Skipping."
+      )
+      continue
+
+    baseline_key = metric_to_baseline_key.get(metric_name)
+    if not baseline_key:
+      summary_messages.append(
+          f"::warning title=Unsupported Metric::Metric '{metric_name}' is not"
+          " supported by this script. Skipping."
+      )
+      continue
+
+    try:
+      baseline_value = float(baseline_info[baseline_key])
+      threshold_percentage = float(baseline_info["threshold"])
+    except ValueError:
+      summary_messages.append(
+          f"::warning title=Invalid Baseline Value::Metric '{metric_name}' in"
+          f" baseline for '{config_id}' has non-numeric '{baseline_key}' or"
+          " 'threshold'. Skipping."
+      )
+      continue
+
+    # Extract the actual metric value from results.json
+    actual_metric_entry = actual_metrics_container.get(metric_name)
+
+    if not actual_metric_entry or "value" not in actual_metric_entry:
+      summary_messages.append(
+          f"Metric '{metric_name}': Actual value or 'value' key not found in"
+          " results, or not a dictionary. Skipping."
+      )
+      # For debugging:
+      # available_keys = (
+      #     list(actual_metrics_container.keys())
+      #     if actual_metrics_container
+      #     else "None"
+      # )
+      # summary_messages.append(
+      #   f"  Available metric keys in results 'metrics' object:"
+      #   f" {available_keys}"
+      # )
+      continue
+
+    try:
+      actual_value = float(actual_metric_entry["value"])
+    except (ValueError, TypeError):
+      summary_messages.append(
+          f"Metric '{metric_name}': Actual value"
+          f" '{actual_metric_entry['value']}' is not a valid number."
+          " Skipping."
+      )
+      continue
+
+    actual_unit = actual_metric_entry.get("unit")
+    if not actual_unit:
+      summary_messages.append(
+          f"Metric '{metric_name}': Actual value unit is not specified."
+          " Skipping."
+      )
+      continue
+
+    summary_messages.append(f"\nComparing metric: {metric_name}")
+    summary_messages.append(f"  Actual Value: {actual_value:.3f} {actual_unit}")
+    summary_messages.append(
+        f"  Baseline Value: {baseline_value:.3f} {actual_unit}"
+    )
+    summary_messages.append(
+        f"  Allowed Threshold: {threshold_percentage*100:.1f}%"
+    )
+
+    # Higher value is worse for time-based metrics
+    allowed_upper_bound = baseline_value * (1.0 + threshold_percentage)
+    summary_messages.append(
+        "  Allowed Upper Bound (Baseline * (1 + Threshold)):"
+        f" {allowed_upper_bound:.3f} {actual_unit}"
+    )
+
+    if actual_value > allowed_upper_bound:
+      percentage_diff = 0.0
+      if (
+          abs(baseline_value) > 1e-9
+      ):  # Avoid division by zero for very small baselines
+        percentage_diff = (
+            (actual_value - baseline_value) / baseline_value
+        ) * 100.0
+      elif (
+          actual_value > 0
+      ):  # If baseline is effectively zero, any positive value is infinitely
+        # worse.
+        percentage_diff = float("inf")
+
+      # Use GitHub Actions error annotation for better visibility
+      error_title = f"REGRESSION: {metric_name}"
+      error_details = (
+          f"Value {actual_value:.3f} {actual_unit} is {percentage_diff:.2f}%"
+          f" worse than baseline {baseline_value:.3f} {actual_unit}. Exceeds"
+          f" threshold of {threshold_percentage*100:.1f}% (max allowed:"
+          f" {allowed_upper_bound:.3f} {actual_unit})."
+      )
+      summary_messages.append(
+          "  ::error"
+          f" file={results_json_file},title={error_title}::{error_details}"
+      )
+      regressions_found = True
+    else:
+      summary_messages.append(
+          f"  Metric '{metric_name}' is within threshold. PASSED."
+      )
+  return regressions_found, summary_messages
+
+
+def main():
+  parser = argparse.ArgumentParser(
+      description=(
+          "Compare benchmark results to baselines and fail on regression."
+      ),
+      formatter_class=argparse.RawTextHelpFormatter,
+  )  # For better help text formatting
+  parser.add_argument(
+      "--results-json-file",
+      required=True,
+      help=(
+          "Path to the benchmark results JSON file (e.g., output/results.json)."
+      ),
+  )
+  parser.add_argument(
+      "--baseline-yaml-file",
+      required=True,
+      help=(
+          "Path to the baseline YAML file (e.g.,"
+          " xla/tools/benchmarks/baseline/presubmit_baseline.yml)."
+      ),
+  )
+  parser.add_argument(
+      "--config-id",
+      required=True,
+      help=(
+          "The configuration ID for the current benchmark run. \nThis ID must"
+          " exactly match a top-level key in the baseline YAML file. \nExample:"
+          " 'gemma3_1b_flax_call_gpu_b200_1_host_1_device'"
+      ),
+  )
+
+  args = parser.parse_args()
+
+  print("--- Benchmark Baseline Comparison ---")
+  print(f"Results JSON: {args.results_json_file}")
+  print(f"Baseline YAML: {args.baseline_yaml_file}")
+  print(f"Config ID for Baseline Lookup: {args.config_id}")
+
+  if not os.path.exists(args.results_json_file):
+    print(
+        f"::error file={args.results_json_file}::Results JSON file not found."
+    )
+    sys.exit(1)
+
+  if not os.path.exists(args.baseline_yaml_file):
+    print(
+        f"::notice file={args.baseline_yaml_file}::Baseline YAML file not"
+        " found. Skipping comparison."
+    )
+    # Exiting 0 because no baseline means no regression to detect.
+    # If a baseline is mandatory, this could be sys.exit(1).
+    sys.exit(0)
+
+  results_data = load_results_data(args.results_json_file)
+  print("\nLoaded Results Data:")
+  print(json.dumps(results_data, indent=2))
+
+  baseline_data_full = load_baseline_data(args.baseline_yaml_file)
+  config_baselines = validate_baseline_data(
+      baseline_data_full, args.config_id, args.baseline_yaml_file
+  )
+  print("\nLoaded Baseline Data for this config_id:")
+  print(json.dumps(config_baselines, indent=2))
+
+  regressions_found, summary_messages = compare_metrics(
+      results_data,
+      config_baselines,
+      args.results_json_file,
+      args.config_id,
+  )
+
+  print("\n".join(summary_messages))
+
+  if regressions_found:
+    print(
+        "\n::error::One or more benchmark metrics regressed beyond the allowed"
+        " threshold. Failing the check."
+    )
+    sys.exit(1)  # Exit with error code 1 to fail the GitHub Actions step
+  else:
+    print(
+        "\nAll benchmark metrics are within allowed thresholds (or no"
+        " applicable baselines found)."
+    )
+    sys.exit(0)  # Exit with success
+
+
+if __name__ == "__main__":
+  main()

third_party/xla/.github/workflows/benchmarks/prepare_artifact.sh

@@ -0,0 +1,86 @@
+# Copyright 2025 The OpenXLA Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+# .github/workflows/benchmarks/prepare_artifact.sh
+# TODO(juliagmt): convert this to a python script.
+#!/bin/bash
+set -e # Exit immediately if a command exits with a non-zero status.
+set -u # Treat unset variables as an error when substituting.
+
+echo "--- prepare_artifact.sh (Self-creating directory version) ---"
+echo "SCRIPT: Current PWD: $(pwd)"
+echo "SCRIPT: GITHUB_WORKSPACE is: $GITHUB_WORKSPACE"
+echo "SCRIPT: Intended OUTPUT_DIR is: $OUTPUT_DIR"
+
+# Create the directory HERE, inside this script, right before using it.
+echo "SCRIPT: Ensuring directory '$OUTPUT_DIR' exists by creating it with mkdir -p."
+mkdir -p "$OUTPUT_DIR"
+
+# Verify creation immediately
+echo "SCRIPT: Verifying directory '$OUTPUT_DIR' after mkdir with 'ls -ld':"
+ls -ld "$OUTPUT_DIR" || echo "SCRIPT: 'ls -ld ""$OUTPUT_DIR""' FAILED even after mkdir in script!"
+
+# Now, check with [ -d ... ]
+if [ ! -d "$OUTPUT_DIR" ]; then
+  echo "::error::SCRIPT: Output directory '$OUTPUT_DIR' STILL NOT found with [ -d ... ] even after mkdir in this script."
+  echo "SCRIPT: Listing parent directory '$(dirname "$OUTPUT_DIR")' using 'ls -la':"
+  ls -la "$(dirname "$OUTPUT_DIR")" || echo "SCRIPT: Failed to list parent directory."
+  exit 1
+else
+  echo "SCRIPT: Output directory '$OUTPUT_DIR' IS now found with [ -d ... ]."
+fi
+
+# --- Original script logic from here ---
+echo "--- Preparing Artifact (main logic) ---"
+
+ARTIFACT_FILE_NAME=$(basename "$ARTIFACT_LOCATION")
+LOCAL_ARTIFACT_PATH="$OUTPUT_DIR/$ARTIFACT_FILE_NAME"
+
+echo "Target local path: ${LOCAL_ARTIFACT_PATH}"
+
+if [ "$IS_GCS_ARTIFACT" == "true" ]; then
+  echo "Downloading GCS artifact from: $ARTIFACT_LOCATION"
+   if ! command -v wget &> /dev/null; then
+     echo "::error::wget command not found in container. Cannot download GCS artifact."
+     exit 1
+   fi
+
+   wget -q -nv -O "$LOCAL_ARTIFACT_PATH" "$ARTIFACT_LOCATION"
+   WGET_EXIT_CODE=$?
+    if [ $WGET_EXIT_CODE -ne 0 ]; then
+      echo "::error::wget failed to download GCS artifact from $ARTIFACT_LOCATION (Exit code: $WGET_EXIT_CODE)"
+      rm -f "$LOCAL_ARTIFACT_PATH" # Clean up partial file
+      exit $WGET_EXIT_CODE
+    fi
+   echo "GCS artifact downloaded."
+else
+   REPO_ARTIFACT_PATH="$GITHUB_WORKSPACE/$ARTIFACT_LOCATION" # ARTIFACT_LOCATION is the relative repo path here
+   echo "Copying local artifact from workspace path: $REPO_ARTIFACT_PATH (IS_GCS_ARTIFACT was false)"
+    if [ ! -f "$REPO_ARTIFACT_PATH" ]; then
+       echo "::error::Local artifact not found at repository path: $REPO_ARTIFACT_PATH"
+       exit 1
+    fi
+    cp -v "$REPO_ARTIFACT_PATH" "$LOCAL_ARTIFACT_PATH" || exit 1 # Exit if copy fails
+   echo "Local artifact copied successfully."
+fi
+
+# Verify the final destination file exists
+if [ ! -f "$LOCAL_ARTIFACT_PATH" ]; then
+   echo "::error::Final artifact file not found at destination: $LOCAL_ARTIFACT_PATH"
+   exit 1
+fi
+echo "Artifact successfully prepared at $LOCAL_ARTIFACT_PATH."
+
+echo "artifact_local_path=$LOCAL_ARTIFACT_PATH" >> "$GITHUB_OUTPUT"
+echo "--- Artifact Prep Finished ---"

third_party/xla/.github/workflows/benchmarks/run_benchmark.sh

@@ -0,0 +1,276 @@
+# Copyright 2025 The OpenXLA Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+# .github/workflows/benchmarks/run_benchmark.sh
+# TODO(juliagmt): convert this to a python script.
+#!/bin/bash
+
+# This script encapsulates the logic to run a benchmark and generate results.json.
+
+set -u # Treat unset variables as an error when substituting.
+# IMPORTANT: pipefail is handled specifically around the runner command.
+set -e # Exit on errors, EXCEPT where explicitly handled.
+
+echo "--- Running Benchmark Script ---"
+echo "Benchmark Name: $BENCHMARK_NAME"
+echo "Config ID: $CONFIG_ID"
+echo "Hardware Category: $HARDWARE_CATEGORY"
+echo "Output Directory: $OUTPUT_DIR"
+echo "Runner Binary: $RUNNER_BINARY"
+echo "Stats Binary: $STATS_BINARY"
+echo "Device Type Flag: $DEVICE_TYPE_FLAG"
+echo "Local Artifact Path: $LOCAL_ARTIFACT_PATH"
+echo "Input Format: $INPUT_FORMAT"
+echo "XLA Flags JSON: $XLA_FLAGS_JSON"
+echo "Runtime Flags JSON: $RUNTIME_FLAGS_JSON"
+echo "Commit SHA: $COMMIT_SHA"
+echo "Workflow Run ID: $WORKFLOW_RUN_ID"
+
+
+# --- Validate Inputs ---
+if [ -z "$LOCAL_ARTIFACT_PATH" ] || [ ! -f "$LOCAL_ARTIFACT_PATH" ]; then echo "::error::LOCAL_ARTIFACT_PATH path is invalid or file not found: '$LOCAL_ARTIFACT_PATH'"; exit 1; fi
+if [ -z "$RUNNER_BINARY" ] || [ ! -x "$RUNNER_BINARY" ]; then echo "::error::RUNNER_BINARY path is invalid or file not executable: '$RUNNER_BINARY'"; exit 1; fi
+if [ -z "$DEVICE_TYPE_FLAG" ]; then echo "::error::DEVICE_TYPE_FLAG is empty"; exit 1; fi
+if [ -z "$STATS_BINARY" ] || [ ! -x "$STATS_BINARY" ]; then echo "::error::STATS_BINARY path is invalid or file not executable: '$STATS_BINARY'"; exit 1; fi
+if ! command -v jq &> /dev/null; then echo "::error::jq command not found."; exit 1; fi
+
+RUNNER_STDOUT_FILE="$OUTPUT_DIR/runner_stdout.txt"
+XSPACE_FILE_PATH="$OUTPUT_DIR/xspace.pb"
+RESULTS_JSON_FILE="$OUTPUT_DIR/results.json"
+
+# --- Prepare flags ---
+declare -a xla_flags_array=()
+declare -a runtime_flags_array=()
+
+# Use JQ to safely parse JSON and populate bash arrays
+if echo "$XLA_FLAGS_JSON" | jq -e '. | arrays and length > 0' > /dev/null; then
+    mapfile -t xla_flags_array < <(echo "$XLA_FLAGS_JSON" | jq -r '.[]')
+fi
+if echo "$RUNTIME_FLAGS_JSON" | jq -e '. | arrays and length > 0' > /dev/null; then
+   mapfile -t runtime_flags_array < <(echo "$RUNTIME_FLAGS_JSON" | jq -r '.[]')
+fi
+
+needs_xspace_dump_flag=true # By default, we enable profiler and xspace dump.
+
+# --- Build Runner Command ---
+declare -a runner_command_array=(
+    "$RUNNER_BINARY"
+    "--device_type=$DEVICE_TYPE_FLAG"
+)
+if [ ${#runtime_flags_array[@]} -gt 0 ]; then runner_command_array+=("${runtime_flags_array[@]}"); fi
+if [ ${#xla_flags_array[@]} -gt 0 ]; then runner_command_array+=("${xla_flags_array[@]}"); fi
+if $needs_xspace_dump_flag; then
+   runner_command_array+=("--xla_gpu_dump_xspace_to=$XSPACE_FILE_PATH")
+fi
+runner_command_array+=("$LOCAL_ARTIFACT_PATH")
+
+# --- Execute Runner ---
+echo "Executing HLO Runner command:"
+printf "%q " "${runner_command_array[@]}"; echo
+
+set +e # Disable exit-on-error temporarily to capture exit code
+set -o pipefail # Ensure tee doesn't mask the runner's exit code
+"${runner_command_array[@]}" 2>&1 | tee "$RUNNER_STDOUT_FILE"
+RUNNER_EXIT_CODE=${PIPESTATUS[0]}
+set +o pipefail
+set -e # Re-enable exit-on-error
+
+echo "Runner stdout/stderr saved to $RUNNER_STDOUT_FILE"
+echo "Runner exited with code: $RUNNER_EXIT_CODE"
+
+
+# --- Process Stats and Generate results.json ---
+STATS_RUN_SUCCESSFUL=false
+METRICS_JSON_CONTENT="{}" # Initialize as an empty JSON object
+
+if [ -f "$XSPACE_FILE_PATH" ] && [ $RUNNER_EXIT_CODE -eq 0 ]; then
+  echo "XSpace file found. Running compute_xspace_stats_main..."
+  STATS_PLATFORM_TYPE=$([[ "$HARDWARE_CATEGORY" == GPU* ]] && echo "GPU" || echo "CPU")
+
+  # Capture the output of compute_xspace_stats_main to parse it
+  # Do not write its output directly to results.json yet
+  echo "Executing Stats command and capturing its output:"
+
+  set +e # Temporarily disable exit-on-error for stats command
+  STATS_OUTPUT=$("$STATS_BINARY" --input="$XSPACE_FILE_PATH" --device_type="$STATS_PLATFORM_TYPE" 2>&1)
+  STATS_EXIT_CODE=$?
+  set -e
+
+  echo "compute_xspace_stats_main output:"
+  echo "$STATS_OUTPUT"
+  echo "compute_xspace_stats_main exited with code: $STATS_EXIT_CODE"
+
+  # Append stats tool's raw output to the main runner log for complete record
+  echo -e "\n--- compute_xspace_stats_main Raw Output ---" >> "$RUNNER_STDOUT_FILE"
+  echo "$STATS_OUTPUT" >> "$RUNNER_STDOUT_FILE"
+  echo "--- End compute_xspace_stats_main Raw Output ---" >> "$RUNNER_STDOUT_FILE"
+
+
+  if [ $STATS_EXIT_CODE -eq 0 ]; then
+    STATS_RUN_SUCCESSFUL=true
+    metrics_obj_str="{"
+    first_metric=true
+
+    # Process each line of STATS_OUTPUT
+    while IFS=':' read -r key value; do
+        # Trim leading/trailing whitespace from key and value
+        key=$(echo "$key" | sed 's/^[[:space:]]*//;s/[[:space:]]*$//')
+        value=$(echo "$value" | sed 's/^[[:space:]]*//;s/[[:space:]]*$//')
+
+        # Sanitize base metric key (e.g., "Device Time" -> "DEVICE_TIME")
+        base_metric_key=$(echo "$key" | tr ' ' '_' | tr '[:lower:]' '[:upper:]')
+        final_metric_key=""
+
+        # Determine the final metric key based on HARDWARE_CATEGORY for baseline matching
+        if [[ "$HARDWARE_CATEGORY" == GPU* ]]; then
+            if [[ "$base_metric_key" == "DEVICE_TIME" ]]; then
+                final_metric_key="GPU_DEVICE_TIME"
+            elif [[ "$base_metric_key" == "DEVICE_MEMCPY_TIME" ]]; then
+                final_metric_key="GPU_DEVICE_MEMCPY_TIME"
+            elif  [[ "$base_metric_key" == "PEAK_MEMORY" ]]; then
+                final_metric_key="PEAK_GPU_MEMORY"
+            # Add other specific GPU mappings here if needed
+            # else
+            #    final_metric_key="GPU_${base_metric_key}" # Generic prefix
+            else
+                final_metric_key="$base_metric_key" # If no specific GPU mapping, use base
+            fi
+        elif [[ "$HARDWARE_CATEGORY" == CPU* ]]; then
+            if [[ "$base_metric_key" == "CPU_TIME" ]] || [[ "$base_metric_key" == "TIME" ]]; then # Handle "CPU Time" or just "Time"
+                final_metric_key="CPU_TIME"
+            elif [[ "$base_metric_key" == "WALL_TIME" ]]; then
+                final_metric_key="WALL_TIME" # Wall time is generic
+            # Add other specific CPU mappings here if needed
+            # else
+            #    final_metric_key="CPU_${base_metric_key}" # Generic prefix
+            else
+                final_metric_key="$base_metric_key" # If no specific CPU mapping, use base
+            fi
+        else
+            final_metric_key="$base_metric_key" # For unknown/other categories
+        fi
+
+        # Expecting lines like "Metric Name: 123.45 us" or "Metric Name: 12345 bytes"
+        read number unit <<< $(echo "$value" | sed -E 's/([0-9]+\.?[0-9]*)\s*([a-zA-Z]+).*/\1 \2/')
+        # Convert microseconds to milliseconds
+        if [[ "$unit" == "us" ]]; then
+            final_metric_value=$(LC_ALL=C awk -v num="$number" 'BEGIN { printf "%.3f", num / 1000 }')
+            final_unit="ms"
+        elif [[ "$unit" == "bytes" ]]; then
+            # Convert bytes to GB
+            final_metric_value=$(echo "$number" | awk '{printf "%.2f", $1/1024^3}')
+            final_unit="GB"
+        else
+          echo "::warning::Skipping unsupported unit: $unit"
+          continue
+        fi
+        if ! $first_metric; then metrics_obj_str+=","; fi
+        metrics_obj_str+="\"$final_metric_key\": {\"value\": $final_metric_value, \"unit\": \"$final_unit\"}"
+        first_metric=false
+        echo "INFO: Parsed metric: OriginalKey='$key', BaseKey='$base_metric_key', FinalKey='$final_metric_key', Value='$final_metric_value $final_unit'"
+    done <<< "$STATS_OUTPUT"
+    metrics_obj_str+="}"
+
+    if echo "$metrics_obj_str" | jq -e . > /dev/null 2>&1; then
+        METRICS_JSON_CONTENT=$(echo "$metrics_obj_str" | jq '.')
+        echo "Successfully parsed metrics from stats output."
+    else
+        echo "::warning::Could not construct valid JSON from stats output. Metrics object will be empty."
+        echo "Problematic metrics string constructed: $metrics_obj_str"
+        METRICS_JSON_CONTENT="{}"
+        STATS_RUN_SUCCESSFUL=false 
+    fi
+  else
+    echo "::warning::compute_xspace_stats_main failed with code $STATS_EXIT_CODE. No metrics will be parsed from its output."
+  fi
+else
+   if [ $RUNNER_EXIT_CODE -ne 0 ]; then 
+      echo "::warning::Runner failed (Exit Code: $RUNNER_EXIT_CODE), skipping stats processing."
+   else 
+     echo "::warning::XSpace file missing at $XSPACE_FILE_PATH, skipping stats processing."
+   fi
+fi
+
+# --- Construct Final results.json ---
+RUN_STATUS_MSG=""
+ERROR_MSG_CONTENT=""
+
+if [ $RUNNER_EXIT_CODE -ne 0 ]; then
+    RUN_STATUS_MSG="FAILURE"
+    ERROR_MSG_CONTENT="Runner failed with code $RUNNER_EXIT_CODE"
+elif [ ! -f "$XSPACE_FILE_PATH" ]; then
+    RUN_STATUS_MSG="SUCCESS_NO_PROFILE"
+    ERROR_MSG_CONTENT="XSpace file not generated by successful run."
+elif [ $STATS_EXIT_CODE -ne 0 ] || [ "$STATS_RUN_SUCCESSFUL" = false ] ; then
+    RUN_STATUS_MSG="STATS_FAILURE"
+    ERROR_MSG_CONTENT="compute_xspace_stats_main failed (code $STATS_EXIT_CODE) or metrics parsing failed. Runner was successful."
+else
+    RUN_STATUS_MSG="SUCCESS"
+    ERROR_MSG_CONTENT=""
+fi
+
+# Use jq to build the final JSON, incorporating the parsed metrics
+jq -n \
+  --arg bn "$BENCHMARK_NAME" \
+  --arg cid "$CONFIG_ID" \
+  --arg hc "$HARDWARE_CATEGORY" \
+  --arg rs "$RUN_STATUS_MSG" \
+  --arg em "$ERROR_MSG_CONTENT" \
+  --arg cs "$COMMIT_SHA" \
+  --arg wrid "$WORKFLOW_RUN_ID" \
+  --argjson metrics "$METRICS_JSON_CONTENT" \
+  '{
+     benchmark_name: $bn,
+     config_id: $cid,
+     hardware_category: $hc,
+     run_status: $rs,
+     error_message: $em,
+     commit_sha: $cs,
+     workflow_run_id: $wrid,
+     metrics: $metrics
+   }' > "$RESULTS_JSON_FILE"
+
+if [ $? -eq 0 ]; then
+    echo "Final results.json created at $RESULTS_JSON_FILE."
+else
+    echo "::error::FATAL: Failed to create final results.json using jq."
+    echo "FATAL JQ ERROR. Benchmark Name: $BENCHMARK_NAME, Run Status: $RUN_STATUS_MSG, Error: $ERROR_MSG_CONTENT" > "$RESULTS_JSON_FILE.txt"
+    exit 1 
+fi
+
+# --- Debug: Verify file creation ---
+echo "DEBUG: Listing contents of OUTPUT_DIR ($OUTPUT_DIR):"
+ls -la "$OUTPUT_DIR"
+echo "DEBUG: Checking for RESULTS_JSON_FILE ($RESULTS_JSON_FILE):"
+if [ -f "$RESULTS_JSON_FILE" ]; then
+  echo "DEBUG: RESULTS_JSON_FILE exists. Content (first 20 lines):"
+  head -n 20 "$RESULTS_JSON_FILE"
+else
+  echo "DEBUG: RESULTS_JSON_FILE does NOT exist."
+  if [ -f "${RESULTS_JSON_FILE}.txt" ]; then
+    echo "DEBUG: RESULTS_JSON_FILE.txt exists. Content:"
+    cat "${RESULTS_JSON_FILE}.txt"
+  else
+    echo "DEBUG: RESULTS_JSON_FILE.txt also does NOT exist."
+  fi
+fi
+echo "DEBUG: End of file check."
+
+# --- Final Exit Status of the script ---
+if [ $RUNNER_EXIT_CODE -ne 0 ]; then
+  echo "::error::Benchmark run failed (HLO Runner Exit Code: $RUNNER_EXIT_CODE)."
+  exit $RUNNER_EXIT_CODE
+fi
+
+echo "--- Run Benchmark Script Finished ---"

third_party/xla/.github/workflows/benchmarks/run_comparison.sh

@@ -0,0 +1,59 @@
+# Copyright 2025 The OpenXLA Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+# .github/workflows/benchmarks/run_comparison.sh
+# TODO(juliagmt): convert this to a python script.
+#!/bin/bash
+
+# This script encapsulates the logic to compare benchmark results against a baseline.
+
+set -e  # Exit immediately if a command exits with a non-zero status.
+set -u # Treat unset variables as an error when substituting.
+
+echo "--- Running Comparison Script ---"
+echo "CONFIG_ID for baseline lookup: $CONFIG_ID" # This is now the comprehensive ID
+echo "Results Directory: $RESOLVED_OUTPUT_DIR"
+echo "Baseline File: $RESOLVED_BASELINE_YAML"
+echo "Comparison Python Script: $RESOLVED_COMPARISON_SCRIPT"
+
+ACTUAL_RESULTS_JSON_PATH="${RESOLVED_OUTPUT_DIR}/results.json"
+
+if [ ! -f "$ACTUAL_RESULTS_JSON_PATH" ]; then
+  echo "::warning::Primary results file '$ACTUAL_RESULTS_JSON_PATH' not found. Cannot perform baseline comparison."
+  # Check for fallback .txt file for more info if primary is missing
+  if [ -f "${ACTUAL_RESULTS_JSON_PATH}.txt" ]; then
+    echo "Fallback results.json.txt found:"
+    cat "${ACTUAL_RESULTS_JSON_PATH}.txt"
+  fi
+  exit 0 # Exiting cleanly to not block if results are missing; comparison script handles this too
+fi
+SCRIPT_CONFIG_ID_FOR_BASELINE_LOOKUP="${CONFIG_ID}"
+
+echo "Using Config ID for baseline lookup: $SCRIPT_CONFIG_ID_FOR_BASELINE_LOOKUP"
+
+echo "Constructed Config ID for baseline lookup: $SCRIPT_CONFIG_ID_FOR_BASELINE_LOOKUP"
+
+python3 "$RESOLVED_COMPARISON_SCRIPT" \
+  --results-json-file="$ACTUAL_RESULTS_JSON_PATH" \
+  --baseline-yaml-file="$RESOLVED_BASELINE_YAML" \
+  --config-id="$SCRIPT_CONFIG_ID_FOR_BASELINE_LOOKUP"
+
+COMPARISON_EXIT_CODE=$?
+if [ $COMPARISON_EXIT_CODE -ne 0 ]; then
+  echo "::error::Baseline comparison script failed or regressions detected (Exit Code: $COMPARISON_EXIT_CODE)."
+  exit $COMPARISON_EXIT_CODE
+fi
+
+echo "Baseline comparison successful: No regressions detected or no applicable baselines found."
+echo "--- Comparison Script Finished ---"