[WIP] Allow SampleOp to take in a dynamic shots value and return a dynamically-shaped result
#4165
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| 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 }} |