build-wheel-linux-arm64.yaml

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 }}