name: Roll PyTorch

on:
  schedule:
    - cron: '0 13 * * *'
  workflow_dispatch:

jobs:
  build_linux:
    name: Manylinux Build
    runs-on: ubuntu-latest

    # Don't run this in everyone's forks.
    if: github.repository == 'llvm/torch-mlir'

    steps:
    - name: Get torch-mlir
      uses: actions/checkout@v3
      with:
        submodules: 'true'
        token: ${{ secrets.WORKFLOW_INVOCATION_TOKEN }}

    - name: Setup ccache
      uses: ./.github/actions/setup-build
      with:
        cache-suffix: 'rollPyTorch'

    - name: Determine nightly PyTorch version
      run: |
        set -eo pipefail

        cd ${GITHUB_WORKSPACE}
        python -m pip install wheel

        # Fetch the most recent nightly torchvision release
        VISION_RELEASE=$(python -m pip index versions -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html --pre torchvision | grep "Available versions" | tr ' ' '\n' | grep "^[0-9]" | sort --version-sort --reverse | head -n1 | tr -d ',' | sed 's/\([^+]*\).*/\1/')
        echo "Found torchvision release ${VISION_RELEASE}"

        # Fetch the whl file associated with the nightly torchvision release
        rm -f torch*.whl
        python -m pip download -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html --pre "torchvision==${VISION_RELEASE}"

        # Downloading the torchvision WHL also downloads the PyTorch WHL file
        # Read the version from the downloaded whl file without extracting it
        PT_RELEASE=$(unzip -p torch-*.whl 'torch-*/METADATA' | grep "^Version:" | awk '{ print $2 }' | sed 's/\([^+]*\).*/\1/')
        echo "Found torch release ${PT_RELEASE}"
        printf -- "-f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html\n--pre\ntorch==%s\ntorchvision==%s\n" "${PT_RELEASE}" "${VISION_RELEASE}" > pytorch-requirements.txt

        # Read the commit hash from the downloaded whl file without extracting it
        PT_HASH=$(unzip -p torch-"${PT_RELEASE}"*.whl torch/version.py | grep git_version | awk '{ print $3 }' | tr -d "'")
        echo "Found torch commit hash ${PT_HASH}"

        PT_HASH_CHANGED=0
        echo "${PT_HASH}" | cmp - pytorch-hash.txt --quiet || PT_HASH_CHANGED=$?
        echo "${PT_HASH}" > pytorch-hash.txt
        rm torch-"${PT_RELEASE}"*.whl

        # Write the release and hash to the environment file so that we can
        # retrieve them when creating a PR
        echo "PT_HASH=${PT_HASH}" >> ${GITHUB_ENV}
        echo "PT_RELEASE=${PT_RELEASE}" >> ${GITHUB_ENV}
        echo "PT_HASH_CHANGED=${PT_HASH_CHANGED}" >> ${GITHUB_ENV}

    - name: Build and test (in-tree), also update ODS and shape library
      if: env.PT_HASH_CHANGED != '0'
      run: |
        cd ${GITHUB_WORKSPACE}
        TM_PACKAGES="in-tree" TM_USE_PYTORCH_BINARY="OFF" \
        TORCH_MLIR_SRC_PYTORCH_BRANCH="${{ env.PT_HASH }}" \
        TORCH_MLIR_SRC_PYTORCH_RELEASE="${{ env.PT_RELEASE }}" \
        TM_UPDATE_ODS_AND_SHAPE_LIB="ON" \
        ./build_tools/python_deploy/build_linux_packages.sh

    - name: Push changes to main branch
      if: env.PT_HASH_CHANGED != '0'
      run: |
        cd ${GITHUB_WORKSPACE}
        git config user.email "torch-mlir@users.noreply.github.com"
        git config user.name "Roll PyTorch Action"
        git fetch --recurse-submodules=no
        git checkout main
        git add pytorch-hash.txt pytorch-requirements.txt lib/Dialect/Torch/Transforms/ShapeLibrary.cpp include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td
        git diff --cached --exit-code || (git commit -m "update PyTorch version to ${{ env.PT_RELEASE }}" && git push --set-upstream origin main)

    - name: Update PyTorch Build Cache (if running on main branch)
      if: github.ref_name == 'main'
      id: cache-pytorch
      uses: ashay/cache@v1
      with:
        path: ${{ github.workspace }}/build_tools/python_deploy/wheelhouse
        key: ${{ runner.os }}-pytorch-${{ env.PT_HASH }}