Cache PyTorch source builds to reduce CI time (#1500)

* ci: cache PyTorch source builds

This patch reduces the time spent in regular CI builds by caching
PyTorch source builds.  Specifically, this patch:

1. Makes CI lookup the cache entry for the PyTorch commit hash in
   pytorch-version.txt
2. If lookup was successful, CI fetches the previously-generated WHL
   file into the build_tools/python/wheelhouse directory
3. CI sets the `TM_PYTORCH_INSTALL_WITHOUT_REBUILD` variable to `true`
4. The build_libtorch.sh script then uses the downloaded WHL file
   instead of rebuilding PyTorch

* ci: warm up PyTorch source cache during daily RollPyTorch action

This patch makes the RollPyTorch action write the updated WHL file to
the cache, so that it can be later retrieved by CI that runs for each
PR.  We deliberately add the caching step to the end of the action since
the RollPyTorch action never needs to read from the cache, although
executing this step earlier in the process should not cause problems
either.

.github/workflows/RollPyTorch.yml

@@ -64,3 +64,10 @@ jobs:
         git checkout main
         git add pytorch-version.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
+      id: cache-pytorch
+      uses: actions/cache@v3
+      with:
+        path: ${{ github.workspace }}/build_tools/python_deploy/wheelhouse
+        key: ubuntu-x86_64-pytorch-${{ env.PT_HASH }}

.github/workflows/buildAndTest.yml

@@ -52,16 +52,32 @@ jobs:
       with:
         submodules: 'true'
 
+    - name: Fetch PyTorch commit hash
+      run: |
+        PT_HASH="$(cat ${GITHUB_WORKSPACE}/pytorch-version.txt)"
+        echo "PT_HASH=${PT_HASH}" >> ${GITHUB_ENV}
+
     - name: Setup ccache
       uses: ./.github/actions/setup-build
       with:
         cache-suffix: ${{ matrix.os-arch }}-${{ matrix.llvm-build }}-${{ matrix.torch-binary }}
 
+    - name: Cache PyTorch Build
+      id: cache-pytorch
+      uses: actions/cache@v3
+      with:
+        path: ${{ github.workspace }}/build_tools/python_deploy/wheelhouse
+        key: ${{ matrix.os-arch }}-pytorch-${{ env.PT_HASH }}
+
     - name: Build and Test os-arch='ubuntu-x86_64' llvm-build='${{ matrix.llvm-build }}' torch-binary='${{ matrix.torch-binary }}'
       if: ${{ matrix.os-arch == 'ubuntu-x86_64' }}
       run: |
         cd $GITHUB_WORKSPACE
-        TORCH_MLIR_SRC_PYTORCH_BRANCH="$(cat pytorch-version.txt)" TM_PACKAGES="${{ matrix.llvm-build }}" TM_USE_PYTORCH_BINARY="${{ matrix.torch-binary }}" ./build_tools/python_deploy/build_linux_packages.sh
+        TORCH_MLIR_SRC_PYTORCH_BRANCH="$(cat pytorch-version.txt)" \
+        TM_PACKAGES="${{ matrix.llvm-build }}" \
+        TM_USE_PYTORCH_BINARY="${{ matrix.torch-binary }}" \
+        TM_PYTORCH_INSTALL_WITHOUT_REBUILD="${{ steps.cache-pytorch.outputs.cache-hit }}" \
+        ./build_tools/python_deploy/build_linux_packages.sh
     - name: Configure os-arch='macos-arm64' llvm-build='in-tree' torch-binary='${{ matrix.torch-binary }}'
       # cross compile, can't test arm64
       if: ${{ matrix.os-arch == 'macos-arm64' && matrix.llvm-build == 'in-tree' }}

build_tools/build_libtorch.sh

@@ -7,6 +7,7 @@ PYTORCH_ROOT=${PYTORCH_ROOT:-$SRC_ROOT/externals/pytorch}
 PYTORCH_INSTALL_PATH=${PYTORCH_INSTALL_PATH:-$SRC_ROOT/libtorch}
 TORCH_MLIR_SRC_PYTORCH_REPO="${TORCH_MLIR_SRC_PYTORCH_REPO:-pytorch/pytorch}"
 TORCH_MLIR_SRC_PYTORCH_BRANCH="${TORCH_MLIR_SRC_PYTORCH_BRANCH:-master}"
+TM_PYTORCH_INSTALL_WITHOUT_REBUILD="${TM_PYTORCH_INSTALL_WITHOUT_REBUILD:-false}"
 PT_C_COMPILER="${PT_C_COMPILER:-clang}"
 PT_CXX_COMPILER="${PT_CXX_COMPILER:-clang++}"
 CMAKE_OSX_ARCHITECTURES="${CMAKE_OSX_ARCHITECTURES:-x86_64}"
@@ -151,10 +152,12 @@ unpack_pytorch() {
 
 #main
 echo "Building libtorch from source"
-checkout_pytorch
-install_requirements
-build_pytorch
-package_pytorch
+if [[ $TM_PYTORCH_INSTALL_WITHOUT_REBUILD != "true" ]]; then
+  checkout_pytorch
+  install_requirements
+  build_pytorch
+  package_pytorch
+fi
 if [[ $CMAKE_OSX_ARCHITECTURES = "arm64" ]]; then
   echo "${Yellow} Cross compiling for arm64 so unpacking PyTorch wheel for libs${NC}"
   unpack_pytorch

build_tools/python_deploy/build_linux_packages.sh

@@ -65,10 +65,18 @@ echo "Setting torch-mlir PyTorch Repo for source builds to: ${TORCH_MLIR_SRC_PYT
 export TORCH_MLIR_SRC_PYTORCH_BRANCH="${TORCH_MLIR_SRC_PYTORCH_BRANCH:-master}"
 echo "Setting torch-mlir PyTorch version for source builds to: ${TORCH_MLIR_SRC_PYTORCH_BRANCH}"
 
+# If using PyTorch source, install from the existing build instead of rebuilding
+# all of PyTorch.  This option is useful in CI, when it determines that the
+# PyTorch version has not changed between consecutive runs.
+export TM_PYTORCH_INSTALL_WITHOUT_REBUILD="${TM_PYTORCH_INSTALL_WITHOUT_REBUILD:-false}"
+
 function run_on_host() {
   echo "Running on host for $1:$@"
   echo "Outputting to ${TM_OUTPUT_DIR}"
-  rm -rf "${TM_OUTPUT_DIR}"
+  if [[ $TM_PYTORCH_INSTALL_WITHOUT_REBUILD != "true" ]]; then
+    # We want to use the cached files, so don't remove them.
+    rm -rf "${TM_OUTPUT_DIR}"
+  fi
   mkdir -p "${TM_OUTPUT_DIR}"
   case "$package" in
     torch-mlir)
@@ -115,6 +123,7 @@ function run_on_host() {
     -e "TM_USE_PYTORCH_BINARY=${TM_USE_PYTORCH_BINARY}" \
     -e "TORCH_MLIR_SRC_PYTORCH_REPO=${TORCH_MLIR_SRC_PYTORCH_REPO}" \
     -e "TORCH_MLIR_SRC_PYTORCH_BRANCH=${TORCH_MLIR_SRC_PYTORCH_BRANCH}" \
+    -e "TM_PYTORCH_INSTALL_WITHOUT_REBUILD=${TM_PYTORCH_INSTALL_WITHOUT_REBUILD}" \
     -e "CCACHE_DIR=/main_checkout/torch-mlir/.ccache" \
     "${TM_CURRENT_DOCKER_IMAGE}" \
     /bin/bash /main_checkout/torch-mlir/build_tools/python_deploy/build_linux_packages.sh
@@ -197,6 +206,7 @@ function build_in_tree() {
       -DTORCH_MLIR_USE_INSTALLED_PYTORCH="$torch_from_bin" \
       -DTORCH_MLIR_SRC_PYTORCH_REPO=${TORCH_MLIR_SRC_PYTORCH_REPO} \
       -DTORCH_MLIR_SRC_PYTORCH_BRANCH=${TORCH_MLIR_SRC_PYTORCH_BRANCH} \
+      -DTM_PYTORCH_INSTALL_WITHOUT_REBUILD=${TM_PYTORCH_INSTALL_WITHOUT_REBUILD} \
       -DPython3_EXECUTABLE="$(which python3)" \
       /main_checkout/torch-mlir/externals/llvm-project/llvm
   cmake --build /main_checkout/torch-mlir/build
@@ -310,6 +320,7 @@ function build_out_of_tree() {
       -DTORCH_MLIR_USE_INSTALLED_PYTORCH="$torch_from_bin" \
       -DTORCH_MLIR_SRC_PYTORCH_REPO=${TORCH_MLIR_SRC_PYTORCH_REPO} \
       -DTORCH_MLIR_SRC_PYTORCH_BRANCH=${TORCH_MLIR_SRC_PYTORCH_BRANCH} \
+      -DTM_PYTORCH_INSTALL_WITHOUT_REBUILD=${TM_PYTORCH_INSTALL_WITHOUT_REBUILD} \
       -DPython3_EXECUTABLE="$(which python3)" \
       /main_checkout/torch-mlir
   cmake --build /main_checkout/torch-mlir/build_oot

python/CMakeLists.txt

@@ -23,6 +23,7 @@ if (NOT TORCH_MLIR_USE_INSTALLED_PYTORCH)
   # Source builds
   set(ENV{TORCH_MLIR_SRC_PYTORCH_REPO} ${TORCH_MLIR_SRC_PYTORCH_REPO})
   set(ENV{TORCH_MLIR_SRC_PYTORCH_BRANCH} ${TORCH_MLIR_SRC_PYTORCH_BRANCH})
+  set(ENV{TM_PYTORCH_INSTALL_WITHOUT_REBUILD} ${TM_PYTORCH_INSTALL_WITHOUT_REBUILD})
   set(ENV{MACOSX_DEPLOYMENT_TARGET} ${MACOSX_DEPLOYMENT_TARGET})
   set(ENV{CMAKE_OSX_ARCHITECTURES} ${CMAKE_OSX_ARCHITECTURES})
   set(ENV{CMAKE_C_COMPILER_LAUNCHER} ${CMAKE_C_COMPILER_LAUNCHER})