Skip to content

[Frontend] Corrected output of PyTree when using qml.counts() #3251

[Frontend] Corrected output of PyTree when using qml.counts()

[Frontend] Corrected output of PyTree when using qml.counts() #3251

name: Build Catalyst Wheel on Linux (arm64)
on:
pull_request:
types:
- opened
- reopened
- synchronize
- labeled
- ready_for_review
push:
branches: [ main ]
workflow_dispatch:
workflow_call:
env:
MACOSX_DEPLOYMENT_TARGET: 14.0
concurrency:
group: Build Catalyst Wheel on Linux (arm64)-${{ github.ref }}
cancel-in-progress: true
jobs:
check_if_wheel_build_required:
if: github.event.pull_request.draft == false
uses: ./.github/workflows/check-for-wheel-build.yml
constants:
needs: [check_if_wheel_build_required]
if: needs.check_if_wheel_build_required.outputs.build-wheels == 'true'
name: "Set build matrix"
uses: ./.github/workflows/constants.yaml
build-dependencies:
needs: [constants, check_if_wheel_build_required]
strategy:
fail-fast: false
matrix:
python_version: [{major_minor: "3.10", patch: "14", package: "python3.10"}]
container_img: ["quay.io/pypa/manylinux_2_28_aarch64"]
container_name: ["manylinux_2_28_aarch64"]
name: Build Dependencies (Python ${{ matrix.python_version.major_minor }})
runs-on:
group: 'Office 24th floor M2 Mac'
if: needs.check_if_wheel_build_required.outputs.build-wheels == 'true'
steps:
- name: Checkout Catalyst repo
uses: actions/checkout@v4
- name: Setup Runner Environment
id: setup_env
uses: ./.github/workflows/utils/setup_self_hosted_macos_env_linux_arm64
with:
python_version: ${{ matrix.python_version.major_minor }}
# Cache external project sources
- name: Cache LLVM Source
id: cache-llvm-source
uses: actions/cache@v4
with:
path: mlir/llvm-project
key: llvm-${{ needs.constants.outputs.llvm_version }}-default-source
enableCrossOsArchive: True
- name: Cache MHLO Source
id: cache-mhlo-source
uses: actions/cache@v4
with:
path: mlir/mlir-hlo
key: mhlo-${{ needs.constants.outputs.mhlo_version }}-default-source
enableCrossOsArchive: True
- name: Cache Enzyme Source
id: cache-enzyme-source
uses: actions/cache@v4
with:
path: mlir/Enzyme
key: enzyme-${{ needs.constants.outputs.enzyme_version }}-default-source
enableCrossOsArchive: True
- name: Clone LLVM Submodule
if: steps.cache-llvm-source.outputs.cache-hit != 'true'
uses: actions/checkout@v4
with:
repository: llvm/llvm-project
ref: ${{ needs.constants.outputs.llvm_version }}
path: mlir/llvm-project
- name: Clone MHLO Submodule
if: steps.cache-mhlo-source.outputs.cache-hit != 'true'
uses: actions/checkout@v4
with:
repository: tensorflow/mlir-hlo
ref: ${{ needs.constants.outputs.mhlo_version }}
path: mlir/mlir-hlo
- name: Clone Enzyme Submodule
if: steps.cache-enzyme-source.outputs.cache-hit != 'true'
uses: actions/checkout@v4
with:
repository: EnzymeAD/Enzyme
ref: ${{ needs.constants.outputs.enzyme_version }}
path: mlir/Enzyme
# Cache external project builds
- name: Restore LLVM Build
id: cache-llvm-build
uses: actions/cache/restore@v4
with:
path: llvm-build
key: ${{ matrix.container_name }}-llvm-${{ needs.constants.outputs.llvm_version }}-wheel-build
- name: Restore MHLO Build
id: cache-mhlo-build
uses: actions/cache/restore@v4
with:
path: mhlo-build
key: ${{ matrix.container_name }}-mhlo-${{ needs.constants.outputs.mhlo_version }}-wheel-build
lookup-only: True
- name: Restore Enzyme Build
id: cache-enzyme-build
uses: actions/cache/restore@v4
with:
path: enzyme-build
key: ${{ matrix.container_name }}-enzyme-${{ needs.constants.outputs.llvm_version }}-${{ needs.constants.outputs.enzyme_version }}-wheel-build
lookup-only: True
- name: Build LLD
if: steps.cache-llvm-build.outputs.cache-hit != 'true'
run: |
set -x
# With GCC 13, LLVM fails some tests, then we use GCC 12 instead (Copied from below)
export GCC_VERSION=12
docker run --rm --platform linux/aarch64 \
-v /var/run/docker.sock:/var/run/docker.sock \
-v `pwd`:/catalyst \
-i ${{ matrix.container_img }} \
bash /catalyst/.github/workflows/scripts/linux_arm64/rh8/build_lld.sh $GCC_VERSION ${{ matrix.python_version.major_minor }} ${{ matrix.python_version.patch }} ${{ matrix.python_version.package }}
- name: Build LLVM / MLIR
if: steps.cache-llvm-build.outputs.cache-hit != 'true'
run: |
set -x
# With GCC 13, LLVM fails some tests, then we use GCC 12 instead
export GCC_VERSION=12
docker run --rm --platform linux/aarch64 \
-v /var/run/docker.sock:/var/run/docker.sock \
-v `pwd`:/catalyst \
-i ${{ matrix.container_img }} \
bash /catalyst/.github/workflows/scripts/linux_arm64/rh8/build_llvm.sh $GCC_VERSION ${{ matrix.python_version.major_minor }} ${{ matrix.python_version.patch }} ${{ matrix.python_version.package }}
- name: Save LLVM Build
id: save-llvm-build
if: steps.cache-llvm-build.outputs.cache-hit != 'true'
uses: actions/cache/save@v4
with:
path: llvm-build
key: ${{ matrix.container_name }}-llvm-${{ needs.constants.outputs.llvm_version }}-wheel-build
- name: Build MHLO Dialect
if: steps.cache-mhlo-build.outputs.cache-hit != 'true'
run: |
set -x
export GCC_VERSION=13
docker run --rm --platform linux/aarch64 \
-v /var/run/docker.sock:/var/run/docker.sock \
-v `pwd`:/catalyst \
-i ${{ matrix.container_img }} \
bash /catalyst/.github/workflows/scripts/linux_arm64/rh8/build_mhlo.sh $GCC_VERSION ${{ matrix.python_version.major_minor }} ${{ matrix.python_version.patch }} ${{ matrix.python_version.package }}
- name: Save MHLO Build
id: save-mhlo-build
if: steps.cache-mhlo-build.outputs.cache-hit != 'true'
uses: actions/cache/save@v4
with:
path: mhlo-build
key: ${{ matrix.container_name }}-mhlo-${{ needs.constants.outputs.mhlo_version }}-wheel-build
- name: Build Enzyme
if: steps.cache-enzyme-build.outputs.cache-hit != 'true'
run: |
set -x
export GCC_VERSION=13
docker run --rm --platform linux/aarch64 \
-v /var/run/docker.sock:/var/run/docker.sock \
-v `pwd`:/catalyst \
-i ${{ matrix.container_img }} \
bash /catalyst/.github/workflows/scripts/linux_arm64/rh8/build_enzyme.sh $GCC_VERSION ${{ matrix.python_version.major_minor }} ${{ matrix.python_version.patch }} ${{ matrix.python_version.package }}
- name: Save Enzyme Build
id: save-enzyme-build
if: steps.cache-enzyme-build.outputs.cache-hit != 'true'
uses: actions/cache/save@v4
with:
path: enzyme-build
key: ${{ matrix.container_name }}-enzyme-${{ needs.constants.outputs.llvm_version }}-${{ needs.constants.outputs.enzyme_version }}-wheel-build
catalyst-linux-wheels-arm64:
needs: [constants, build-dependencies]
strategy:
fail-fast: false
max-parallel: 2
matrix:
python_version: [{major_minor: "3.10", patch: "14", package: "python3.10", alternative: "310"},
{major_minor: "3.11", patch: "9", package: "python3.11", alternative: "311"},
{major_minor: "3.12", patch: "3", package: "python3.12", alternative: "312"}]
container_img: ["quay.io/pypa/manylinux_2_28_aarch64"]
container_name: ["manylinux_2_28_aarch64"]
name: Build Wheels (Python ${{ matrix.python_version.major_minor }})
runs-on:
group: 'Office 24th floor M2 Mac'
steps:
- name: Checkout Catalyst repo
uses: actions/checkout@v4
- name: Setup Runner Environment
id: setup_env
uses: ./.github/workflows/utils/setup_self_hosted_macos_env_linux_arm64
with:
python_version: ${{ matrix.python_version.major_minor }}
- name: Get Cached LLVM Source
id: cache-llvm-source
uses: actions/cache/restore@v4
with:
path: mlir/llvm-project
key: llvm-${{ needs.constants.outputs.llvm_version }}-default-source
enableCrossOsArchive: True
fail-on-cache-miss: True
- name: Get Cached LLVM Build
id: cache-llvm-build
uses: actions/cache/restore@v4
with:
path: llvm-build
key: ${{ matrix.container_name }}-llvm-${{ needs.constants.outputs.llvm_version }}-wheel-build
fail-on-cache-miss: True
- name: Get Cached MHLO Source
id: cache-mhlo-source
uses: actions/cache/restore@v4
with:
path: mlir/mlir-hlo
key: mhlo-${{ needs.constants.outputs.mhlo_version }}-default-source
enableCrossOsArchive: True
fail-on-cache-miss: True
- name: Get Cached MHLO Build
id: cache-mhlo-build
uses: actions/cache/restore@v4
with:
path: mhlo-build
key: ${{ matrix.container_name }}-mhlo-${{ needs.constants.outputs.mhlo_version }}-wheel-build
fail-on-cache-miss: True
- name: Get Cached Enzyme Source
id: cache-enzyme-source
uses: actions/cache/restore@v4
with:
path: mlir/Enzyme
key: enzyme-${{ needs.constants.outputs.enzyme_version }}-default-source
enableCrossOsArchive: True
fail-on-cache-miss: True
- name: Get Cached Enzyme Build
id: cache-enzyme-build
uses: actions/cache/restore@v4
with:
path: enzyme-build
key: ${{ matrix.container_name }}-enzyme-${{ needs.constants.outputs.llvm_version }}-${{ needs.constants.outputs.enzyme_version }}-wheel-build
fail-on-cache-miss: True
# Build Catalyst Wheel
- name: Build Catalyst Wheel
run: |
set -x
export GCC_VERSION=13
docker run --rm --platform linux/aarch64 \
-v /var/run/docker.sock:/var/run/docker.sock \
-v `pwd`:/catalyst \
-i ${{ matrix.container_img }} \
bash /catalyst/.github/workflows/scripts/linux_arm64/rh8/build_catalyst.sh $GCC_VERSION ${{ matrix.python_version.major_minor }} ${{ matrix.python_version.patch }} ${{ matrix.python_version.package }} ${{ matrix.python_version.alternative }}
- name: Upload Wheel Artifact
uses: actions/upload-artifact@v4
with:
name: catalyst-linux_arm64-wheel-py-${{ matrix.python_version.major_minor}}.zip
path: wheel/
retention-days: 14
test-wheels:
needs: [constants, catalyst-linux-wheels-arm64]
strategy:
fail-fast: false
max-parallel: 2
matrix:
python_version: [{major_minor: "3.10", patch: "14", package: "python3.10"},
{major_minor: "3.11", patch: "9", package: "python3.11"},
{major_minor: "3.12", patch: "3", package: "python3.12"}]
container_img: ["quay.io/pypa/manylinux_2_28_aarch64"]
container_name: ["manylinux_2_28_aarch64"]
# To check all wheels for supported python3 versions
name: Test Wheels (Python ${{ matrix.python_version.major_minor }})
runs-on:
group: 'Office 24th floor M2 Mac'
steps:
- name: Checkout Catalyst repo
uses: actions/checkout@v4
- name: Setup Runner Environment
id: setup_env
uses: ./.github/workflows/utils/setup_self_hosted_macos_env_linux_arm64
with:
python_version: ${{ matrix.python_version.major_minor }}
- name: Download Wheel Artifact
uses: actions/download-artifact@v4
with:
name: catalyst-linux_arm64-wheel-py-${{ matrix.python_version.major_minor }}.zip
path: dist
# Needed for accessing llvm-symbolizer
- name: Get Cached LLVM Build
id: cache-llvm-build
uses: actions/cache@v4
with:
path: llvm-build
key: ${{ matrix.container_name }}-llvm-${{ needs.constants.outputs.llvm_version }}-wheel-build
fail-on-cache-miss: True
- name: Run Python Pytest Tests
run: |
set -x
export GCC_VERSION=13
docker run --rm --platform linux/aarch64 \
-v /var/run/docker.sock:/var/run/docker.sock \
-v `pwd`:/catalyst \
-i ${{ matrix.container_img }} \
bash /catalyst/.github/workflows/scripts/linux_arm64/rh8/test_wheels.sh $GCC_VERSION ${{ matrix.python_version.major_minor }} ${{ matrix.python_version.patch }} ${{ matrix.python_version.package }}