Merge branch 'pytorch:main' into main

Author: diro5t

.ci/docker/README.md

@@ -1,7 +1,7 @@
-# Docker images for Jenkins
+# Docker images for GitHub CI
 
 This directory contains everything needed to build the Docker images
-that are used in our CI
+that are used in our CI.
 
 The Dockerfiles located in subdirectories are parameterized to
 conditionally run build stages depending on build arguments passed to
@@ -12,13 +12,14 @@ each image as the `BUILD_ENVIRONMENT` environment variable.
 
 See `build.sh` for valid build environments (it's the giant switch).
 
-Docker builds are now defined with `.circleci/cimodel/data/simple/docker_definitions.py`
-
 ## Contents
 
 * `build.sh` -- dispatch script to launch all builds
 * `common` -- scripts used to execute individual Docker build stages
+* `ubuntu` -- Dockerfile for Ubuntu image for CPU build and test jobs
 * `ubuntu-cuda` -- Dockerfile for Ubuntu image with CUDA support for nvidia-docker
+* `ubuntu-rocm` -- Dockerfile for Ubuntu image with ROCm support
+* `ubuntu-xpu` -- Dockerfile for Ubuntu image with XPU support
 
 ## Usage
 

.ci/docker/build.sh

@@ -46,9 +46,7 @@ if [[ "$image" == *xla* ]]; then
   exit 0
 fi
 
-if [[ "$image" == *-bionic* ]]; then
-  UBUNTU_VERSION=18.04
-elif [[ "$image" == *-focal* ]]; then
+if [[ "$image" == *-focal* ]]; then
   UBUNTU_VERSION=20.04
 elif [[ "$image" == *-jammy* ]]; then
   UBUNTU_VERSION=22.04
@@ -73,6 +71,11 @@ if [[ "$image" == *cuda* && "$UBUNTU_VERSION" != "22.04" ]]; then
   DOCKERFILE="${OS}-cuda/Dockerfile"
 elif [[ "$image" == *rocm* ]]; then
   DOCKERFILE="${OS}-rocm/Dockerfile"
+elif [[ "$image" == *xpu* ]]; then
+  DOCKERFILE="${OS}-xpu/Dockerfile"
+elif [[ "$image" == *cuda*linter* ]]; then
+  # Use a separate Dockerfile for linter to keep a small image size
+  DOCKERFILE="linter-cuda/Dockerfile"
 elif [[ "$image" == *linter* ]]; then
   # Use a separate Dockerfile for linter to keep a small image size
   DOCKERFILE="linter/Dockerfile"
@@ -88,8 +91,8 @@ _UCC_COMMIT=7cb07a76ccedad7e56ceb136b865eb9319c258ea
 # configuration, so we hardcode everything here rather than do it
 # from scratch
 case "$image" in
-  pytorch-linux-bionic-cuda12.1-cudnn8-py3-gcc9)
-    CUDA_VERSION=12.1.0
+  pytorch-linux-focal-cuda12.1-cudnn8-py3-gcc9)
+    CUDA_VERSION=12.1.1
     CUDNN_VERSION=8
     ANACONDA_PYTHON_VERSION=3.10
     GCC_VERSION=9
@@ -102,8 +105,8 @@ case "$image" in
     CONDA_CMAKE=yes
     TRITON=yes
     ;;
-  pytorch-linux-bionic-cuda12.1-cudnn8-py3-gcc9-inductor-benchmarks)
-    CUDA_VERSION=12.1.0
+  pytorch-linux-focal-cuda12.1-cudnn8-py3-gcc9-inductor-benchmarks)
+    CUDA_VERSION=12.1.1
     CUDNN_VERSION=8
     ANACONDA_PYTHON_VERSION=3.10
     GCC_VERSION=9
@@ -117,7 +120,7 @@ case "$image" in
     TRITON=yes
     INDUCTOR_BENCHMARKS=yes
     ;;
