wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
third_party/xla/.github/actionlint.yml
Vendored
+30
View File
@@ -0,0 +1,30 @@
# Copyright 2026 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.
# ============================================================================
self-hosted-runner:
labels:
- "linux-x86-n2-16" # Linux X86 runner using the 16 vcpu n2-standard-16 machine.
- "linux-x86-n2-64" # Linux X86 runner using the 64 vcpu n2-standard-64 machine.
- "linux-x86-n2-128" # Linux X86 runner using the 128 vcpu n2-standard-128 machine.
- "linux-x86-g2-16-l4-1gpu" # Linux X86 GPU runner using g2-standard-16 machine with 1 NVIDIA L4 GPU attached.
- "linux-x86-g2-48-l4-4gpu" # Linux X86 GPU runner using g2-standard-48 machine with 4 NVIDIA L4 GPUs attached.
- "linux-arm64-c4a-16" # Linux ARM64 CPU Runner using the 16 vcpu c4a-standard-16 machine.
- "linux-arm64-t2a-48" # Linux ARM64 CPU Runner using the 48 vcpu t2a-standard-48 machine.
- "windows-x86-n2-16" # Windows X86 runner using n2-standard-16 machine.
- "amd-do-linux-xla-gpu-gfx950-1" # AMD runner 1 GPU.
- "amd-do-linux-xla-gpu-gfx950-4" # AMD runner 4 GPU.
- "amd-do-linux-xla-gpu-gfx950-8" # AMD runner 8 GPUs.
- "linux-x86-a4-224-b200-1gpu" # Linux X86 GPU runner using 1 B200 GPU and 1/8 the resources of a a4-highgpu-8g machine
- "linux-x86-a3-8g-h100-8gpu" # Linux X86 GPU runner using a3-highgpu-8g machine with 8 NVIDIA H100 GPUs attached.
- "linux-x64-battlemage-256-1gpu-intel" # Linux X86 GPU runner using an EMR4148 machine with 2 Intel BMG GPUs attached.
third_party/xla/.github/dependabot.yml
Vendored
+33
View File
@@ -0,0 +1,33 @@
# 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.
# ============================================================================
# To get started with Dependabot version updates, you'll need to specify which
# package ecosystems to update and where the package manifests are located.
# Please see the documentation for all configuration options:
# https://help.github.com/github/administering-a-repository/configuration-options-for-dependency-updates
version: 2
updates:
- package-ecosystem: pip
directory: /
schedule:
interval: weekly
cooldown:
default-days: 7
- package-ecosystem: github-actions
directory: /
schedule:
interval: weekly
cooldown:
default-days: 7
third_party/xla/.github/pull_request_template.md
Vendored
+24
View File
@@ -0,0 +1,24 @@
📝 Summary of Changes
Please provide a clear and concise summary of the changes you've made.
🎯 Justification
Explain why this change is important and which workload benefits from this
change.
🚀 Kind of Contribution
Please remove what does not apply: 🐛 Bug Fix, ⚡️ Performance Improvement,
✨ New Feature, ♻️ Cleanup, 📚 Documentation, 🧪 Tests
📊 Benchmark (for Performance Improvements)
Please measure and include speedups for one of the public HLOs in
`compiler/xla/tools/benchmarks/hlo/`.
🧪 Unit Tests:
What unit tests were added? For example, a new pass should be tested on minimal
HLO. The transformation can be tested with FileCheck tests or assertions on the
transformed HLO.
🧪 Execution Tests:
What execution tests were added? For example, a new optimization should be
tested with an end-to-end execution test triggering the optimization and
asserting correctness. Please provide test cases running with at most 2 GPUs.
third_party/xla/.github/workflows/actions-lint.yml
Vendored
+39
View File
@@ -0,0 +1,39 @@
# Copyright 2026 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.
# ============================================================================
name: actionlint
on:
pull_request:
push:
branches:
- main
workflow_dispatch:
permissions: {}
jobs:
actionlint:
runs-on: ubuntu-latest
steps:
- name: Checkout
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
with:
persist-credentials: false
- name: actionlint
uses: docker://rhysd/actionlint@sha256:b1934ee5f1c509618f2508e6eb47ee0d3520686341fec936f3b79331f9315667 # v1.7.12
with:
args: -color
third_party/xla/.github/workflows/bazel_dependency_violations.yml
Vendored
+71
View File
@@ -0,0 +1,71 @@
# Copyright 2024 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.
# ============================================================================
name: Bazel Dependency Violations
permissions:
contents: read
on:
pull_request:
push:
branches:
- main
jobs:
dependency-violations:
strategy:
matrix:
tag: [gpu, cuda-only, rocm-only]
name: no-${{ matrix.tag }}-targets-in-cpu-build
runs-on: "linux-x86-n2-16"
container:
image: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest" # zizmor: ignore[unpinned-images]
defaults:
run:
shell: bash
timeout-minutes: 10
continue-on-error: true
steps:
- name: "Checking out repository"
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
with:
persist-credentials: false
- name: "Run bazel build --nobuild //xla/... with retries"
run: parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --nobuild //xla/...
- name: "Run bazel cquery ... //xla/..."
run: |
set -euo pipefail
TAG_WITH_UNDERSCORES="${{ matrix.tag }}"
TAG_WITH_UNDERSCORES="${TAG_WITH_UNDERSCORES/-/_}"
if ! OUTPUT=$(bazel cquery --aspects "build_tools/dependencies/aspects.bzl%validate_${TAG_WITH_UNDERSCORES}_tag" //xla/... 2>&1); then
echo "Failed to run bazel cquery. Output:"
echo "$OUTPUT"
exit 1
fi
if echo "$OUTPUT" | grep 'Violation' >/dev/null; then
echo "The following dependency violations were found:"
echo "$OUTPUT" | grep 'Violation' | sed -e 's/^.*\[Violation\]/ -/'
echo ""
echo ""
echo "There are a couple of potential solutions for this/these violation(s):"
echo ""
echo "1. Tag the dependent target with the same tag as the dependee."
echo ""
echo "2. If unavoidable make the dependency selective using the"
echo " 'if_{gpu|cuda|rocm}_is_configured' macro. This is discouraged"
echo " outside of stream_executor."
echo ""
exit 1
fi
echo "No dependency violations found for tag '${{ matrix.tag }}'."
third_party/xla/.github/workflows/bazel_query.yml
Vendored
+47
View File
@@ -0,0 +1,47 @@
# Copyright 2024 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.
# ============================================================================
name: Bazel Query
permissions:
contents: read
on:
pull_request:
push:
branches:
- main
jobs:
bazel-query:
runs-on: "linux-x86-n2-16"
container:
image: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest" # zizmor: ignore[unpinned-images]
defaults:
run:
shell: bash
timeout-minutes: 10
steps:
- name: "Checking out repository"
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
with:
persist-credentials: false
- name: "Run bazel build --nobuild //xla/... with retries"
run: parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --nobuild //xla/...
- name: "Run bazel query //xla/..."
run: bazel query //xla/... > /dev/null
- name: "Run bazel query deps(//xla/...)"
run: bazel query "deps(//xla/...)" > /dev/null
- name: "Check that TSL doesn't depend on XLA"
run: |
bazel query "somepath(//xla/tsl/..., //xla/... -//xla/tsl/...)" 2>&1 | tee query_output.log
grep -q "INFO: Empty results" query_output.log
third_party/xla/.github/workflows/bazel_tags.yml
Vendored
+41
View File
@@ -0,0 +1,41 @@
# Copyright 2024 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.
# ============================================================================
name: Bazel Tags
permissions:
contents: read
on:
pull_request:
push:
branches:
- main
jobs:
bazel-tags:
runs-on: "linux-x86-n2-16"
container:
image: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest" # zizmor: ignore[unpinned-images]
defaults:
run:
shell: bash
timeout-minutes: 10
steps:
- name: "Checking out repository"
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
with:
persist-credentials: false
- name: "Run bazel build --nobuild //xla/... with retries"
run: parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --nobuild //xla/...
- name: "Assert all tags are documented"
run: bazel query //xla/... --output=build | python3 build_tools/lint/tags.py
third_party/xla/.github/workflows/benchmark_postsubmit.yml
Vendored
+279
View File
@@ -0,0 +1,279 @@
# 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.
# ============================================================================
name: Postsubmit Benchmarks
permissions:
contents: read
on:
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'
push:
branches:
- main
concurrency:
group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
cancel-in-progress: ${{ github.ref != 'main' }}
jobs:
Tests:
strategy:
# Don't fail fast - want to see results for all builds even if one fails.
fail-fast: false
matrix:
job_info: [
{
pool: "linux-arm64-t2a-48",
container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build-arm64:latest",
pretty_name: "XLA Linux ARM64 CPU 48 vcpu Presubmit",
bazel_arch_dir: "aarch64-opt",
platform: "CPU",
runner_binary: "hlo_runner_main",
hlo_test_file: "xla/tools/hlo_opt/tests/cpu_hlo.hlo"
},
{
pool: "linux-x86-n2-128",
container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest",
pretty_name: "XLA Linux x86 CPU 128 vcpu Presubmit",
bazel_arch_dir: "k8-opt",
platform: "CPU",
runner_binary: "hlo_runner_main",
hlo_test_file: "xla/tools/hlo_opt/tests/cpu_hlo.hlo"
},
{
pool: "linux-x86-g2-16-l4-1gpu",
container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build-cuda12.8-cudnn9.8:latest",
pretty_name: "XLA Linux x86 GPU L4 16 vcpu Presubmit",
bazel_arch_dir: "k8-opt",
platform: "GPU",
runner_binary: "hlo_runner_main_gpu",
hlo_test_file: "xla/tools/hlo_opt/tests/gpu_hlo_backend.hlo"
},
]
name: ${{ matrix.job_info.pretty_name }}
runs-on: ${{ matrix.job_info.pool }}
container: ${{ matrix.job_info.container }} # zizmor: ignore[unpinned-images]
defaults:
run:
shell: bash
timeout-minutes: 60
env:
OUTPUT_FILE_PATH: ${{ github.workspace }}/hlo_test_output.txt
steps:
- name: Print GitHub Context
run: |
echo "GitHub SHA: ${{ github.sha }}"
echo "GitHub Ref: ${GITHUB_REF}"
echo "GitHub Ref Name: ${GITHUB_REF_NAME}"
echo "GitHub Head Ref: ${GITHUB_HEAD_REF})"
echo "GitHub Base Ref: ${GITHUB_BASE_REF})"
echo "GitHub Repository: ${{ github.repository }}"
echo "GitHub Run ID: ${{ github.run_id }}"
echo "GitHub Run Number: ${{ github.run_number }}"
echo "GitHub Workflow: ${GITHUB_WORKFLOW}"
echo "GitHub Actor: ${GITHUB_ACTOR}"
echo "GitHub Event Name: ${{ github.event_name }}"
if [[ "${{ github.event_name }}" == "pull_request" ]]; then
echo "Pull Request Number: ${{ github.event.pull_request.number }}"
echo "Pull Request Head Ref: ${GITHUB_EVENT_PULL_REQUEST_HEAD_REF}"
echo "Pull Request Base Ref: ${GITHUB_EVENT_PULL_REQUEST_BASE_REF}"
fi
env:
GITHUB_EVENT_PULL_REQUEST_HEAD_REF: ${{ github.event.pull_request.head.ref }}
GITHUB_EVENT_PULL_REQUEST_BASE_REF: ${{ github.event.pull_request.base.ref }}
# Find the current PR number, if any, because github context doesn't have it for push events.
- uses: jwalton/gh-find-current-pr@f3d61b485d2801773f7a07b2aaa3306bd8f8e653 # v1.3.5
id: find_pr
with:
# Can be "open", "closed", or "all". Defaults to "open".
state: all
# This will echo "Your PR is 7", or be skipped if there is no current PR.
- run: echo "PR_NUMBER=${STEPS_FIND_PR_OUTPUTS_PR}"
if: success() && steps.find_pr.outputs.number
env:
PR: ${{ steps.find_pr.outputs.pr }}
STEPS_FIND_PR_OUTPUTS_PR: ${{ steps.find_pr.outputs.pr }}
- name: Checkout OpenXLA
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
with:
persist-credentials: false
- name: Wait For Connection
uses: google-ml-infra/actions/ci_connection@7f5ca0c263a81ed09ea276524c1b9192f1304e3c
with:
halt-dispatch-input: ${{ inputs.halt-for-connection }}
- name: Configure GPU backend
if: ${{ matrix.job_info.platform == 'GPU' }}
run: |
./configure.py --backend=CUDA --cuda_compiler=nvcc
- name: "Run build.py"
run: |
./build_tools/ci/build.py --build="${{ matrix.job_info.pretty_name }}_github_actions"
# Run the corresponding HLO tests based on platform
- name: Run HLO tests
run: |
bazel_arch_dir="${{ matrix.job_info.bazel_arch_dir }}"
binary_dir="./bazel-out/${bazel_arch_dir}/bin/xla/tools"
runner_binary="${{ matrix.job_info.runner_binary }}"
test_hlo_file="${{ matrix.job_info.hlo_test_file }}"
platform="${{ matrix.job_info.platform }}"
output_file="$GITHUB_WORKSPACE/${test_hlo_file}_output.txt"
xspace_file="$GITHUB_WORKSPACE/xspace.pb"
echo "Running ${platform} test with binary: $binary_dir/$runner_binary"
pwd #print working directory
if [[ "$platform" == "CPU" ]]; then
$binary_dir/multihost_hlo_runner/$runner_binary --device_type=host --num_repeats=5 --use_spmd_partitioning --profile_execution=True --xla_gpu_dump_xspace_to="$xspace_file" "$test_hlo_file" > "$output_file"
$binary_dir/compute_xspace_stats_main --input="$xspace_file" --device_type=CPU >> "$output_file"
elif [[ "$platform" == "GPU" ]]; then
$binary_dir/multihost_hlo_runner/$runner_binary --device_type=gpu --num_repeats=5 --use_spmd_partitioning --profile_execution=True --xla_gpu_dump_xspace_to="$xspace_file" "$test_hlo_file" > "$output_file"
$binary_dir/compute_xspace_stats_main_gpu --input="$xspace_file" --device_type=GPU >> "$output_file"
else
echo "Unsupported platform: $platform"
exit 1
fi
cat "$output_file"
echo "Output written to: $output_file"
echo "XSpace file written to: $xspace_file"
echo "OUTPUT_FILE=$output_file" >> "$GITHUB_ENV"
echo "XSPACE_FILE=$xspace_file" >> "$GITHUB_ENV"
- name: Upload HLO test output to a GCS bucket
run: |
GCS_BUCKET="gs://openxla-postsubmit-transient"
TIMESTAMP=$(date +%Y%m%d_%H%M%S)
DATE_FOLDER=$(date +%Y%m%d)
FILENAME=$(basename "$OUTPUT_FILE")
COMMIT_SHA="${{ github.sha }}"
RUN_ID="${{ github.run_id }}"
PR_NUMBER="${STEPS_FIND_PR_OUTPUTS_PR}"
PLATFORM="${{ matrix.job_info.platform }}"
POOL="${{ matrix.job_info.pool }}"
PR_SUFFIX=""
if [[ -n "$PR_NUMBER" ]]; then
PR_SUFFIX="_pr_${PR_NUMBER}"
fi
GCS_OBJECT_NAME="${DATE_FOLDER}/${TIMESTAMP}_${POOL}_run_${RUN_ID}_commit_${COMMIT_SHA}${PR_SUFFIX}_${PLATFORM}_${FILENAME}"
echo "Uploading $OUTPUT_FILE to $GCS_BUCKET/$GCS_OBJECT_NAME"
ls -l "$OUTPUT_FILE" # Verify the file exists
gcloud storage cp "$OUTPUT_FILE" "$GCS_BUCKET/$GCS_OBJECT_NAME"
env:
STEPS_FIND_PR_OUTPUTS_PR: ${{ steps.find_pr.outputs.pr }}
- name: Download Gemma Hlo Files
run: |
mkdir -p tmp_hlo
cd tmp_hlo
wget https://storage.googleapis.com/xla-benchmarking-temp/gemma3_1b_flax_call.hlo
cd ..
- name: Run Gemma HLO tests
if: ${{ matrix.job_info.pool != 'linux-arm64-t2a-48' }}
run: |
bazel_arch_dir="${{ matrix.job_info.bazel_arch_dir }}"
binary_dir="./bazel-out/${bazel_arch_dir}/bin/xla/tools"
runner_binary="${{ matrix.job_info.runner_binary }}"
test_hlo_file="tmp_hlo/gemma3_1b_flax_call.hlo"
platform="${{ matrix.job_info.platform }}"
output_file="$GITHUB_WORKSPACE/${test_hlo_file}_output.txt"
xspace_file="$GITHUB_WORKSPACE/xspace.pb"
if [[ "$platform" == "CPU" ]]; then
$binary_dir/multihost_hlo_runner/$runner_binary --device_type=host --num_repeats=5 --use_spmd_partitioning --profile_execution=True --xla_gpu_dump_xspace_to="$xspace_file" "$test_hlo_file" > "$output_file"
$binary_dir/compute_xspace_stats_main --input="$xspace_file" --device_type=CPU >> "$output_file"
elif [[ "$platform" == "GPU" ]]; then
$binary_dir/multihost_hlo_runner/$runner_binary --device_type=gpu --num_repeats=5 --use_spmd_partitioning --profile_execution=True --xla_gpu_dump_xspace_to="$xspace_file" "$test_hlo_file" > "$output_file"
$binary_dir/compute_xspace_stats_main_gpu --input="$xspace_file" --device_type=GPU >> "$output_file"
else
echo "Unsupported platform: $platform"
exit 1
fi
cat "$output_file"
echo "Output written to: $output_file"
echo "XSpace file written to: $xspace_file"
echo "OUTPUT_FILE=$output_file" >> "$GITHUB_ENV"
echo "XSPACE_FILE=$xspace_file" >> "$GITHUB_ENV"
- name: Upload gemma HLO test output to a GCS bucket
if: ${{ matrix.job_info.pool != 'linux-arm64-t2a-48' }}
run: |
GCS_BUCKET="gs://openxla-postsubmit-transient"
TIMESTAMP=$(date +%Y%m%d_%H%M%S)
DATE_FOLDER=$(date +%Y%m%d)
FILENAME="gemma3_1b_flax_call.hlo"
COMMIT_SHA="${{ github.sha }}"
RUN_ID="${{ github.run_id }}"
PR_NUMBER="${STEPS_FIND_PR_OUTPUTS_PR}"
PLATFORM="${{ matrix.job_info.platform }}"
POOL="${{ matrix.job_info.pool }}"
PR_SUFFIX=""
if [[ -n "$PR_NUMBER" ]]; then
PR_SUFFIX="_pr_${PR_NUMBER}"
fi
GCS_OBJECT_NAME="${DATE_FOLDER}/${TIMESTAMP}_${POOL}_run_${RUN_ID}_commit_${COMMIT_SHA}${PR_SUFFIX}_${PLATFORM}_${FILENAME}"
echo "Uploading $OUTPUT_FILE to $GCS_BUCKET/$GCS_OBJECT_NAME"
ls -l "$OUTPUT_FILE" # Verify the file exists
gcloud storage cp "$OUTPUT_FILE" "$GCS_BUCKET/$GCS_OBJECT_NAME"
env:
STEPS_FIND_PR_OUTPUTS_PR: ${{ steps.find_pr.outputs.pr }}
- name: Upload XSpace artifacts
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: xspace-artifacts-${{ matrix.job_info.pool }}-${{ matrix.job_info.platform }}
path: ${{ env.XSPACE_FILE }}
third_party/xla/.github/workflows/benchmark_presubmit.yml
Vendored
+154
View File
@@ -0,0 +1,154 @@
# 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.
# ============================================================================
name: Presubmit Benchmarks
permissions:
contents: read
on:
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'
pull_request:
branches:
- main
concurrency:
group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
cancel-in-progress: ${{ github.ref != 'main' }}
jobs:
Tests:
strategy:
# Don't fail fast - want to see results for all builds even if one fails.
fail-fast: false
matrix:
job_info: [
{
pool: "linux-x86-n2-128",
container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest",
pretty_name: "XLA Linux x86 CPU 128 vcpu Presubmit",
bazel_arch_dir: "k8-opt",
platform: "CPU"
},
{
pool: "linux-x86-g2-16-l4-1gpu",
container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build-cuda12.8-cudnn9.8:latest",
pretty_name: "XLA Linux x86 GPU L4 16 vcpu Presubmit",
bazel_arch_dir: "k8-opt",
platform: "GPU"
},
]
name: ${{ matrix.job_info.pretty_name }}
runs-on: ${{ matrix.job_info.pool }}
container: ${{ matrix.job_info.container }} # zizmor: ignore[unpinned-images]
defaults:
run:
shell: bash
timeout-minutes: 15
env:
OUTPUT_DIR: ${{ github.workspace }}/output
steps:
- name: Print GitHub Context
run: |
echo "GitHub SHA: ${{ github.sha }}"
echo "GitHub Ref: ${GITHUB_REF}"
echo "GitHub Ref Name: ${GITHUB_REF_NAME}"
echo "GitHub Head Ref: ${GITHUB_HEAD_REF})"
echo "GitHub Base Ref: ${GITHUB_BASE_REF})"
echo "GitHub Repository: ${{ github.repository }}"
echo "GitHub Run ID: ${{ github.run_id }}"
echo "GitHub Run Number: ${{ github.run_number }}"
echo "GitHub Workflow: ${GITHUB_WORKFLOW}"
echo "GitHub Actor: ${GITHUB_ACTOR}"
echo "GitHub Event Name: ${{ github.event_name }}"
if [[ "${{ github.event_name }}" == "pull_request" ]]; then
echo "Pull Request Number: ${{ github.event.pull_request.number }}"
echo "Pull Request Head Ref: ${GITHUB_EVENT_PULL_REQUEST_HEAD_REF}"
echo "Pull Request Base Ref: ${GITHUB_EVENT_PULL_REQUEST_BASE_REF}"
fi
env:
GITHUB_EVENT_PULL_REQUEST_HEAD_REF: ${{ github.event.pull_request.head.ref }}
GITHUB_EVENT_PULL_REQUEST_BASE_REF: ${{ github.event.pull_request.base.ref }}
- name: Checkout OpenXLA
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
with:
persist-credentials: false
- name: Configure GPU backend
if: ${{ matrix.job_info.platform == 'GPU' }}
run: |
./configure.py --backend=CUDA --cuda_compiler=nvcc
- name: "Run build.py"
run: |
./build_tools/ci/build.py --build="${{ matrix.job_info.pretty_name }}_github_actions"
- name: Download Gemma Hlo Files
run: |
mkdir -p "$OUTPUT_DIR/tmp_hlo"
cd "$OUTPUT_DIR/tmp_hlo"
wget https://storage.googleapis.com/xla-benchmarking-temp/gemma3_1b_flax_call.hlo
cd ..
# Run the corresponding HLO tests based on platform
- name: Run HLO tests
run: |
BAZEL_ARCH_DIR="${{ matrix.job_info.bazel_arch_dir }}"
BINARY_DIR="./bazel-out/${BAZEL_ARCH_DIR}/bin/xla/tools"
OUTPUT_FILE_PATH="$OUTPUT_DIR/gemma3_1b_flax_call_output.txt"
TEST_HLO_FILE="$OUTPUT_DIR/tmp_hlo/gemma3_1b_flax_call.hlo"
XSPACE_FILE_PATH="$OUTPUT_DIR/gemma3_1b_flax_call_xspace.pb"
DEVICE_TYPE=""
if [[ ${{ matrix.job_info.platform }} == "CPU" ]]; then
DEVICE_TYPE="host"
HLO_RUNNER_BINARY="$BINARY_DIR/multihost_hlo_runner/hlo_runner_main"
COMPUTE_XSPACE_STATS_BINARY="$BINARY_DIR/compute_xspace_stats_main"
elif [[ ${{ matrix.job_info.platform }} == "GPU" ]]; then
DEVICE_TYPE="gpu"
HLO_RUNNER_BINARY="$BINARY_DIR/multihost_hlo_runner/hlo_runner_main_gpu"
COMPUTE_XSPACE_STATS_BINARY="$BINARY_DIR/compute_xspace_stats_main_gpu"
else
echo "Unsupported platform: ${{ matrix.job_info.platform }}"
exit 1
fi
echo "Running test with binary: $HLO_RUNNER_BINARY"
pwd #print working directory
$HLO_RUNNER_BINARY --device_type=$DEVICE_TYPE --use_spmd_partitioning --num_repeats=5 --profile_execution=True --xla_gpu_dump_xspace_to="$XSPACE_FILE_PATH" "$TEST_HLO_FILE" > "$OUTPUT_FILE_PATH"
$COMPUTE_XSPACE_STATS_BINARY --input="$XSPACE_FILE_PATH" --device_type="${{ matrix.job_info.platform }}" >> "$OUTPUT_FILE_PATH"
cat "$OUTPUT_FILE_PATH"
third_party/xla/.github/workflows/benchmarks/README.md
Vendored
+237
View File
@@ -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
Vendored Executable
+135
View File
@@ -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
Vendored
+340
View File
@@ -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
Vendored Executable
+86
View File
@@ -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
Vendored Executable
+276
View File
@@ -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
Vendored
+59
View File
@@ -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 ---"
third_party/xla/.github/workflows/buildifier.yml
Vendored
+43
View File
@@ -0,0 +1,43 @@
# Copyright 2023 The TensorFlow 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.
# ============================================================================
name: Buildifier
permissions:
contents: read
on:
pull_request:
push:
branches:
- main
env:
# Have `go install` place binaries in $PATH
GOBIN: "/usr/local/bin"
jobs:
buildifier-lint:
runs-on: ubuntu-22.04
defaults:
run:
shell: bash
timeout-minutes: 6
steps:
- name: "Checking out repository"
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
with:
persist-credentials: false
- name: "Install buildifier"
run: parallel --ungroup --retries 3 --delay 15 --nonall -- go install github.com/bazelbuild/buildtools/buildifier@433ea85 # 6.4.0
- name: "Run buildifier"
run: buildifier --lint=warn --warnings=-out-of-order-load -r xla/
third_party/xla/.github/workflows/check_contents.yml
Vendored
+140
View File
@@ -0,0 +1,140 @@
# Copyright 2023 The TensorFlow 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.
# ============================================================================
# The goal of these checks is to externalize all check_contents checks that run
# internally. For now, no automation will try to keep these in sync as many
# METADATA files are doing different checks, and we'll have to consolidate those
# files once XLA moves out of Tensorflow internally.
# TODO(ddunleavy): Update this after METADATA files are consolidated.
name: Check Contents
permissions:
contents: read
on:
pull_request:
push:
branches:
- main
env:
# A bit tricky here: this lets us invoke python files outside of `bazel run`
# as long as the python file has imports as if it expects bazel to invoke it,
# and only depends on other python files in this workspace or the standard
# library.
PYTHONPATH: "."
jobs:
check-contents:
runs-on: ubuntu-22.04
defaults:
run:
shell: bash
timeout-minutes: 6
if: |
github.event.sender.type == 'User' ||
contains(github.event.pull_request.body, 'FORCE_TEST_ACTIONS')
steps:
- name: "Checking out repository"
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
with:
persist-credentials: false
- name: "Fetch HEAD of main branch"
run: git fetch origin main --depth=1
# ----------------C++ checks----------------
- name: "Check for tsl::{Status,StatusOr}"
if: always() # We want to run all checks, not just til first failure
run: "python3 build_tools/lint/check_contents.py --path_regex_exclusion '.github/workflows/check_contents\\.yml' --prohibited_regex 'tsl::Status(Or)?' --suppression_regex 'TENSORFLOW_STATUS_OK' --failure_message 'Use plain Status(Or) rather than tsl::Status(Or) within XLA.'"
- name: "Check for tsl::Status::OK()"
if: always() # We want to run all checks, not just til first failure
run: "python3 build_tools/lint/check_contents.py --path_regex_exclusion '.github/workflows/check_contents\\.yml' --prohibited_regex 'tsl::Status::OK' --failure_message 'Use OkStatus instead of Status::OK.'"
- name: "Check for std::call_once"
if: always() # We want to run all checks, not just til first failure
run: "python3 build_tools/lint/check_contents.py --path_regex_exclusion '.github/workflows/check_contents\\.yml' --prohibited_regex 'std::call_once' --failure_message 'Use absl::call_once instead of std::call_once'"
- name: "Check for Abseil compatibility API"
if: always() # We want to run all checks, not just til first failure
run: "python3 build_tools/lint/check_contents.py --path_regex_exclusion '.github/workflows/check_contents\\.yml' --prohibited_regex '\\b(absl::(any|get|get_if|make_any|make_unique|nullopt|nullopt_t|MakeUnique|make_optional|optional|underlying_type_t|variant|visit))\\b' --suppression_regex '(?i)non-std ok' --failure_message 'Use standard library version instead of the Abseil compatibility API.'"
- name: "Use absl instead of Tensorflow/TSL functions/types"
if: always() # We want to run all checks, not just til first failure
run: "python3 build_tools/lint/check_contents.py --path_regex_exclusion '.github/workflows/check_contents\\.yml' --prohibited_regex '\\b(gtl::FlatMap|gtl::FlatSet|gtl::InlinedVector|gtl::optional|strings::Appendf|strings::Printf|strings::safe_strto64|strings::safe_strtof|strings::StrAppend|strings::StrCat|str_util::|tensorflow::StringPiece)\\b' --suppression_regex '(?i)non-absl ok' --failure_message 'Use Abseil types and functions instead of those in Tensorflow/TSL.'"
- name: "Check for Header Guards"
if: always()
run: "python3 build_tools/lint/check_header_guards.py"
- name: "Check for TF Status Macros"
if: always()
run: "python3 build_tools/lint/check_contents.py --path_regex '.*\\.(h|cc|cpp)' --path_regex_exclusion 'xla/tsl/.*' 'xla/python/.*' 'xla/pjrt/.*' --prohibited_regex '\\bTF_(RETURN_IF_ERROR|ASSIGN_OR_RETURN)\\b' --failure_message 'Within XLA, you can now use RETURN_IF_ERROR/ASSIGN_OR_RETURN from xla/tsl/platform/status_macros.h rather than their TSL TF_ equivalents.'"
# ----------------Python checks----------------
- name: "Check for python print()"
if: always() # We want to run all checks, not just til first failure
run: "python3 build_tools/lint/check_contents.py --path_regex '.*\\.py' --prohibited_regex '(^|\\n|\\s)print\\(.*\\)' --suppression_regex 'DISABLE_DEBUG_PRINT_CHECK' --failure_message 'Found debug print() statement left in the code. Add DISABLE_DEBUG_PRINT_CHECK to disable'"
- name: "Check for mock.patch"
if: always() # We want to run all checks, not just til first failure
run: "python3 build_tools/lint/check_contents.py --path_regex '.*\\.py' --prohibited_regex 'mock\\.patch\\(' --failure_message 'Please use match.patch.object instead of mock.patch.'"
third_party/xla/.github/workflows/ci.yml
Vendored
+224
View File
@@ -0,0 +1,224 @@
# 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.
# ============================================================================
name: CI
permissions:
contents: read
on:
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'
pull_request:
push:
branches:
- main
concurrency:
group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
cancel-in-progress: ${{ github.ref != 'main' }}
jobs:
Tests:
strategy:
# Don't fail fast - want to see results for all builds even if one fails.
fail-fast: false
matrix:
job_info: [
{
pool: "linux-x86-n2-16",
container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest",
name: "XLA Linux x86 CPU",
repo: "openxla/xla",
},
{
pool: "windows-x86-n2-16",
container: null,
name: "XLA Windows x86 CPU",
repo: "openxla/xla",
},
{
pool: "linux-x86-n2-16",
container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest",
name: "XLA Linux x86 CPU Bzlmod",
repo: "openxla/xla",
},
{
pool: "linux-x86-n2-16",
container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest",
name: "XLA Linux X86 GPU ONEAPI",
repo: "openxla/xla",
},
{
pool: "linux-x86-n2-16",
container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest",
name: "XLA Linux x86 GPU Hermetic ROCm",
repo: "openxla/xla",
},
{
pool: "linux-x86-g2-16-l4-1gpu",
container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest",
name: "XLA Linux x86 GPU L4",
repo: "openxla/xla",
},
{
pool: "linux-arm64-c4a-16",
container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build-arm64:latest",
name: "XLA Linux ARM64 CPU",
repo: "openxla/xla",
},
{
pool: "linux-x86-n2-16",
container: "gcr.io/tensorflow-sigs/build:latest-python3.11",
name: "JAX Linux x86 CPU Bzlmod",
repo: "jax-ml/jax",
},
{
pool: "windows-x86-n2-16",
container: null,
name: "JAX Windows x86 CPU",
repo: "jax-ml/jax",
},
{
pool: "linux-x86-g2-16-l4-1gpu",
container: "gcr.io/tensorflow-sigs/build:latest-python3.12",
name: "JAX Linux x86 GPU L4",
repo: "jax-ml/jax",
},
{
pool: "linux-x86-n2-16",
container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest",
name: "TensorFlow Linux x86 CPU",
repo: "tensorflow/tensorflow",
},
{
pool: "linux-x86-g2-16-l4-1gpu",
container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest",
name: "TensorFlow Linux x86 GPU L4",
repo: "tensorflow/tensorflow",
},
]
name: ${{ matrix.job_info.name }}
runs-on: ${{ matrix.job_info.pool }}
container: ${{ matrix.job_info.container }} # zizmor: ignore[unpinned-images]
defaults:
run:
shell: bash
steps:
- name: "Checking out openxla/xla"
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
with:
path: "openxla/xla"
persist-credentials: false
- name: Checking out ${{ matrix.job_info.repo }}
if: ${{ matrix.job_info.repo != 'openxla/xla' }}
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
with:
repository: ${{ matrix.job_info.repo }}
path: ${{ matrix.job_info.repo }}
persist-credentials: false
- name: Checking out rules_ml_toolchain # TODO(alekstheod): remove once jax is migrated to the latest rules_ml_toolchain changes
run: |
git clone https://github.com/google-ml-infra/rules_ml_toolchain.git openxla/rules_ml_toolchain
cd openxla/rules_ml_toolchain
git checkout b11745590f513ec55b32e2d126073576fde18c71
- name: "Wait For Connection"
uses: google-ml-infra/actions/ci_connection@7f5ca0c263a81ed09ea276524c1b9192f1304e3c
with:
halt-dispatch-input: ${{ inputs.halt-for-connection }}
- name: "Run build.py"
working-directory: ${{ matrix.job_info.repo }}
timeout-minutes: 60
run: |
# Shellcheck doesn't understand GitHub Actions expressions and incorrectly
# thinks this comparison will always fail.
# shellcheck disable=SC2193
if [[ "${{ matrix.job_info.pool }}" == *"windows"* ]]; then
python "$GITHUB_WORKSPACE"\\openxla\\xla\\build_tools\\ci\\build.py --build="${{ matrix.job_info.name }}_github_actions"
else
"$GITHUB_WORKSPACE"/openxla/xla/build_tools/ci/build.py --build="${{ matrix.job_info.name }}_github_actions"
fi
third_party/xla/.github/workflows/ci_multi_device.yml
Vendored
+85
View File
@@ -0,0 +1,85 @@
# 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.
# ============================================================================
name: Multi-Device CI
permissions:
contents: read
on:
workflow_dispatch: # Allows manual triggering
inputs:
halt-for-connection:
description: 'Should this workflow run wait for a remote connection?'
type: choice
required: true
default: 'no'
options:
- 'yes'
- 'no'
schedule:
- cron: "0 0 * * *" # Run at midnight every day
jobs:
Tests:
strategy:
# Don't fail fast - want to see results for all builds even if one fails.
fail-fast: false
matrix:
job_info: [
{
pool: "linux-x86-a3-8g-h100-8gpu",
container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest",
name: "XLA Linux x86 GPU 8X H100",
repo: "openxla/xla",
},
]
name: ${{ matrix.job_info.name }}
runs-on: ${{ matrix.job_info.pool }}
container: ${{ matrix.job_info.container }} # zizmor: ignore[unpinned-images]
defaults:
run:
shell: bash
timeout-minutes: 60
steps:
- name: "Checking out openxla/xla"
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
with:
path: "openxla/xla"
persist-credentials: false
- name: Checking out ${{ matrix.job_info.repo }}
if: ${{ matrix.job_info.repo != 'openxla/xla' }}
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
with:
repository: ${{ matrix.job_info.repo }}
path: ${{ matrix.job_info.repo }}
persist-credentials: false
- name: "Wait For Connection"
uses: google-ml-infra/actions/ci_connection@7f5ca0c263a81ed09ea276524c1b9192f1304e3c
with:
halt-dispatch-input: ${{ inputs.halt-for-connection }}
- name: "Run build.py"
working-directory: ${{ matrix.job_info.repo }}
run: |
# shellcheck disable=SC2193
if [[ "${{ matrix.job_info.pool }}" == *"windows"* ]]; then
python "$GITHUB_WORKSPACE"\\openxla\\xla\\build_tools\\ci\\build.py --build="${{ matrix.job_info.name }}_github_actions"
else
"$GITHUB_WORKSPACE"/openxla/xla/build_tools/ci/build.py --build="${{ matrix.job_info.name }}_github_actions"
fi
third_party/xla/.github/workflows/clang_format.yml
Vendored
+48
View File
@@ -0,0 +1,48 @@
# Copyright 2023 The TensorFlow 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.
# ============================================================================
name: Clang Format
permissions:
contents: read
on:
pull_request:
push:
branches:
- main
jobs:
clang-format:
runs-on: ubuntu-22.04
defaults:
run:
shell: bash
timeout-minutes: 6
if: |
github.event.sender.type == 'User' ||
contains(github.event.pull_request.body, 'FORCE_TEST_ACTIONS')
steps:
- name: "Checking out repository"
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
with:
fetch-depth: '0'
persist-credentials: false
- name: "Fetch HEAD of main branch"
run: git fetch origin main --depth=1
- name: Install uv dependencies
uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 # v8.2.0
- name: "Run clang-format check"
env:
TARGET_REF: "origin/${{ github.base_ref || 'main' }}"
run: ./build_tools/ci/run_clang_format.sh
third_party/xla/.github/workflows/clang_tidy.yml
Vendored
+63
View File
@@ -0,0 +1,63 @@
# Copyright 2026 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.
# ============================================================================
name: Clang Tidy
permissions:
contents: read
on:
pull_request:
jobs:
clang-tidy-cuda:
runs-on: "linux-x86-n2-16"
container:
image: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest" # zizmor: ignore[unpinned-images]
defaults:
run:
shell: bash
timeout-minutes: 30
continue-on-error: true
if: |
github.event.sender.type == 'User' ||
contains(github.event.pull_request.body, 'FORCE_TEST_ACTIONS')
steps:
- name: "Checking out repository"
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
with:
# Fetch more commits to ensure we can find the merge base for clang-tidy.
# This is better than 0 since we don't want to fetch everything as well.
fetch-depth: '100'
persist-credentials: false
- name: "Update base branch"
run: |
git config --global --add safe.directory "$GITHUB_WORKSPACE"
BASE_REF="${GITHUB_BASE_REF:-main}"
# Keep this fetch unshallow so clang-tidy can find the merge base even
# when the PR branch is more than checkout fetch-depth commits behind.
git fetch origin "refs/heads/${BASE_REF}:refs/remotes/origin/${BASE_REF}"
- name: Configure Workspace
# The actual compute capability does not matter. This allows deviceless run for this pipeline.
# The common configs ensure that bazel query and bazel build don't refetch repo mappings.
run: |
set -euo pipefail
./configure.py --backend=CUDA --cuda_compute_capabilities="9.0"
echo "common --config cuda_nvcc" >> xla_configure.bazelrc
echo "common --config rbe_linux_cpu" >> xla_configure.bazelrc
- name: "Run Clang Tidy"
env:
TARGET_REF: "origin/${{ github.base_ref || 'main' }}"
TAGS_TO_IGNORE: "no_oss|rocm-only|oneapi-only"
BAZEL_CMD: "bazel"
run: |
set -euo pipefail
./build_tools/ci/run_clang_tidy.sh
third_party/xla/.github/workflows/cpu_benchmarks_nightly.yml
Vendored
+248
View File
@@ -0,0 +1,248 @@
# 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.
# ============================================================================
name: Nightly CPU Benchmarks
permissions:
contents: read
on:
workflow_dispatch: # Allows manual triggering
schedule:
- cron: "0 0 * * *" # Run at midnight every day
jobs:
Tests:
strategy:
# Don't fail fast - want to see results for all builds even if one fails.
fail-fast: false
matrix:
job_info: [
{
pool: "linux-arm64-t2a-48",
container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build-arm64:latest",
pretty_name: "XLA Linux ARM64 CPU 48 vcpu Presubmit",
bazel_arch_dir: "aarch64-opt",
platform: "CPU",
output_dir: "cpu_48_output"
},
{
pool: "linux-x86-n2-128",
container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest",
pretty_name: "XLA Linux x86 CPU 128 vcpu Presubmit",
bazel_arch_dir: "k8-opt",
platform: "CPU",
output_dir: "cpu_128_output"
}
]
name: ${{ matrix.job_info.pretty_name }}
runs-on: ${{ matrix.job_info.pool }}
container: ${{ matrix.job_info.container }} # zizmor: ignore[unpinned-images]
defaults:
run:
shell: bash
timeout-minutes: 540
env:
OUTPUT_DIR: ${{ github.workspace }}/${{ matrix.job_info.output_dir }}
steps:
- name: Print machine specs
run: |
lscpu
free -h # Memory information
df -h # Disk space information
uname -a # Kernel information
- name: Check Python Version in Container
run: python3 --version
- name: Install Python 3.10 if not present (IN CONTAINER)
run: |
if ! python3 --version > /dev/null 2>&1; then # check for python3
echo "Python 3 not found, installing..."
apt-get update
apt-get install -y python3.10 python3-pip
else
echo "Python 3 found."
fi
- name: Checkout OpenXLA
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
with:
persist-credentials: false
- name: Create results directory
run: mkdir results
- name: Configure CPU backend
run: |
./configure.py --backend=CPU
- name: Download and run Gemma HLO files and collect runtime and compile time stats
if: ${{ matrix.job_info.pool != 'linux-arm64-t2a-48' }}
run: |
mkdir -p "$OUTPUT_DIR/tmp_hlo"
cd "$OUTPUT_DIR/tmp_hlo"
wget https://storage.googleapis.com/xla-benchmarking-temp/gemma2_2b_keras_jax.hlo
wget https://storage.googleapis.com/xla-benchmarking-temp/gemma3_1b_flax_call.hlo
cd ..
- name: Print GitHub Context
run: |
echo "GitHub SHA: ${{ github.sha }}"
echo "GitHub Ref: ${GITHUB_REF}"
echo "GitHub Ref Name: ${GITHUB_REF_NAME}"
echo "GitHub Head Ref: ${GITHUB_HEAD_REF})"
echo "GitHub Base Ref: ${GITHUB_BASE_REF})"
echo "GitHub Repository: ${{ github.repository }}"
echo "GitHub Run ID: ${{ github.run_id }}"
echo "GitHub Run Number: ${{ github.run_number }}"
echo "GitHub Workflow: ${GITHUB_WORKFLOW}"
echo "GitHub Actor: ${GITHUB_ACTOR}"
echo "GitHub Event Name: ${{ github.event_name }}"
if [[ "${{ github.event_name }}" == "pull_request" ]]; then
echo "Pull Request Number: ${{ github.event.pull_request.number }}"
echo "Pull Request Head Ref: ${GITHUB_EVENT_PULL_REQUEST_HEAD_REF}"
echo "Pull Request Base Ref: ${GITHUB_EVENT_PULL_REQUEST_BASE_REF}"
fi
env:
GITHUB_EVENT_PULL_REQUEST_HEAD_REF: ${{ github.event.pull_request.head.ref }}
GITHUB_EVENT_PULL_REQUEST_BASE_REF: ${{ github.event.pull_request.base.ref }}
# Find the current PR number, if any, because github context doesn't have it for push events.
- uses: jwalton/gh-find-current-pr@f3d61b485d2801773f7a07b2aaa3306bd8f8e653 # v1.3.5
id: findPr
with:
# Can be "open", "closed", or "all". Defaults to "open".
state: all
# This will echo "Your PR is 7", or be skipped if there is no current PR.
- run: echo "PR_NUMBER=${STEPS_FINDPR_OUTPUTS_PR}"
if: success() && steps.findPr.outputs.number
env:
PR: ${{ steps.findPr.outputs.pr }}
STEPS_FINDPR_OUTPUTS_PR: ${{ steps.findPr.outputs.pr }}
- name: "Run build.py"
run: |
./build_tools/ci/build.py --build="${{ matrix.job_info.pretty_name }}_github_actions"
- name: Run hlo_runner_main on cpu_hlo.hlo and collect profile stats
run: |
bazel_arch_dir="${{ matrix.job_info.bazel_arch_dir }}"
binary_dir="./bazel-out/${bazel_arch_dir}/bin/xla/tools"
mkdir -p "$OUTPUT_DIR"
OUTPUT_FILE_PATH="$OUTPUT_DIR/cpu_hlo_output.txt"
XSPACE_FILE_PATH="$OUTPUT_DIR/cpu_hlo_xspace.pb"
test_hlo_file="xla/tools/hlo_opt/tests/cpu_hlo.hlo"
echo "Running CPU test with binary: $binary_dir"
pwd #print working directory
$binary_dir/multihost_hlo_runner/hlo_runner_main --device_type=host --num_repeats=5 --use_spmd_partitioning --profile_execution=True --xla_gpu_dump_xspace_to="$XSPACE_FILE_PATH" "$test_hlo_file" > "$OUTPUT_FILE_PATH"
$binary_dir/compute_xspace_stats_main --input="$XSPACE_FILE_PATH" --device_type=CPU >> "$OUTPUT_FILE_PATH"
cat "$OUTPUT_FILE_PATH"
- name: Run hlo_runner_main on gemma2_2b_keras_jax.hlo and collect profile stats
if: ${{ matrix.job_info.pool != 'linux-arm64-t2a-48' }}
run: |
bazel_arch_dir="${{ matrix.job_info.bazel_arch_dir }}"
binary_dir="./bazel-out/${bazel_arch_dir}/bin/xla/tools"
OUTPUT_FILE_PATH="$OUTPUT_DIR/gemma2_2b_keras_jax_output.txt"
XSPACE_FILE_PATH="$OUTPUT_DIR/gemma2_2b_xspace.pb"
test_hlo_file="$OUTPUT_DIR/tmp_hlo/gemma2_2b_keras_jax.hlo"
echo "Running CPU test with binary: $binary_dir"
pwd #print working directory
$binary_dir/multihost_hlo_runner/hlo_runner_main --device_type=host --num_repeats=5 --use_spmd_partitioning --profile_execution=True --xla_gpu_dump_xspace_to="$XSPACE_FILE_PATH" "$test_hlo_file" > "$OUTPUT_FILE_PATH"
$binary_dir/compute_xspace_stats_main --input="$XSPACE_FILE_PATH" --device_type=CPU >> "$OUTPUT_FILE_PATH"
cat "$OUTPUT_FILE_PATH"
- name: Run hlo_runner_main on gemma3_1b_flax_call.hlo and collect profile stats
if: ${{ matrix.job_info.pool != 'linux-arm64-t2a-48' }}
run: |
bazel_arch_dir="${{ matrix.job_info.bazel_arch_dir }}"
binary_dir="./bazel-out/${bazel_arch_dir}/bin/xla/tools"
OUTPUT_FILE_PATH="$OUTPUT_DIR/gemma3_1b_flax_call_output.txt"
XSPACE_FILE_PATH="$OUTPUT_DIR/gemma3_1b_xspace.pb"
test_hlo_file="$OUTPUT_DIR/tmp_hlo/gemma3_1b_flax_call.hlo"
echo "Running CPU test with binary: $binary_dir"
pwd #print working directory
$binary_dir/multihost_hlo_runner/hlo_runner_main --device_type=host --num_repeats=5 --use_spmd_partitioning --profile_execution=True --xla_gpu_dump_xspace_to="$XSPACE_FILE_PATH" "$test_hlo_file" > "$OUTPUT_FILE_PATH"
$binary_dir/compute_xspace_stats_main --input="$XSPACE_FILE_PATH" --device_type=CPU >> "$OUTPUT_FILE_PATH"
cat "$OUTPUT_FILE_PATH"
- name: Upload HLO test output to a GCS bucket
if: ${{ matrix.job_info.pool != 'linux-arm64-t2a-48' }}
run: |
GCS_BUCKET="gs://openxla-nightly-transient"
TIMESTAMP=$(date +%Y%m%d_%H%M%S)
DATE_FOLDER=$(date +%Y%m%d)
FILENAME_GEMMA="gemma2_2b_keras_jax_output.txt"
FILENAME_CPU="cpu_hlo_output.txt"
FILENAME_GEMMA3="gemma3_1b_flax_call_output.txt"
COMMIT_SHA="${{ github.sha }}"
RUN_ID="${{ github.run_id }}"
PR_NUMBER="${STEPS_FINDPR_OUTPUTS_PR}"
OUTPUT_SUBDIR="${{ matrix.job_info.output_dir }}"
PR_SUFFIX=""
if [[ -n "$PR_NUMBER" ]]; then
PR_SUFFIX="_pr_${PR_NUMBER}"
fi
GEMMA_GCS_OBJECT_NAME="${DATE_FOLDER}/${TIMESTAMP}_${{ matrix.job_info.pool }}_run_${RUN_ID}_commit_${COMMIT_SHA}${PR_SUFFIX}_${OUTPUT_SUBDIR}_${FILENAME_GEMMA}"
CPU_GCS_OBJECT_NAME="${DATE_FOLDER}/${TIMESTAMP}_${{ matrix.job_info.pool }}_run_${RUN_ID}_commit_${COMMIT_SHA}${PR_SUFFIX}_${OUTPUT_SUBDIR}_${FILENAME_CPU}"
GEMMA3_GCS_OBJECT_NAME="${DATE_FOLDER}/${TIMESTAMP}_${{ matrix.job_info.pool }}_run_${RUN_ID}_commit_${COMMIT_SHA}${PR_SUFFIX}_${OUTPUT_SUBDIR}_${FILENAME_GEMMA3}"
echo "Uploading $OUTPUT_DIR/$FILENAME_GEMMA to $GCS_BUCKET/$GEMMA_GCS_OBJECT_NAME"
ls -l "$OUTPUT_DIR/$FILENAME_GEMMA" # Verify the file exists
gcloud storage cp "$OUTPUT_DIR/$FILENAME_GEMMA" "$GCS_BUCKET/$GEMMA_GCS_OBJECT_NAME"
echo "Uploading $OUTPUT_DIR/$FILENAME_CPU to $GCS_BUCKET/$CPU_GCS_OBJECT_NAME"
ls -l "$OUTPUT_DIR/$FILENAME_CPU" # Verify the file exists
gcloud storage cp "$OUTPUT_DIR/$FILENAME_CPU" "$GCS_BUCKET/$CPU_GCS_OBJECT_NAME"
echo "Uploading $OUTPUT_DIR/$FILENAME_GEMMA3 to $GCS_BUCKET/$GEMMA3_GCS_OBJECT_NAME"
ls -l "$OUTPUT_DIR/$FILENAME_GEMMA3" # Verify the file exists
gcloud storage cp "$OUTPUT_DIR/$FILENAME_GEMMA3" "$GCS_BUCKET/$GEMMA3_GCS_OBJECT_NAME"
env:
STEPS_FINDPR_OUTPUTS_PR: ${{ steps.findPr.outputs.pr }}
- name: Upload XSpace artifacts
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: cpu-xla-benchmarks-xspace-${{ matrix.job_info.pool }}
path: ${{ github.workspace }}/${{ matrix.job_info.output_dir }}/*_xspace.pb
- name: Run E2E benchmarks flax_2b
timeout-minutes: 60
run: |
cd xla/backends/cpu/benchmarks/e2e/gemma2/flax_2b
bash setup.sh
bash run.sh | tee -a ../../../../../../../results/flax_2b.log
# TODO(juliagmt): Add CPU XSpace collection support to hlo_runner_main.
third_party/xla/.github/workflows/dwyu.yml
Vendored
+64
View File
@@ -0,0 +1,64 @@
# Copyright 2026 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.
# ============================================================================
name: DWYU (Depend on What You Use)
permissions:
contents: read
on:
pull_request:
push:
branches:
- main
env:
GOBIN: "/usr/local/bin"
PYTHONPATH: "."
BANT_VERSION: "0.2.8"
jobs:
dwyu:
runs-on: ubuntu-22.04
defaults:
run:
shell: bash
timeout-minutes: 15
continue-on-error: true
steps:
- name: "Checking out repository"
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
with:
fetch-depth: '0'
persist-credentials: false
- name: "Fetch HEAD of main branch"
run: git fetch origin main --depth=1
- name: "Install bant"
run: |
curl -fSL "https://github.com/hzeller/bant/releases/download/v${BANT_VERSION}/bant-v${BANT_VERSION}-linux-static-x86_64.tar.gz" -o /tmp/bant.tar.gz
tar xzf /tmp/bant.tar.gz -C /tmp
sudo install "/tmp/bant-v${BANT_VERSION}-linux-static-x86_64/bin/bant" /usr/local/bin/bant
- name: "Install bazelisk"
run: parallel --ungroup --retries 3 --delay 15 --nonall -- go install github.com/bazelbuild/bazelisk@24651ab # v1.20.0
- name: "Run bazel build --nobuild //xla/... with retries"
run: parallel --ungroup --retries 3 --delay 15 --nonall -- bazelisk build --nobuild //xla/...
- name: "Find affected targets and run bant dwyu"
run: |
set -euo pipefail
TARGETS=$(python3 build_tools/lint/check_dwyu.py)
if [ -z "$TARGETS" ]; then
echo "No targets to check."
exit 0
fi
echo "Running: bant dwyu $TARGETS"
# shellcheck disable=SC2086
bant dwyu $TARGETS
third_party/xla/.github/workflows/generate_benchmark_matrix.yml
Vendored
+187
View File
@@ -0,0 +1,187 @@
# 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/generate_benchmark_matrix.yml
name: Generate Benchmark Matrix
permissions:
contents: read
on:
workflow_call:
inputs:
registry_file:
description: 'Path to the benchmark registry file (relative to repository root)'
required: false
type: string
default: 'xla/tools/benchmarks/registries/default_registry.yml'
workflow_type:
description: 'The type of workflow triggering this generation (e.g. PRESUBMIT, POSTSUBMIT, SCHEDULED, MANUAL, or aliases like nightly)'
required: true
type: string
checkout_ref:
description: 'The Git ref (branch, tag, or SHA) to checkout'
required: false
type: string
default: ''
outputs:
matrix_include_json:
description: 'JSON string representing the array for the "include" matrix strategy'
value: ${{ jobs.generate.outputs.matrix_json_output }}
jobs:
generate:
name: Generate Matrix (${{ inputs.workflow_type }})
runs-on: linux-x86-n2-64
container: us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest # zizmor: ignore[unpinned-images]
outputs:
matrix_json_output: ${{ steps.run_generator.outputs.matrix_json }}
defaults:
run:
shell: bash
timeout-minutes: 60
steps:
- name: Checkout OpenXLA
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
with:
# Use inputs.checkout_ref if provided, otherwise default to the event's ref
# (e.g., PR's HEAD SHA or caller's commit SHA)
ref: ${{ inputs.checkout_ref || github.ref }}
fetch-depth: 0 # Fetch all history, might be needed by build/configure scripts
persist-credentials: false
- name: Configure OpenXLA
run: "echo \"Configuring OpenXLA for CPU to build the generator tool...\"\nif [ -f \"./configure.py\" ]; then\n ./configure.py --backend=CPU \nelse\n echo \"::warning::configure.py not found. Assuming C++ tool build doesn't require it or is pre-configured.\"\nfi\n"
# TODO(juliagmt): Use build.py to build the binary.
- name: Build generate_benchmark_matrices_main
run: |
echo "Building generator..."
if ! bazel build \
--build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd \
--test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd \
--config=warnings \
--config=nonccl \
--config=rbe_linux_cpu \
--color=yes \
--test_output=errors \
--verbose_failures \
--keep_going \
--nobuild_tests_only \
--profile=profile.json.gz \
--flaky_test_attempts=3 \
--jobs=150 \
--bes_upload_mode=fully_async \
//xla/tools/benchmarks/utils:generate_benchmark_matrices_main; then
echo "::error::Failed to build generate_benchmark_matrices_main"
exit 1
fi
- name: Run generate_benchmark_matrices_main
id: run_generator
run: |
BINARY_PATH="./bazel-bin/xla/tools/benchmarks/utils/generate_benchmark_matrices_main"
REGISTRY_PATH="${INPUTS_REGISTRY_FILE}"
# Convert workflow_type input to uppercase for the C++ binary,
# as the binary's GetWorkflowTypeFromStr expects uppercase or known aliases.
WORKFLOW_TYPE_ARG_UPPER=$(echo "${INPUTS_WORKFLOW_TYPE}" | tr '[:lower:]' '[:upper:]')
if [ ! -f "$BINARY_PATH" ]; then
echo "::error::Generator binary not found at $BINARY_PATH after build."
# List bazel-bin for debugging if binary not found
echo "Listing bazel-bin directory:"
ls -R bazel-bin || echo "bazel-bin directory not found."
exit 1
fi
echo "Generating matrix for workflow type: $WORKFLOW_TYPE_ARG_UPPER (original input: ${INPUTS_WORKFLOW_TYPE}) using registry: $REGISTRY_PATH"
# Execute and capture output
JSON_ARRAY_STRING=$("$BINARY_PATH" --registry_file="$REGISTRY_PATH" --workflow_type="$WORKFLOW_TYPE_ARG_UPPER")
BINARY_EXIT_CODE=$?
if [ $BINARY_EXIT_CODE -ne 0 ]; then
echo "::error::generate_benchmark_matrices_main failed with exit code $BINARY_EXIT_CODE."
echo "Output from binary was:"
echo "$JSON_ARRAY_STRING" # This might contain error messages from the binary's LOG(QFATAL)
exit 1
fi
# Basic validation: must be an array (even if empty "[]")
if ! echo "$JSON_ARRAY_STRING" | jq -e '. | type == "array"' > /dev/null; then
echo "::error::Generator output is not a valid JSON array."
echo "Output was: $JSON_ARRAY_STRING"
exit 1
fi
echo "Generated matrix JSON array string:"
echo "$JSON_ARRAY_STRING"
{
echo "matrix_json<<EOF_MATRIX_JSON"
echo "$JSON_ARRAY_STRING"
echo "EOF_MATRIX_JSON"
} >> "$GITHUB_OUTPUT"
env:
INPUTS_REGISTRY_FILE: ${{ inputs.registry_file }}
INPUTS_WORKFLOW_TYPE: ${{ inputs.workflow_type }}
third_party/xla/.github/workflows/github_api.py
Vendored
+160
View File
@@ -0,0 +1,160 @@
# Copyright 2023 The TensorFlow 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.
# ============================================================================
"""Provides a Python interface to various parts of the GitHub API.
NOTE: not using PyGithub https://github.com/PyGithub/PyGithub as it doesn't
have full API coverage and it's easy enough to use the endpoints we need like
this (with zero dependencies as a bonus!)
"""
from typing import Any, Optional
import urllib
import requests
class GitHubAPI:
"""Wraps the GitHub REST API."""
def __init__(self, token: Optional[str] = None):
self._session = requests.Session()
self._session.headers["Accept"] = "application/vnd.github+json"
if token:
self._session.headers["Authorization"] = f"token {token}"
def _make_request(
self, verb: str, endpoint: str, **kwargs: dict[str, Any]
) -> requests.Response:
"""Helper method to make a request and raise an HTTPError if one occurred.
Arguments:
verb: The HTTP verb to use
endpoint: The endpoint to make the request to
**kwargs: The json that will be sent as the body of the request.
Returns:
a requests.Response object containing the response from the API.
Raises:
requests.exceptions.HTTPError
"""
res = self._session.request(
verb,
urllib.parse.urljoin("https://api.github.com", endpoint),
json=kwargs,
)
res.raise_for_status()
return res.json()
def get_commit(self, repo: str, commit_id: str) -> requests.Response:
"""Gets a commit by it's SHA-1 hash.
https://docs.github.com/en/rest/commits/commits?apiVersion=2022-11-28#get-a-
commit
Arguments:
repo: a string of the form `owner/repo_name`, e.g. openxla/xla.
commit_id: a string describing the commit to get, e.g. `deadbeef` or
`HEAD`.
Returns:
a requests.Response object containing the response from the API.
Raises:
requests.exceptions.HTTPError
"""
endpoint = f"repos/{repo}/commits/{commit_id}"
return self._make_request("GET", endpoint)
def write_issue_comment(
self, repo: str, issue_number: int, body: str
) -> requests.Response:
"""Writes a comment on an issue (or PR).
https://docs.github.com/en/rest/issues/comments?apiVersion=2022-11-
28#create-an-issue-comment
Arguments:
repo: a string of the form `owner/repo_name`, e.g. openxla/xla
issue_number: the issue (or PR) to comment on
body: the body of the comment
Returns:
a requests.Response object containing the response from the API.
Raises:
requests.exceptions.HTTPError
"""
endpoint = f"repos/{repo}/issues/{issue_number}/comments"
return self._make_request("POST", endpoint, body=body)
def set_issue_status(
self, repo: str, issue_number: int, status: str
) -> requests.Response:
"""Sets the status of an issue (or PR).
https://docs.github.com/en/rest/issues/issues?apiVersion=2022-11-28#update-
an-issue
Arguments:
repo: a string of the form `owner/repo_name`, e.g. openxla/xla
issue_number: the issue (or PR) to set the status of
status: the status to set
Returns:
a requests.Response object containing the response from the API.
Raises:
requests.exceptions.HTTPError
"""
endpoint = f"repos/{repo}/issues/{issue_number}"
return self._make_request("POST", endpoint, status=status)
def add_issue_labels(
self, repo: str, issue_number: int, labels: list[str]
) -> requests.Response:
"""Adds labels to an issue (or PR).
https://docs.github.com/en/actions/managing-issues-and-pull-requests/adding-labels-to-issues
Arguments:
repo: a string of the form `owner/repo_name`, e.g. openxla/xla
issue_number: the issue (or PR) to set the status of
labels: the labels to add to the issue
Returns:
a requests.Response object containing the response from the API.
Raises:
requests.exceptions.HTTPError
"""
endpoint = f"repos/{repo}/issues/{issue_number}/labels"
return self._make_request("POST", endpoint, labels=labels)
def get_user_orgs(self, username: str) -> requests.Response:
"""Gets all public org memberships for a user.
https://docs.github.com/en/rest/orgs/orgs?apiVersion=2022-11-28#list-organizations-for-a-user
Arguments:
username: The user's GitHub username as a string.
Returns:
a requests.Response object containing the response from the API.
Raises:
requests.exceptions.HTTPError
"""
endpoint = f"users/{username}/orgs"
return self._make_request("GET", endpoint, username=username)
third_party/xla/.github/workflows/gpu_benchmarks_nightly.yml
Vendored
+211
View File
@@ -0,0 +1,211 @@
# 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.
# ============================================================================
name: Nightly GPU Benchmarks
permissions:
contents: read
on:
workflow_dispatch: # Allows manual triggering
schedule:
- cron: "0 0 * * *" # Run at midnight every day
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-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest",
pretty_name: "XLA Linux x86 GPU L4 48 vcpu Presubmit",
},
{
os: "linux-x86-g2-16-l4-1gpu",
container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest",
pretty_name: "XLA Linux x86 GPU L4 16 vcpu Presubmit",
},
{
os: "linux-x86-a4-224-b200-1gpu",
container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build-cuda12.8-cudnn9.8:latest",
pretty_name: "XLA Linux x86 GPU A4 224 vcpu Presubmit",
},
# Expect more GPU types in the future.
]
name: ${{ matrix.job_info.pretty_name }}
runs-on: ${{ matrix.job_info.os }}
container: ${{ matrix.job_info.container }} # zizmor: ignore[unpinned-images]
defaults:
run:
shell: bash
timeout-minutes: 240
env:
OUTPUT_DIR: ${{ github.workspace }}/output
steps:
- name: Checkout XLA
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
with:
persist-credentials: false
- name: Download Gemma Hlo Files
run: |
mkdir -p tmp_hlo
cd tmp_hlo
wget https://storage.googleapis.com/xla-benchmarking-temp/gemma2_2b_keras_jax.hlo
wget https://storage.googleapis.com/xla-benchmarking-temp/gemma3_1b_flax_call.hlo
wget https://storage.googleapis.com/xla-benchmarking-temp/gemma3_1b_flax_sample_loop.hlo
cd ..
- name: Configure GPU backend
run: |
./configure.py --backend=CUDA --cuda_compiler=nvcc
- name: "Run build.py"
run: |
./build_tools/ci/build.py --build="${{ matrix.job_info.pretty_name }}_github_actions"
- name: Run HLO tests and collect data
run: |
binary_dir="./bazel-out/k8-opt/bin/xla/tools"
mkdir -p "$OUTPUT_DIR"
# Run gpu_hlo_backend.hlo
HLO_FILE_GB="xla/tools/hlo_opt/tests/gpu_hlo_backend.hlo"
OUTPUT_PREFIX_GB="$OUTPUT_DIR/gpu_hlo_backend"
echo "Running GPU test: $HLO_FILE_GB"
$binary_dir/multihost_hlo_runner/hlo_runner_main_gpu \
--device_type=gpu \
--num_repeats=5 \
--use_spmd_partitioning \
--profile_execution=True \
--xla_gpu_dump_xspace_to="${OUTPUT_PREFIX_GB}_xspace.pb" \
"$HLO_FILE_GB" \
> "${OUTPUT_PREFIX_GB}.txt"
$binary_dir/compute_xspace_stats_main_gpu \
--input="${OUTPUT_PREFIX_GB}_xspace.pb" \
--device_type=GPU \
>> "${OUTPUT_PREFIX_GB}.txt"
cat "${OUTPUT_PREFIX_GB}.txt"
# Run gemma2_2b_keras_jax.hlo
HLO_FILE_GEMMA="tmp_hlo/gemma2_2b_keras_jax.hlo"
OUTPUT_PREFIX_GEMMA="$OUTPUT_DIR/gemma2_2b_keras_jax"
echo "Running GPU test: $HLO_FILE_GEMMA"
$binary_dir/multihost_hlo_runner/hlo_runner_main_gpu \
--device_type=gpu \
--num_repeats=5 \
--use_spmd_partitioning \
--profile_execution=True \
--xla_gpu_dump_xspace_to="${OUTPUT_PREFIX_GEMMA}_xspace.pb" \
"$HLO_FILE_GEMMA" \
> "${OUTPUT_PREFIX_GEMMA}.txt"
$binary_dir/compute_xspace_stats_main_gpu \
--input="${OUTPUT_PREFIX_GEMMA}_xspace.pb" \
--device_type=GPU \
>> "${OUTPUT_PREFIX_GEMMA}.txt"
cat "${OUTPUT_PREFIX_GEMMA}.txt"
echo "Output written to: ${OUTPUT_PREFIX_GB}.txt and ${OUTPUT_PREFIX_GEMMA}.txt"
# Run gemma3_1b_flax_call.hlo
HLO_FILE_GEMMA3_CALL="tmp_hlo/gemma3_1b_flax_call.hlo"
OUTPUT_PREFIX_GEMMA3_CALL="$OUTPUT_DIR/gemma3_1b_flax_call"
echo "Running GPU test: $HLO_FILE_GEMMA3_CALL"
$binary_dir/multihost_hlo_runner/hlo_runner_main_gpu \
--device_type=gpu \
--num_repeats=5 \
--use_spmd_partitioning \
--profile_execution=True \
--xla_gpu_dump_xspace_to="${OUTPUT_PREFIX_GEMMA3_CALL}_xspace.pb" \
"$HLO_FILE_GEMMA3_CALL" \
> "${OUTPUT_PREFIX_GEMMA3_CALL}.txt"
$binary_dir/compute_xspace_stats_main_gpu \
--input="${OUTPUT_PREFIX_GEMMA3_CALL}_xspace.pb" \
--device_type=GPU \
>> "${OUTPUT_PREFIX_GEMMA3_CALL}.txt"
cat "${OUTPUT_PREFIX_GEMMA3_CALL}.txt"
echo "Output written to: ${OUTPUT_PREFIX_GEMMA3_CALL}.txt"
# Run gemma3_1b_flax_sample_loop.hlo
HLO_FILE_GEMMA3_SAMPLE_LOOP="tmp_hlo/gemma3_1b_flax_sample_loop.hlo"
OUTPUT_PREFIX_GEMMA3_SAMPLE_LOOP="$OUTPUT_DIR/gemma3_1b_flax_sample_loop"
echo "Running GPU test: $HLO_FILE_GEMMA3_SAMPLE_LOOP"
$binary_dir/multihost_hlo_runner/hlo_runner_main_gpu \
--device_type=gpu \
--num_repeats=5 \
--use_spmd_partitioning \
--profile_execution=True \
--xla_gpu_dump_xspace_to="${OUTPUT_PREFIX_GEMMA3_SAMPLE_LOOP}_xspace.pb" \
"$HLO_FILE_GEMMA3_SAMPLE_LOOP" \
> "${OUTPUT_PREFIX_GEMMA3_SAMPLE_LOOP}.txt"
$binary_dir/compute_xspace_stats_main_gpu \
--input="${OUTPUT_PREFIX_GEMMA3_SAMPLE_LOOP}_xspace.pb" \
--device_type=GPU \
>> "${OUTPUT_PREFIX_GEMMA3_SAMPLE_LOOP}.txt"
cat "${OUTPUT_PREFIX_GEMMA3_SAMPLE_LOOP}.txt"
echo "Output written to: ${OUTPUT_PREFIX_GEMMA3_SAMPLE_LOOP}.txt"
- name: Upload HLO test output to GCS
run: |
GCS_BUCKET="gs://openxla-nightly-transient"
TIMESTAMP=$(date +%Y%m%d_%H%M%S)
DATE_FOLDER=$(date +%Y%m%d)
COMMIT_SHA="${{ github.sha }}"
RUN_ID="${{ github.run_id }}"
upload_to_gcs() {
local base_name="$1"
local gcs_file_name="${DATE_FOLDER}/${TIMESTAMP}_${{ matrix.job_info.os }}_run_${RUN_ID}_commit_${COMMIT_SHA}_${base_name}.txt"
echo "Uploading $OUTPUT_DIR/${base_name}.txt to $GCS_BUCKET/$gcs_file_name"
gcloud storage cp "$OUTPUT_DIR/${base_name}.txt" "$GCS_BUCKET/$gcs_file_name"
}
# Upload output for gpu_hlo_backend
GB_BASE_NAME="gpu_hlo_backend"
upload_to_gcs "$GB_BASE_NAME"
# Upload output for gemma2_2b_keras_jax
GEMMA_BASE_NAME="gemma2_2b_keras_jax"
upload_to_gcs "$GEMMA_BASE_NAME"
# Upload output for gemma3_1b_flax_call
GEMMA3_CALL_BASE_NAME="gemma3_1b_flax_call"
upload_to_gcs "$GEMMA3_CALL_BASE_NAME"
# Upload output for gemma3_1b_flax_sample_loop
GEMMA3_SAMPLE_LOOP_BASE_NAME="gemma3_1b_flax_sample_loop"
upload_to_gcs "$GEMMA3_SAMPLE_LOOP_BASE_NAME"
- name: Upload XSpace artifacts
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: gpu-xla-benchmarks-xspace-${{ matrix.job_info.os }}
path: ${{ github.workspace }}/output/*_xspace.pb
third_party/xla/.github/workflows/license_check.yml
Vendored
+43
View File
@@ -0,0 +1,43 @@
# Copyright 2026 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.
# ============================================================================
name: License Check
permissions: {}
on:
pull_request:
push:
branches:
- main
jobs:
license-check:
runs-on: ubuntu-latest
timeout-minutes: 10
steps:
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
with:
persist-credentials: false
- name: "Set up Go"
uses: actions/setup-go@924ae3a1cded613372ab5595356fb5720e22ba16 # v6.5.0
with:
go-version: 'stable'
- name: "Run license check"
uses: google-ml-infra/actions/ci_check_license@432d36c6ac62d76fa99b3679758ccfbc10140fb6
with:
copyright-holder: "The XLA Authors"
license-type: "apache"
exclude-paths: "third_party/"
exclude-files: "BUILD _toc.yaml MHLOConfig.cmake.in"
third_party/xla/.github/workflows/nightly_benchmarks.yml
Vendored
+202
View File
@@ -0,0 +1,202 @@
# 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/nightly_benchmarks.yml
name: Nightly Benchmarks
permissions:
contents: read
on:
workflow_dispatch: # Allows manual triggering
inputs:
halt-for-connection:
description: 'Should this workflow run wait for a remote connection?'
type: choice
required: true
default: 'no'
options:
- 'yes'
- 'no'
schedule:
- cron: "0 0 * * *" # Run at midnight every day
jobs:
# =================================================================
# Job 1: Generate the matrix specifically for SCHEDULED benchmarks
# =================================================================
generate_matrix:
name: Generate Scheduled Matrix
uses: ./.github/workflows/generate_benchmark_matrix.yml
with:
workflow_type: 'SCHEDULED'
registry_file: 'xla/tools/benchmarks/registries/default_registry.yml'
checkout_ref: ${{ github.sha }} # On push/dispatch to main, github.sha is the commit SHA
run_benchmarks:
name: Run Benchmark (${{ matrix.benchmark_entry.config_id }}) # config_id will indicate the workflow type, e.g., '_scheduled'
needs: generate_matrix
# Also check if the generate_matrix job itself was skipped due to its `if` condition
if: success() && needs.generate_matrix.result != 'skipped' && needs.generate_matrix.outputs.matrix_include_json != '[]' && needs.generate_matrix.outputs.matrix_include_json != ''
strategy:
fail-fast: false # Continue running other benchmarks even if one fails
matrix:
benchmark_entry: ${{ fromJson(needs.generate_matrix.outputs.matrix_include_json || '[]') }}
runs-on: ${{ matrix.benchmark_entry.runner_label }}
container: ${{ matrix.benchmark_entry.container_image }}
defaults:
run:
shell: bash
timeout-minutes: 60
env:
# Variables from the matrix
BENCHMARK_NAME: ${{ matrix.benchmark_entry.benchmark_name }}
CONFIG_ID: ${{ matrix.benchmark_entry.config_id }}
RUNNER_LABEL: ${{ matrix.benchmark_entry.runner_label }}
CONTAINER_IMAGE: ${{ matrix.benchmark_entry.container_image }}
ARTIFACT_LOCATION: ${{ matrix.benchmark_entry.artifact_location }}
IS_GCS_ARTIFACT: ${{ matrix.benchmark_entry.is_gcs_artifact }}
INPUT_FORMAT: ${{ matrix.benchmark_entry.input_format }}
XLA_FLAGS_JSON: ${{ toJson(matrix.benchmark_entry.xla_compilation_flags) }}
RUNTIME_FLAGS_JSON: ${{ toJson(matrix.benchmark_entry.runtime_flags) }}
TARGET_METRICS_JSON: ${{ toJson(matrix.benchmark_entry.target_metrics) }}
TOPOLOGY_JSON: ${{ toJson(matrix.benchmark_entry.topology) }}
HARDWARE_CATEGORY: ${{ matrix.benchmark_entry.hardware_category }} # e.g., NVIDIA_GPU_B200
# Workflow context variables
CHECKOUT_REF: ${{ github.sha }}
COMMIT_SHA: ${{ github.sha }}
WORKFLOW_RUN_ID: ${{ github.run_id }}
# Paths and script locations
OUTPUT_DIR_NAME: benchmark_output_${{ matrix.benchmark_entry.config_id }} # config_id includes workflow type
SCRIPT_DIR_RELATIVE: .github/workflows/benchmarks
BASELINE_YAML_FILE_RELATIVE: xla/tools/benchmarks/baseline/nightly_baseline.yml
COMPARISON_SCRIPT_RELATIVE: .github/workflows/benchmarks/compare_with_baseline.py
steps:
- name: "Wait For Connection"
uses: google-ml-infra/actions/ci_connection@7f5ca0c263a81ed09ea276524c1b9192f1304e3c
with:
halt-dispatch-input: ${{ inputs.halt-for-connection }}
- name: Print Job Info & Set Full Paths in ENV
run: |
# Resolve full paths based on GITHUB_WORKSPACE and relative paths defined in env
FULL_OUTPUT_DIR_PATH="${GITHUB_WORKSPACE}/${OUTPUT_DIR_NAME}"
RESOLVED_SCRIPT_DIR_PATH="${GITHUB_WORKSPACE}/${SCRIPT_DIR_RELATIVE}"
RESOLVED_BASELINE_YAML_PATH="${GITHUB_WORKSPACE}/${BASELINE_YAML_FILE_RELATIVE}" # Will use nightly_baseline.yml
RESOLVED_COMPARISON_SCRIPT_PATH="${GITHUB_WORKSPACE}/${COMPARISON_SCRIPT_RELATIVE}"
RESOLVED_RUN_COMPARISON_SCRIPT_PATH="${GITHUB_WORKSPACE}/${SCRIPT_DIR_RELATIVE}/run_comparison.sh"
echo "--- Benchmark Job Info (Nightly) ---"
echo "Config ID (from matrix): $CONFIG_ID"
echo "Benchmark Name: $BENCHMARK_NAME"
echo "Runner Label: $RUNNER_LABEL"
echo "Hardware Category: $HARDWARE_CATEGORY"
echo "Output Directory Name (relative): $OUTPUT_DIR_NAME"
echo "Full Output Directory Path: $FULL_OUTPUT_DIR_PATH"
echo "Full Script Directory Path: $RESOLVED_SCRIPT_DIR_PATH"
echo "Full Baseline YAML Path: $RESOLVED_BASELINE_YAML_PATH" # This will point to nightly_baseline.yml
echo "Full Comparison Script Path: $RESOLVED_COMPARISON_SCRIPT_PATH"
echo "Full Run Comparison Script Path: $RESOLVED_RUN_COMPARISON_SCRIPT_PATH"
echo "GITHUB_WORKSPACE: ${GITHUB_WORKSPACE}"
echo "Current PWD: $(pwd)"
echo "---------------------------------------"
# Make resolved paths available to subsequent steps via GITHUB_ENV
{
echo "RESOLVED_OUTPUT_DIR=$FULL_OUTPUT_DIR_PATH"
echo "RESOLVED_SCRIPT_DIR=$RESOLVED_SCRIPT_DIR_PATH"
echo "RESOLVED_BASELINE_YAML=$RESOLVED_BASELINE_YAML_PATH"
echo "RESOLVED_COMPARISON_SCRIPT=$RESOLVED_COMPARISON_SCRIPT_PATH"
echo "RESOLVED_RUN_COMPARISON_SCRIPT=$RESOLVED_RUN_COMPARISON_SCRIPT_PATH"
} >> "$GITHUB_ENV"
# Create the output directory; scripts will assume it exists
mkdir -p "$FULL_OUTPUT_DIR_PATH"
if [ ! -d "$FULL_OUTPUT_DIR_PATH" ]; then
echo "::error::Failed to create output directory: $FULL_OUTPUT_DIR_PATH"
exit 1
fi
- name: Checkout OpenXLA Repository
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
with:
ref: ${{ env.CHECKOUT_REF }}
persist-credentials: false
- name: Build Binaries
id: build_binaries
run: |
bash "${RESOLVED_SCRIPT_DIR}/build_binaries.sh"
env:
OUTPUT_DIR: ${{ env.RESOLVED_OUTPUT_DIR }}
- name: Prepare Benchmark Artifact
id: prep_artifact
run: |
bash "${RESOLVED_SCRIPT_DIR}/prepare_artifact.sh"
env:
OUTPUT_DIR: ${{ env.RESOLVED_OUTPUT_DIR }}
- name: Run Benchmark and Generate Stats
id: run_hlo
env:
RUNNER_BINARY: "${{ steps.build_binaries.outputs.runner_binary }}"
STATS_BINARY: "${{ steps.build_binaries.outputs.stats_binary }}"
DEVICE_TYPE_FLAG: "${{ steps.build_binaries.outputs.device_type_flag }}"
LOCAL_ARTIFACT_PATH: "${{ steps.prep_artifact.outputs.artifact_local_path }}"
OUTPUT_DIR: ${{ env.RESOLVED_OUTPUT_DIR }}
# Other job-level env vars are automatically inherited
run: |
# The run_benchmark.sh script is expected to create $RESOLVED_OUTPUT_DIR/results.json
bash "${RESOLVED_SCRIPT_DIR}/run_benchmark.sh"
# Print the content of results.json for debugging purposes
RESULTS_JSON_FILE_PATH="${RESOLVED_OUTPUT_DIR}/results.json"
echo "--- Content of results.json (from workflow) ---"
if [ -f "$RESULTS_JSON_FILE_PATH" ]; then
if command -v jq &> /dev/null && jq '.' "$RESULTS_JSON_FILE_PATH" > /dev/null 2>&1; then
jq '.' "$RESULTS_JSON_FILE_PATH"
else
echo "results.json may not be valid JSON or jq failed, printing with cat:"
cat "$RESULTS_JSON_FILE_PATH"
fi
elif [ -f "${RESULTS_JSON_FILE_PATH}.txt" ]; then # Fallback for non-JSON or renamed output
echo "results.json not found, printing fallback .txt file:"
cat "${RESULTS_JSON_FILE_PATH}.txt"
else
echo "::warning::Neither results.json nor results.json.txt found at $RESULTS_JSON_FILE_PATH"
fi
echo "---------------------------------------------"
- name: Compare Results to Baseline
if: (steps.run_hlo.outcome == 'success' || (steps.run_hlo.outcome == 'failure' && (hashFiles(format('{0}/results.json', env.RESOLVED_OUTPUT_DIR)) != '' || hashFiles(format('{0}/results.json.txt', env.RESOLVED_OUTPUT_DIR)) != ''))) && env.skip_comparison_due_to_yaml_install_failure != 'true'
run: |
echo "Starting baseline comparison (against nightly baseline)..."
bash "$RESOLVED_RUN_COMPARISON_SCRIPT" # This script calls the Python comparison script
echo "Baseline comparison finished."
echo "---------------------------------------------"
- name: Upload results.json directly to GCS
run: |
GCS_BUCKET="gs://openxla-nightly-transient"
RESULTS_JSON_FILE_PATH="${RESOLVED_OUTPUT_DIR}/results.json"
# Check if the results file exists
if [ ! -f "$RESULTS_JSON_FILE_PATH" ]; then
echo "::error::results.json not found at $RESULTS_JSON_FILE_PATH"
exit 1
fi
# Construct a GCS object name
TIMESTAMP=$(date +%Y%m%d_%H%M%S)
DATE_FOLDER=$(date +%Y%m%d)
COMMIT_SHA_SHORT=$(echo "${{ github.sha }}" | cut -c1-8)
GCS_OBJECT_NAME="${BENCHMARK_NAME}/${DATE_FOLDER}/${TIMESTAMP}_run_${WORKFLOW_RUN_ID}_commit_${COMMIT_SHA_SHORT}.json"
echo "Uploading $RESULTS_JSON_FILE_PATH to $GCS_BUCKET/$GCS_OBJECT_NAME"
gcloud storage cp "$RESULTS_JSON_FILE_PATH" "$GCS_BUCKET/$GCS_OBJECT_NAME"
- name: Upload Benchmark Artifacts
if: always()
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: results-${{ env.CONFIG_ID }}
path: ${{ env.RESOLVED_OUTPUT_DIR }}
retention-days: 14
if-no-files-found: error
third_party/xla/.github/workflows/oneAPI_Postsubmit.yml
Vendored
+53
View File
@@ -0,0 +1,53 @@
# Copyright 2026 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.
name: Postsubmit CI
permissions:
contents: read
on:
workflow_dispatch:
push:
branches:
- main
concurrency:
group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
cancel-in-progress: ${{ github.ref != 'main' }}
jobs:
Tests:
name: XLA Linux X86 GPU ONEAPI POSTSUBMIT
runs-on: linux-x64-battlemage-256-1gpu-intel
container:
image: us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest # zizmor: ignore[unpinned-images]
options: --device=/dev/dri
timeout-minutes: 120
defaults:
run:
shell: bash
steps:
- name: Checking out openxla/xla
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0
with:
path: openxla/xla
persist-credentials: false
- name: Run oneAPI CI
working-directory: openxla/xla
run: bash build_tools/sycl/ci_docker_test.sh
third_party/xla/.github/workflows/postsubmit_benchmark.yml
Vendored
+235
View File
@@ -0,0 +1,235 @@
# 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/postsubmit_benchmarks.yml
name: Postsubmit - Run Benchmarks
permissions:
contents: read
on:
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'
push:
branches:
- main
concurrency:
# Run every push to main and do not cancel in-progress jobs; the timeout is 60 minutes.
group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
cancel-in-progress: false
jobs:
# =================================================================
# Job 1: Generate the matrix specifically for POSTSUBMIT benchmarks
# =================================================================
generate_matrix:
name: Generate Postsubmit Matrix
# Condition: Run if manually dispatched OR if it's a push event to the main branch.
if: |
github.event_name == 'workflow_dispatch' ||
(github.event_name == 'push' && github.ref == 'refs/heads/main')
uses: ./.github/workflows/generate_benchmark_matrix.yml
with:
workflow_type: 'POSTSUBMIT'
registry_file: 'xla/tools/benchmarks/registries/default_registry.yml'
checkout_ref: ${{ github.sha }} # On push/dispatch to main, github.sha is the commit SHA
run_benchmarks:
name: Run Benchmark (${{ matrix.benchmark_entry.config_id }}) # config_id will indicate the workflow type, e.g., '_postsubmit'
needs: generate_matrix
# Also check if the generate_matrix job itself was skipped due to its `if` condition
if: success() && needs.generate_matrix.result != 'skipped' && needs.generate_matrix.outputs.matrix_include_json != '[]' && needs.generate_matrix.outputs.matrix_include_json != ''
strategy:
fail-fast: false # Continue running other benchmarks even if one fails
matrix:
benchmark_entry: ${{ fromJson(needs.generate_matrix.outputs.matrix_include_json || '[]') }}
runs-on: ${{ matrix.benchmark_entry.runner_label }}
container: ${{ matrix.benchmark_entry.container_image }}
defaults:
run:
shell: bash
timeout-minutes: 60
env:
# Variables from the matrix
BENCHMARK_NAME: ${{ matrix.benchmark_entry.benchmark_name }}
CONFIG_ID: ${{ matrix.benchmark_entry.config_id }}
RUNNER_LABEL: ${{ matrix.benchmark_entry.runner_label }}
CONTAINER_IMAGE: ${{ matrix.benchmark_entry.container_image }}
ARTIFACT_LOCATION: ${{ matrix.benchmark_entry.artifact_location }}
IS_GCS_ARTIFACT: ${{ matrix.benchmark_entry.is_gcs_artifact }}
INPUT_FORMAT: ${{ matrix.benchmark_entry.input_format }}
XLA_FLAGS_JSON: ${{ toJson(matrix.benchmark_entry.xla_compilation_flags) }}
RUNTIME_FLAGS_JSON: ${{ toJson(matrix.benchmark_entry.runtime_flags) }}
TARGET_METRICS_JSON: ${{ toJson(matrix.benchmark_entry.target_metrics) }}
TOPOLOGY_JSON: ${{ toJson(matrix.benchmark_entry.topology) }}
HARDWARE_CATEGORY: ${{ matrix.benchmark_entry.hardware_category }} # e.g., NVIDIA_GPU_B200
# Workflow context variables
CHECKOUT_REF: ${{ github.sha }} # For postsubmit, use the commit SHA from the push/dispatch event
COMMIT_SHA: ${{ github.sha }} # For postsubmit, use the commit SHA
WORKFLOW_RUN_ID: ${{ github.run_id }}
# Paths and script locations
OUTPUT_DIR_NAME: benchmark_output_${{ matrix.benchmark_entry.config_id }} # config_id includes workflow type
SCRIPT_DIR_RELATIVE: .github/workflows/benchmarks
BASELINE_YAML_FILE_RELATIVE: xla/tools/benchmarks/baseline/postsubmit_baseline.yml
COMPARISON_SCRIPT_RELATIVE: .github/workflows/benchmarks/compare_with_baseline.py
steps:
- name: "Wait For Connection"
uses: google-ml-infra/actions/ci_connection@7f5ca0c263a81ed09ea276524c1b9192f1304e3c
with:
halt-dispatch-input: ${{ inputs.halt-for-connection }}
- name: Print Job Info & Set Full Paths in ENV
run: |
# Resolve full paths based on GITHUB_WORKSPACE and relative paths defined in env
FULL_OUTPUT_DIR_PATH="${GITHUB_WORKSPACE}/${OUTPUT_DIR_NAME}"
RESOLVED_SCRIPT_DIR_PATH="${GITHUB_WORKSPACE}/${SCRIPT_DIR_RELATIVE}"
RESOLVED_BASELINE_YAML_PATH="${GITHUB_WORKSPACE}/${BASELINE_YAML_FILE_RELATIVE}" # Will use postsubmit_baseline.yml
RESOLVED_COMPARISON_SCRIPT_PATH="${GITHUB_WORKSPACE}/${COMPARISON_SCRIPT_RELATIVE}"
RESOLVED_RUN_COMPARISON_SCRIPT_PATH="${GITHUB_WORKSPACE}/${SCRIPT_DIR_RELATIVE}/run_comparison.sh"
echo "--- Benchmark Job Info (Postsubmit) ---"
echo "Config ID (from matrix): $CONFIG_ID"
echo "Benchmark Name: $BENCHMARK_NAME"
echo "Runner Label: $RUNNER_LABEL"
echo "Hardware Category: $HARDWARE_CATEGORY"
echo "Output Directory Name (relative): $OUTPUT_DIR_NAME"
echo "Full Output Directory Path: $FULL_OUTPUT_DIR_PATH"
echo "Full Script Directory Path: $RESOLVED_SCRIPT_DIR_PATH"
echo "Full Baseline YAML Path: $RESOLVED_BASELINE_YAML_PATH" # This will point to postsubmit_baseline.yml
echo "Full Comparison Script Path: $RESOLVED_COMPARISON_SCRIPT_PATH"
echo "Full Run Comparison Script Path: $RESOLVED_RUN_COMPARISON_SCRIPT_PATH"
echo "GITHUB_WORKSPACE: ${GITHUB_WORKSPACE}"
echo "Current PWD: $(pwd)"
echo "---------------------------------------"
# Make resolved paths available to subsequent steps via GITHUB_ENV
{
echo "RESOLVED_OUTPUT_DIR=$FULL_OUTPUT_DIR_PATH"
echo "RESOLVED_SCRIPT_DIR=$RESOLVED_SCRIPT_DIR_PATH"
echo "RESOLVED_BASELINE_YAML=$RESOLVED_BASELINE_YAML_PATH"
echo "RESOLVED_COMPARISON_SCRIPT=$RESOLVED_COMPARISON_SCRIPT_PATH"
echo "RESOLVED_RUN_COMPARISON_SCRIPT=$RESOLVED_RUN_COMPARISON_SCRIPT_PATH"
} >> "$GITHUB_ENV"
# Create the output directory; scripts will assume it exists
mkdir -p "$FULL_OUTPUT_DIR_PATH"
if [ ! -d "$FULL_OUTPUT_DIR_PATH" ]; then
echo "::error::Failed to create output directory: $FULL_OUTPUT_DIR_PATH"
exit 1
fi
- name: Checkout OpenXLA Repository
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
with:
ref: ${{ env.CHECKOUT_REF }}
persist-credentials: false
- name: Build Binaries
id: build_binaries
run: |
bash "${RESOLVED_SCRIPT_DIR}/build_binaries.sh"
env:
OUTPUT_DIR: ${{ env.RESOLVED_OUTPUT_DIR }}
- name: Prepare Benchmark Artifact
id: prep_artifact
run: |
bash "${RESOLVED_SCRIPT_DIR}/prepare_artifact.sh"
env:
OUTPUT_DIR: ${{ env.RESOLVED_OUTPUT_DIR }}
- name: Run Benchmark and Generate Stats
id: run_hlo
env:
RUNNER_BINARY: "${{ steps.build_binaries.outputs.runner_binary }}"
STATS_BINARY: "${{ steps.build_binaries.outputs.stats_binary }}"
DEVICE_TYPE_FLAG: "${{ steps.build_binaries.outputs.device_type_flag }}"
LOCAL_ARTIFACT_PATH: "${{ steps.prep_artifact.outputs.artifact_local_path }}"
OUTPUT_DIR: ${{ env.RESOLVED_OUTPUT_DIR }}
# Other job-level env vars are automatically inherited
run: |
# The run_benchmark.sh script is expected to create $RESOLVED_OUTPUT_DIR/results.json
bash "${RESOLVED_SCRIPT_DIR}/run_benchmark.sh"
# Print the content of results.json for debugging purposes
RESULTS_JSON_FILE_PATH="${RESOLVED_OUTPUT_DIR}/results.json"
echo "--- Content of results.json (from workflow) ---"
if [ -f "$RESULTS_JSON_FILE_PATH" ]; then
if command -v jq &> /dev/null && jq '.' "$RESULTS_JSON_FILE_PATH" > /dev/null 2>&1; then
jq '.' "$RESULTS_JSON_FILE_PATH"
else
echo "results.json may not be valid JSON or jq failed, printing with cat:"
cat "$RESULTS_JSON_FILE_PATH"
fi
elif [ -f "${RESULTS_JSON_FILE_PATH}.txt" ]; then # Fallback for non-JSON or renamed output
echo "results.json not found, printing fallback .txt file:"
cat "${RESULTS_JSON_FILE_PATH}.txt"
else
echo "::warning::Neither results.json nor results.json.txt found at $RESULTS_JSON_FILE_PATH"
fi
echo "---------------------------------------------"
- name: Compare Results to Baseline
if: (steps.run_hlo.outcome == 'success' || (steps.run_hlo.outcome == 'failure' && (hashFiles(format('{0}/results.json', env.RESOLVED_OUTPUT_DIR)) != '' || hashFiles(format('{0}/results.json.txt', env.RESOLVED_OUTPUT_DIR)) != ''))) && env.skip_comparison_due_to_yaml_install_failure != 'true'
run: |
echo "Starting baseline comparison (against postsubmit baseline)..."
bash "$RESOLVED_RUN_COMPARISON_SCRIPT" # This script calls the Python comparison script
echo "Baseline comparison finished."
echo "---------------------------------------------"
- name: Upload results.json directly to GCS
run: |
GCS_BUCKET="gs://openxla-postsubmit-transient"
RESULTS_JSON_FILE_PATH="${RESOLVED_OUTPUT_DIR}/results.json"
# Check if the results file exists
if [ ! -f "$RESULTS_JSON_FILE_PATH" ]; then
echo "::error::results.json not found at $RESULTS_JSON_FILE_PATH"
exit 1
fi
# Construct a GCS object name
TIMESTAMP=$(date +%Y%m%d_%H%M%S)
DATE_FOLDER=$(date +%Y%m%d)
COMMIT_SHA_SHORT=$(echo "${{ github.sha }}" | cut -c1-8)
GCS_OBJECT_NAME="${BENCHMARK_NAME}/${DATE_FOLDER}/${TIMESTAMP}_run_${WORKFLOW_RUN_ID}_commit_${COMMIT_SHA_SHORT}.json"
echo "Uploading $RESULTS_JSON_FILE_PATH to $GCS_BUCKET/$GCS_OBJECT_NAME"
gcloud storage cp "$RESULTS_JSON_FILE_PATH" "$GCS_BUCKET/$GCS_OBJECT_NAME"
- name: Upload Benchmark Artifacts
if: always()
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: results-${{ env.CONFIG_ID }}
path: ${{ env.RESOLVED_OUTPUT_DIR }}
retention-days: 14
if-no-files-found: error
third_party/xla/.github/workflows/pr_size_check.py
Vendored
+85
View File
@@ -0,0 +1,85 @@
# Copyright 2026 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.
# ============================================================================
"""Checks PR size and posts a comment if it's too large."""
import os
import sys
import github_api
THRESHOLD_MESSAGES = {
1000: (
"🔴 This PR has a very large delta of over 1000. In order to enable an"
" effective code reivew, please break the PR down into smaller and more"
" focused PRs.\nSee the [Small"
" CLs](https://google.github.io/eng-practices/review/developer/small-cls.html)"
" Google Eng practice for more details on how to write compact PRs."
),
500: (
"⚠️ This PR has a large delta of over 500. Consider breaking the PR"
" down into smaller PRs for a faster code review.\nSee the [Small"
" CLs](https://google.github.io/eng-practices/review/developer/small-cls.html)"
" Google Eng practice for more details on how to write compact PRs."
),
}
def main():
pr_number = os.getenv("PR_NUMBER")
if not pr_number:
print("PR_NUMBER is not set.")
sys.exit(1)
additions = int(os.getenv("ADDITIONS", "0"))
deletions = int(os.getenv("DELETIONS", "0"))
total_delta = additions + deletions
print(
f"PR #{pr_number} additions: {additions}, deletions: {deletions}, total:"
f" {total_delta}"
)
comment = (
f"PR #{pr_number} additions: {additions}, deletions: {deletions}, total:"
f" {total_delta}"
)
exceeded = False
for threshold in sorted(THRESHOLD_MESSAGES.keys(), reverse=True):
if total_delta >= threshold:
comment += f"\n\n{THRESHOLD_MESSAGES[threshold]}"
exceeded = True
break
if exceeded:
token = os.getenv("GH_TOKEN")
repo = os.getenv("GITHUB_REPOSITORY")
if not repo:
print("GITHUB_REPOSITORY is not set.")
sys.exit(1)
try:
api = github_api.GitHubAPI(token)
api.write_issue_comment(repo, int(pr_number), comment)
print("Comment posted.")
except Exception as e: # pylint: disable=broad-except
print(f"Failed to post comment: {e}")
sys.exit(1)
else:
print("PR size is within limits.")
if __name__ == "__main__":
main()
third_party/xla/.github/workflows/pr_size_check.yml
Vendored
+47
View File
@@ -0,0 +1,47 @@
# Copyright 2026 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.
# ============================================================================
name: PR Size Check
permissions:
pull-requests: write
on:
# zizmor: ignore[dangerous-triggers]
pull_request_target:
branches:
- main
jobs:
check-size:
runs-on: ubuntu-22.04
defaults:
run:
shell: bash
env:
GH_TOKEN: ${{ github.token }}
PR_NUMBER: ${{ github.event.pull_request.number }}
ADDITIONS: ${{ github.event.pull_request.additions }}
DELETIONS: ${{ github.event.pull_request.deletions }}
timeout-minutes: 5
if: |
github.repository == 'openxla/xla' && (
github.event.sender.type == 'User' ||
contains(github.event.pull_request.body, 'FORCE_TEST_ACTIONS')
)
steps:
- name: "Checking out repository"
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
with:
persist-credentials: false
- name: "Check PR Size"
run: python3 .github/workflows/pr_size_check.py
third_party/xla/.github/workflows/presubmit_benchmark.yml
Vendored
+210
View File
@@ -0,0 +1,210 @@
# 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/presubmit_benchmarks.yml
name: Presubmit - Run Benchmarks
permissions:
contents: read
on:
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'
pull_request:
branches:
- main
concurrency:
# Group all jobs for a given PR together, and cancel presubmit if a new commit is pushed.
group: ${{ github.workflow }}-${{ github.head_ref }}
cancel-in-progress: true
jobs:
# ================================================================
# Job 1: Generate the matrix specifically for PRESUBMIT benchmarks
# ================================================================
generate_matrix:
name: Generate Presubmit Matrix
uses: ./.github/workflows/generate_benchmark_matrix.yml
with:
workflow_type: 'PRESUBMIT'
registry_file: 'xla/tools/benchmarks/registries/default_registry.yml'
checkout_ref: ${{ github.event.pull_request.head.sha || github.sha }}
run_benchmarks:
name: Run Benchmark (${{ matrix.benchmark_entry.config_id }}) # Use the full config_id: {name}_{hw-short}_{topo}_{workflow} e.g., gemma3_1b_flax_call_l4_1h1d_presubmit
needs: generate_matrix
# Also check if the generate_matrix job itself was skipped due to its `if` condition
if: success() && needs.generate_matrix.result != 'skipped' && needs.generate_matrix.outputs.matrix_include_json != '[]' && needs.generate_matrix.outputs.matrix_include_json != ''
strategy:
fail-fast: false
matrix:
benchmark_entry: ${{ fromJson(needs.generate_matrix.outputs.matrix_include_json || '[]') }}
runs-on: ${{ matrix.benchmark_entry.runner_label }}
container: ${{ matrix.benchmark_entry.container_image }}
defaults:
run:
shell: bash
timeout-minutes: 60
env:
BENCHMARK_NAME: ${{ matrix.benchmark_entry.benchmark_name }}
CONFIG_ID: ${{ matrix.benchmark_entry.config_id }}
RUNNER_LABEL: ${{ matrix.benchmark_entry.runner_label }}
CONTAINER_IMAGE: ${{ matrix.benchmark_entry.container_image }}
ARTIFACT_LOCATION: ${{ matrix.benchmark_entry.artifact_location }}
IS_GCS_ARTIFACT: ${{ matrix.benchmark_entry.is_gcs_artifact }}
INPUT_FORMAT: ${{ matrix.benchmark_entry.input_format }}
XLA_FLAGS_JSON: ${{ toJson(matrix.benchmark_entry.xla_compilation_flags) }}
RUNTIME_FLAGS_JSON: ${{ toJson(matrix.benchmark_entry.runtime_flags) }}
TARGET_METRICS_JSON: ${{ toJson(matrix.benchmark_entry.target_metrics) }}
TOPOLOGY_JSON: ${{ toJson(matrix.benchmark_entry.topology) }}
HARDWARE_CATEGORY: ${{ matrix.benchmark_entry.hardware_category }}
CHECKOUT_REF: ${{ github.event.pull_request.head.sha || github.sha }}
COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.sha }}
WORKFLOW_RUN_ID: ${{ github.run_id }}
OUTPUT_DIR_NAME: benchmark_output_${{ matrix.benchmark_entry.config_id }}
SCRIPT_DIR_RELATIVE: .github/workflows/benchmarks
BASELINE_YAML_FILE_RELATIVE: xla/tools/benchmarks/baseline/presubmit_baseline.yml
COMPARISON_SCRIPT_RELATIVE: .github/workflows/benchmarks/compare_with_baseline.py
steps:
- name: "Wait For Connection"
uses: google-ml-infra/actions/ci_connection@7f5ca0c263a81ed09ea276524c1b9192f1304e3c
with:
halt-dispatch-input: ${{ inputs.halt-for-connection }}
- name: Print Job Info & Set Full Paths in ENV
run: |
FULL_OUTPUT_DIR_PATH="${GITHUB_WORKSPACE}/${OUTPUT_DIR_NAME}"
RESOLVED_SCRIPT_DIR_PATH="${GITHUB_WORKSPACE}/${SCRIPT_DIR_RELATIVE}"
RESOLVED_BASELINE_YAML_PATH="${GITHUB_WORKSPACE}/${BASELINE_YAML_FILE_RELATIVE}"
RESOLVED_COMPARISON_SCRIPT_PATH="${GITHUB_WORKSPACE}/${COMPARISON_SCRIPT_RELATIVE}"
RESOLVED_RUN_COMPARISON_SCRIPT_PATH="${GITHUB_WORKSPACE}/${SCRIPT_DIR_RELATIVE}/run_comparison.sh"
echo "--- Benchmark Job Info ---"
echo "Config ID (from matrix, used for dir): $CONFIG_ID"
echo "Benchmark Name: $BENCHMARK_NAME"
echo "Runner Label: $RUNNER_LABEL"
echo "Hardware Category: $HARDWARE_CATEGORY"
echo "Output Directory Name (relative): $OUTPUT_DIR_NAME"
echo "Full Output Directory Path: $FULL_OUTPUT_DIR_PATH"
echo "Full Script Directory Path: $RESOLVED_SCRIPT_DIR_PATH"
echo "Full Baseline YAML Path: $RESOLVED_BASELINE_YAML_PATH"
echo "Full Comparison Script Path: $RESOLVED_COMPARISON_SCRIPT_PATH"
echo "Full Compare Results Script Path: $RESOLVED_COMPARE_RESULTS_SCRIPT_PATH"
echo "GITHUB_WORKSPACE: ${GITHUB_WORKSPACE}"
echo "Current PWD: $(pwd)"
echo "--------------------------"
# These are now available to subsequent steps via env context
{
echo "RESOLVED_OUTPUT_DIR=$FULL_OUTPUT_DIR_PATH"
echo "RESOLVED_SCRIPT_DIR=$RESOLVED_SCRIPT_DIR_PATH"
echo "RESOLVED_BASELINE_YAML=$RESOLVED_BASELINE_YAML_PATH"
echo "RESOLVED_COMPARISON_SCRIPT=$RESOLVED_COMPARISON_SCRIPT_PATH"
echo "RESOLVED_RUN_COMPARISON_SCRIPT=$RESOLVED_RUN_COMPARISON_SCRIPT_PATH"
} >> "$GITHUB_ENV"
# Create the output directory - scripts will assume it exists
mkdir -p "$FULL_OUTPUT_DIR_PATH"
if [ ! -d "$FULL_OUTPUT_DIR_PATH" ]; then
echo "::error::Failed to create output directory: $FULL_OUTPUT_DIR_PATH"
exit 1
fi
- name: Checkout OpenXLA Repository
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
with:
ref: ${{ env.CHECKOUT_REF }}
persist-credentials: false
- name: Build Binaries
id: build_binaries
# Use the RESOLVED_SCRIPT_DIR and RESOLVED_OUTPUT_DIR from GITHUB_ENV
run: |
bash "${RESOLVED_SCRIPT_DIR}/build_binaries.sh"
env:
OUTPUT_DIR: ${{ env.RESOLVED_OUTPUT_DIR }}
- name: Prepare Benchmark Artifact
id: prep_artifact
run: |
bash "${RESOLVED_SCRIPT_DIR}/prepare_artifact.sh"
env:
OUTPUT_DIR: ${{ env.RESOLVED_OUTPUT_DIR }}
- name: Run Benchmark and Generate Stats
id: run_hlo
env:
RUNNER_BINARY: "${{ steps.build_binaries.outputs.runner_binary }}"
STATS_BINARY: "${{ steps.build_binaries.outputs.stats_binary }}"
DEVICE_TYPE_FLAG: "${{ steps.build_binaries.outputs.device_type_flag }}"
LOCAL_ARTIFACT_PATH: "${{ steps.prep_artifact.outputs.artifact_local_path }}"
OUTPUT_DIR: ${{ env.RESOLVED_OUTPUT_DIR }}
# Other job-level env vars are automatically inherited
run: |
# The run_benchmark.sh script is expected to create $OUTPUT_DIR/results.json
bash "${RESOLVED_SCRIPT_DIR}/run_benchmark.sh"
# Print the content of results.json for debugging
RESULTS_JSON_FILE_PATH="${RESOLVED_OUTPUT_DIR}/results.json"
echo "--- Content of results.json (from workflow) ---"
if [ -f "$RESULTS_JSON_FILE_PATH" ]; then
if command -v jq &> /dev/null && jq '.' "$RESULTS_JSON_FILE_PATH" > /dev/null 2>&1; then
jq '.' "$RESULTS_JSON_FILE_PATH"
else
echo "results.json may not be valid JSON or jq failed, printing with cat:"
cat "$RESULTS_JSON_FILE_PATH"
fi
elif [ -f "${RESULTS_JSON_FILE_PATH}.txt" ]; then
echo "results.json not found, printing fallback .txt file:"
cat "${RESULTS_JSON_FILE_PATH}.txt"
else
echo "::warning::Neither results.json nor results.json.txt found at $RESULTS_JSON_FILE_PATH"
fi
echo "---------------------------------------------"
- name: Compare Results to Baseline
if: (steps.run_hlo.outcome == 'success' || (steps.run_hlo.outcome == 'failure' && (hashFiles(format('{0}/results.json', env.RESOLVED_OUTPUT_DIR)) != '' || hashFiles(format('{0}/results.json.txt', env.RESOLVED_OUTPUT_DIR)) != ''))) && env.skip_comparison_due_to_yaml_install_failure != 'true'
run: |
echo "Starting baseline comparison..."
bash "$RESOLVED_RUN_COMPARISON_SCRIPT"
echo "Baseline comparison finished."
echo "---------------------------------------------"
- name: Upload Benchmark Artifacts
if: always()
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: results-${{ env.CONFIG_ID }}
path: ${{ env.RESOLVED_OUTPUT_DIR }}
retention-days: 7
if-no-files-found: error
third_party/xla/.github/workflows/rocm_ci.yml
Vendored
+175
View File
@@ -0,0 +1,175 @@
# Copyright 2026 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.
# ============================================================================
name: CI ROCm
permissions:
contents: read
on:
pull_request:
workflow_dispatch:
workflow_call:
inputs:
rocm_xla_branch:
description: 'Branch name to fetch ROCm XLA scripts from'
required: false
default: 'rocm-dev-infra'
type: string
runner_label:
description: 'Runner label for GPU jobs'
required: false
default: 'amd-do-linux-xla-gpu-gfx950-1'
type: string
concurrency:
group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
cancel-in-progress: ${{ github.ref != 'main' }}
jobs:
rocm-config:
runs-on: ubuntu-latest
outputs:
docker-image: ${{ steps.out.outputs.docker-image }}
steps:
- id: out
shell: bash
run: |
# hermetic llvm based on ghcr.io/rocm/jax-ubu22.rocm7.2.4:latest
echo "docker-image=ghcr.io/rocm/jax-ubu22.rocm7.2.4@sha256:248db525342fb7438ccbc6f5f1396ec82c783f7e0ae5667cb698edcddc8fe2fd" >> "$GITHUB_OUTPUT"
jax:
needs: rocm-config
if: ${{ github.event_name != 'pull_request' || github.base_ref == 'main' }}
name: JAX Linux x86 AMD Instinct GPU
runs-on: ${{ inputs.runner_label || 'amd-do-linux-xla-gpu-gfx950-1' }}
timeout-minutes: 80
env:
DOCKER_IMAGE: ${{ needs.rocm-config.outputs.docker-image }}
container: &rocm_container
image: ${{ needs.rocm-config.outputs.docker-image }} # zizmor: ignore[unpinned-images] image is pinned in the rocm-config job.
volumes:
- /data/ci-cert.crt:/data/ci-cert.crt
- /data/ci-cert.key:/data/ci-cert.key
- /data:/data
options: >-
--device=/dev/dri
--device=/dev/kfd
--ipc=host
--shm-size=64G
--cap-add=SYS_PTRACE
--security-opt=seccomp=unconfined
--tmpfs /root/.cache/bazel:rw,exec,size=80g
--group-add video
--env-file /etc/podinfo/gha-gpu-isolation-settings
defaults:
run:
shell: bash
steps:
- name: Checking out openxla/xla
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
with:
repository: openxla/xla
persist-credentials: false
- name: Checkout JAX
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
with:
repository: jax-ml/jax
path: jax
persist-credentials: false
- name: CPU Info
run: lscpu
- name: ROCm info
run: /opt/rocm/bin/rocminfo
- name: Run JAX Unit Tests
timeout-minutes: 60
working-directory: jax
run: |
./ci/run_bazel_test_rocm_rbe.sh \
--override_repository=xla="${GITHUB_WORKSPACE}" \
--override_module=xla="${GITHUB_WORKSPACE}" \
--config=single_gpu \
--//jax:build_jaxlib=wheel \
--//jax:build_jax=true \
--local_test_jobs=1 \
--action_env=JAX_ENABLE_X64=0 \
--repo_env=HERMETIC_PYTHON_VERSION=3.14 \
--repo_env=TF_ROCM_RBE_DOCKER_IMAGE="${DOCKER_IMAGE}" \
--bes_keywords=xla \
--bes_keywords=jax \
--bes_keywords=upstream \
--bes_keywords=gpu \
--bes_keywords=pull_request
xla:
needs: rocm-config
name: XLA Linux x86 AMD Instinct GPU
runs-on: ${{ inputs.runner_label || 'amd-do-linux-xla-gpu-gfx950-1' }}
timeout-minutes: 220
env:
DOCKER_IMAGE: ${{ needs.rocm-config.outputs.docker-image }}
EXECUTE_CI_BUILD_URL: https://raw.githubusercontent.com/ROCm/xla/refs/heads/${{ inputs.rocm_xla_branch || 'rocm-dev-infra' }}/build_tools/rocm/execute_ci_build_upstream.sh
container: *rocm_container
defaults:
run:
shell: bash
steps:
- name: Checking out openxla/xla
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
with:
repository: openxla/xla
persist-credentials: false
- name: Fetch ROCm XLA Upstream Script from ROCm/xla
run: |
wget -O build_tools/rocm/execute_ci_build_upstream.sh "${EXECUTE_CI_BUILD_URL}"
chmod +x build_tools/rocm/execute_ci_build_upstream.sh
- name: CPU Info
run: lscpu
- name: ROCm info
run: /opt/rocm/bin/rocminfo
- name: Test XLA [single_gpu]
timeout-minutes: 120
run: |
build_tools/rocm/execute_ci_build_upstream.sh \
--config=rocm_ci \
--config=rocm_rbe \
--config=ci_single_gpu \
--local_test_jobs=1 \
--repo_env=TF_ROCM_RBE_DOCKER_IMAGE="${DOCKER_IMAGE}" \
--local_test_jobs=1 \
--internal_spawn_scheduler \
--strategy=TestRunner=dynamic \
--bes_keywords=xla \
--bes_keywords=upstream \
--bes_keywords=gpu \
--bes_keywords=pull_request
- name: Test XLA [rocm_cpu]
timeout-minutes: 80
run: |
build_tools/rocm/execute_ci_build_upstream.sh \
--config=rocm_ci \
--config=rocm_rbe \
--config=ci_rocm_cpu \
--local_test_jobs=200 \
--repo_env=TF_ROCM_RBE_DOCKER_IMAGE="${DOCKER_IMAGE}" \
--bes_keywords=xla \
--bes_keywords=upstream \
--bes_keywords=cpu \
--bes_keywords=pull_request
third_party/xla/.github/workflows/rollback_notification.py
Vendored
+77
View File
@@ -0,0 +1,77 @@
# Copyright 2023 The TensorFlow 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.
# ============================================================================
"""Notifies PR authors of rollbacks on committed PRs."""
import os
import re
from typing import Generator, Sequence
import github_api
def get_reverted_commit_hashes(message: str) -> list[str]:
"""Searches a commit message for `reverts <commit hash>` and returns the found SHAs.
Arguments:
message: the commit message to search
Returns:
A list of SHAs as strings.
"""
print("Head commit message:", message, sep="\n")
regex = re.compile(r"reverts ([0-9a-f]{5,40})", flags=re.IGNORECASE)
commit_hashes = regex.findall(message)
print(f"Found commit hashes reverted in this commit: {commit_hashes}")
return commit_hashes
def get_associated_prs(
api: github_api.GitHubAPI,
commit_hashes: Sequence[str],
) -> Generator[int, None, None]:
"""Finds PRs associated with commits.
Arguments:
api: GitHubAPI object which will be used to make requests
commit_hashes: A sequence of SHAs which may have PRs associated with them
Yields:
Associated pairs of (PR number, SHA), both as strings
"""
regex = re.compile(r"PR #(\d+)")
for commit_hash in commit_hashes:
response = api.get_commit("openxla/xla", commit_hash)
message = response["commit"]["message"]
if maybe_match := regex.match(message):
pr_number = maybe_match.group(1)
print(f"Found PR #{pr_number} associated with commit hash {commit_hash}")
yield int(pr_number)
print(f"Didn't find any PRs associated with commit hashes: {commit_hashes}")
def main():
api = github_api.GitHubAPI(os.getenv("GH_TOKEN"))
head_commit = api.get_commit("openxla/xla", "HEAD")
commit_hashes = get_reverted_commit_hashes(head_commit["commit"]["message"])
for pr_number in get_associated_prs(api, commit_hashes):
sha = head_commit["sha"]
api.write_issue_comment(
"openxla/xla", pr_number, f"This PR was rolled back in {sha}!"
)
api.set_issue_status("openxla/xla", pr_number, "open")
if __name__ == "__main__":
main()
third_party/xla/.github/workflows/rollback_notification.yml
Vendored
+41
View File
@@ -0,0 +1,41 @@
# Copyright 2023 The TensorFlow 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.
# ============================================================================
name: Rollback Notification
permissions:
pull-requests: write
on:
push:
branches:
- main
jobs:
rollback-notification:
if: contains(github.event.head_commit.message, 'revert')
runs-on: ubuntu-22.04
defaults:
run:
shell: bash
env:
GH_TOKEN: ${{ github.token }}
HEAD_COMMIT_MESSAGE: ${{ github.event.head_commit.message }}
timeout-minutes: 6
steps:
- name: "Checking out repository"
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
with:
persist-credentials: false
- name: "Check if PR was rolled back"
run: python3 .github/workflows/rollback_notification.py
third_party/xla/.github/workflows/scorecards-analysis.yml
Vendored
+72
View File
@@ -0,0 +1,72 @@
# Copyright 2023 The TensorFlow 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.
# ============================================================================
# This workflow uses actions that are not certified by GitHub. They are provided
# by a third-party and are governed by separate terms of service, privacy
# policy, and support documentation.
name: Scorecards supply-chain security
on:
# For Branch-Protection check. Only the default branch is supported. See
# https://github.com/ossf/scorecard/blob/main/docs/checks.md#branch-protection
branch_protection_rule:
# To guarantee Maintained check is occasionally updated. See
# https://github.com/ossf/scorecard/blob/main/docs/checks.md#maintained
schedule:
- cron: '26 3 * * 2'
push:
branches: ["main"]
# Declare default permissions as read only.
permissions: read-all
jobs:
analysis:
if: github.repository == 'openxla/xla' # Don't do this in forks
name: Scorecards analysis
runs-on: ubuntu-latest
permissions:
# Needed to upload the results to code-scanning dashboard.
security-events: write
# Needed to publish results and get a badge (see publish_results below).
id-token: write
steps:
- name: "Checkout code"
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
with:
persist-credentials: false
- name: "Run analysis"
uses: ossf/scorecard-action@4eaacf0543bb3f2c246792bd56e8cdeffafb205a # v2.4.3
with:
results_file: results.sarif
results_format: sarif
publish_results: true
# Upload the results as artifacts (optional). Commenting out will disable uploads of run results in SARIF
# format to the repository Actions tab.
- name: "Upload artifact"
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: SARIF file
path: results.sarif
retention-days: 5
# Upload the results to GitHub's code scanning dashboard (optional).
# Commenting out will disable upload of results to your repo's Code Scanning dashboard
- name: "Upload to code-scanning"
uses: github/codeql-action/upload-sarif@8aad20d150bbac5944a9f9d289da16a4b0d87c1e # v4.36.2
with:
sarif_file: results.sarif
third_party/xla/.github/workflows/zizmor.yml
Vendored
+41
View File
@@ -0,0 +1,41 @@
# Copyright 2026 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.
# ============================================================================
name: Zizmor
on:
pull_request:
push:
branches:
- main
permissions:
contents: read
jobs:
zizmor:
name: GitHub Actions Security Analysis
runs-on: ubuntu-latest
steps:
- name: Checkout code
uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
with:
persist-credentials: false
- name: Install uv dependencies
uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 # v8.2.0
- name: Run zizmor
run: uvx zizmor --format=github .github xla
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}