-  pytorch-linux-bionic-cuda11.8-cudnn8-py3-gcc9)
+  pytorch-linux-focal-cuda11.8-cudnn8-py3-gcc9)
     CUDA_VERSION=11.8.0
     CUDNN_VERSION=8
     ANACONDA_PYTHON_VERSION=3.10
@@ -131,25 +134,11 @@ case "$image" in
     CONDA_CMAKE=yes
     TRITON=yes
     ;;
-  pytorch-linux-bionic-cuda11.8-cudnn8-py3-gcc7)
-    CUDA_VERSION=11.8.0
+  pytorch-linux-focal-cuda12.1-cudnn8-py3-gcc9)
+    CUDA_VERSION=12.1.1
     CUDNN_VERSION=8
     ANACONDA_PYTHON_VERSION=3.10
-    GCC_VERSION=7
-    PROTOBUF=yes
-    DB=yes
-    VISION=yes
-    KATEX=yes
-    UCX_COMMIT=${_UCX_COMMIT}
-    UCC_COMMIT=${_UCC_COMMIT}
-    CONDA_CMAKE=yes
-    TRITON=yes
-    ;;
-    pytorch-linux-bionic-cuda11.8-cudnn8-py3-gcc7-inductor-benchmarks)
-    CUDA_VERSION=11.8.0
-    CUDNN_VERSION=8
-    ANACONDA_PYTHON_VERSION=3.10
-    GCC_VERSION=7
+    GCC_VERSION=9
     PROTOBUF=yes
     DB=yes
     VISION=yes
@@ -158,16 +147,6 @@ case "$image" in
     UCC_COMMIT=${_UCC_COMMIT}
     CONDA_CMAKE=yes
     TRITON=yes
-    INDUCTOR_BENCHMARKS=yes
-    ;;
-  pytorch-linux-focal-py3-clang7-asan)
-    ANACONDA_PYTHON_VERSION=3.9
-    CLANG_VERSION=7
-    PROTOBUF=yes
-    DB=yes
-    VISION=yes
-    CONDA_CMAKE=yes
-    TRITON=yes
     ;;
   pytorch-linux-focal-py3-clang10-onnx)
     ANACONDA_PYTHON_VERSION=3.8
@@ -178,19 +157,19 @@ case "$image" in
     CONDA_CMAKE=yes
     ONNX=yes
     ;;
-  pytorch-linux-focal-py3-clang7-android-ndk-r19c)
+  pytorch-linux-focal-py3-clang9-android-ndk-r21e)
     ANACONDA_PYTHON_VERSION=3.8
-    CLANG_VERSION=7
+    CLANG_VERSION=9
     LLVMDEV=yes
     PROTOBUF=yes
     ANDROID=yes
-    ANDROID_NDK_VERSION=r19c
+    ANDROID_NDK_VERSION=r21e
     GRADLE_VERSION=6.8.3
     NINJA_VERSION=1.9.0
     ;;
-  pytorch-linux-bionic-py3.8-clang9)
+  pytorch-linux-focal-py3.8-clang10)
     ANACONDA_PYTHON_VERSION=3.8
-    CLANG_VERSION=9
+    CLANG_VERSION=10
     PROTOBUF=yes
     DB=yes
     VISION=yes
@@ -199,9 +178,9 @@ case "$image" in
     CONDA_CMAKE=yes
     TRITON=yes
     ;;
-  pytorch-linux-bionic-py3.11-clang9)
+  pytorch-linux-focal-py3.11-clang10)
     ANACONDA_PYTHON_VERSION=3.11
-    CLANG_VERSION=9
+    CLANG_VERSION=10
     PROTOBUF=yes
     DB=yes
     VISION=yes
@@ -210,7 +189,7 @@ case "$image" in
     CONDA_CMAKE=yes
     TRITON=yes
     ;;
-  pytorch-linux-bionic-py3.8-gcc9)
+  pytorch-linux-focal-py3.8-gcc9)
     ANACONDA_PYTHON_VERSION=3.8
     GCC_VERSION=9
     PROTOBUF=yes
@@ -225,7 +204,7 @@ case "$image" in
     PROTOBUF=yes
     DB=yes
     VISION=yes
-    ROCM_VERSION=5.3
+    ROCM_VERSION=5.7
     NINJA_VERSION=1.9.0
     CONDA_CMAKE=yes
     TRITON=yes
@@ -236,25 +215,25 @@ case "$image" in
     PROTOBUF=yes
     DB=yes
     VISION=yes
-    ROCM_VERSION=5.4.2
+    ROCM_VERSION=6.0
     NINJA_VERSION=1.9.0
     CONDA_CMAKE=yes
     TRITON=yes
     ;;
-  pytorch-linux-focal-py3.8-gcc7)
+  pytorch-linux-jammy-xpu-2024.0-py3)
     ANACONDA_PYTHON_VERSION=3.8
-    GCC_VERSION=7
+    GCC_VERSION=11
     PROTOBUF=yes
     DB=yes
     VISION=yes
-    KATEX=yes
+    BASEKIT_VERSION=2024.0.0-49522
+    NINJA_VERSION=1.9.0
     CONDA_CMAKE=yes
     TRITON=yes
-    DOCS=yes
     ;;
-    pytorch-linux-focal-py3.8-gcc7-inductor-benchmarks)
+    pytorch-linux-jammy-py3.8-gcc11-inductor-benchmarks)
     ANACONDA_PYTHON_VERSION=3.8
-    GCC_VERSION=7
+    GCC_VERSION=11
     PROTOBUF=yes
     DB=yes
     VISION=yes
@@ -274,13 +253,60 @@ case "$image" in
     VISION=yes
     TRITON=yes
     ;;
+  pytorch-linux-jammy-py3-clang12-asan)
+    ANACONDA_PYTHON_VERSION=3.9
+    CLANG_VERSION=12
+    PROTOBUF=yes
+    DB=yes
+    VISION=yes
+    CONDA_CMAKE=yes
+    TRITON=yes
+    ;;
+  pytorch-linux-jammy-py3-clang15-asan)
+    ANACONDA_PYTHON_VERSION=3.10
+    CLANG_VERSION=15
+    CONDA_CMAKE=yes
+    VISION=yes
+    ;;
+  pytorch-linux-jammy-py3.8-gcc11)
+    ANACONDA_PYTHON_VERSION=3.8
+    GCC_VERSION=11
+    PROTOBUF=yes
+    DB=yes
+    VISION=yes
+    KATEX=yes
+    CONDA_CMAKE=yes
+    TRITON=yes
+    DOCS=yes
+    UNINSTALL_DILL=yes
+    ;;
+  pytorch-linux-jammy-py3-clang12-executorch)
+    ANACONDA_PYTHON_VERSION=3.10
+    CLANG_VERSION=12
+    CONDA_CMAKE=yes
+    EXECUTORCH=yes
+    ;;
   pytorch-linux-focal-linter)
     # TODO: Use 3.9 here because of this issue https://github.com/python/mypy/issues/13627.
     # We will need to update mypy version eventually, but that's for another day. The task
     # would be to upgrade mypy to 1.0.0 with Python 3.11
     ANACONDA_PYTHON_VERSION=3.9
     CONDA_CMAKE=yes
     ;;
+  pytorch-linux-jammy-cuda11.8-cudnn8-py3.9-linter)
+    ANACONDA_PYTHON_VERSION=3.9
+    CUDA_VERSION=11.8
+    CONDA_CMAKE=yes
+    ;;
+  pytorch-linux-jammy-aarch64-py3.10-gcc11)
+    ANACONDA_PYTHON_VERSION=3.10
+    GCC_VERSION=11
+    ACL=yes
+    PROTOBUF=yes
+    DB=yes
+    VISION=yes
+    CONDA_CMAKE=yes
+    ;;
   *)
     # Catch-all for builds that are not hardcoded.
     PROTOBUF=yes
@@ -298,6 +324,9 @@ case "$image" in
       extract_version_from_image_name rocm ROCM_VERSION
       NINJA_VERSION=1.9.0
       TRITON=yes
+      # To ensure that any ROCm config will build using conda cmake
+      # and thus have LAPACK/MKL enabled
+      CONDA_CMAKE=yes
     fi
     if [[ "$image" == *centos7* ]]; then
       NINJA_VERSION=1.10.2
@@ -331,14 +360,11 @@ if [[ "$image" == *cuda*  && ${OS} == "ubuntu" ]]; then
 fi
 
 # Build image
-# TODO: build-arg THRIFT is not turned on for any image, remove it once we confirm
-# it's no longer needed.
-docker build \
+DOCKER_BUILDKIT=1 docker build \
        --no-cache \
        --progress=plain \
        --build-arg "BUILD_ENVIRONMENT=${image}" \
        --build-arg "PROTOBUF=${PROTOBUF:-}" \
-       --build-arg "THRIFT=${THRIFT:-}" \
        --build-arg "LLVMDEV=${LLVMDEV:-}" \
        --build-arg "DB=${DB:-}" \
        --build-arg "VISION=${VISION:-}" \
@@ -361,7 +387,7 @@ docker build \
        --build-arg "NINJA_VERSION=${NINJA_VERSION:-}" \
        --build-arg "KATEX=${KATEX:-}" \
        --build-arg "ROCM_VERSION=${ROCM_VERSION:-}" \
-       --build-arg "PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH:-gfx906}" \
+       --build-arg "PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH:-gfx906;gfx90a}" \
        --build-arg "IMAGE_NAME=${IMAGE_NAME}" \
        --build-arg "UCX_COMMIT=${UCX_COMMIT}" \
        --build-arg "UCC_COMMIT=${UCC_COMMIT}" \
@@ -370,6 +396,9 @@ docker build \
        --build-arg "ONNX=${ONNX}" \
        --build-arg "DOCS=${DOCS}" \
        --build-arg "INDUCTOR_BENCHMARKS=${INDUCTOR_BENCHMARKS}" \
+       --build-arg "EXECUTORCH=${EXECUTORCH}" \
+       --build-arg "BASEKIT_VERSION=${BASEKIT_VERSION}" \
+       --build-arg "ACL=${ACL:-}" \
        -f $(dirname ${DOCKERFILE})/Dockerfile \
        -t "$tmp_tag" \
        "$@" \

.ci/docker/build_docker.sh

@@ -98,6 +98,18 @@ COPY ./common/install_ninja.sh install_ninja.sh
 RUN if [ -n "${NINJA_VERSION}" ]; then bash ./install_ninja.sh; fi
 RUN rm install_ninja.sh
 
+ARG TRITON
+# Install triton, this needs to be done before sccache because the latter will
+# try to reach out to S3, which docker build runners don't have access
+ENV CMAKE_C_COMPILER cc
+ENV CMAKE_CXX_COMPILER c++
+COPY ./common/install_triton.sh install_triton.sh
+COPY ./common/common_utils.sh common_utils.sh
+COPY ci_commit_pins/triton-rocm.txt triton-rocm.txt
+COPY triton_version.txt triton_version.txt
+RUN if [ -n "${TRITON}" ]; then bash ./install_triton.sh; fi
+RUN rm install_triton.sh common_utils.sh triton-rocm.txt triton_version.txt
+
 # Install ccache/sccache (do this last, so we get priority in PATH)
 COPY ./common/install_cache.sh install_cache.sh
 ENV PATH /opt/cache/bin:$PATH

.ci/docker/centos-rocm/Dockerfile

@@ -0,0 +1 @@
+d4b3e5cc607e97afdba79dc90f8ef968142f347c